From f8117181443f555beb1001d02f73b3bb9ae7d5b1 Mon Sep 17 00:00:00 2001
From: Gheorghe-Teodor Bercea <gt.bercea@gmail.com>
Date: Thu, 25 Jun 2020 13:36:39 -0400
Subject: [PATCH] Add new passes related to Linalg and Affine dialects (#181)

* Reorganize main function.

* Follow review comments.

* Emit constants are globals in Krnl and LLVM dialects.

* Add more passes to onnx-mlir-opt.

* Clean-up.
---
 MLIR.cmake                           | 17 +++++++++++++++++
 src/Tool/ONNXMLIROpt/ONNXMLIROpt.cpp | 10 ++++++++++
 2 files changed, 27 insertions(+)

diff --git a/MLIR.cmake b/MLIR.cmake
index 1a0cf76..815bfd5 100644
--- a/MLIR.cmake
+++ b/MLIR.cmake
@@ -146,6 +146,11 @@ find_mlir_lib(MLIRLLVMIR)
 find_mlir_lib(MLIRLoopAnalysis)
 find_mlir_lib(MLIRSCFToStandard)
 find_mlir_lib(MLIRLoopLikeInterface)
+find_mlir_lib(MLIRLinalgOps)
+find_mlir_lib(MLIRLinalgEDSC)
+find_mlir_lib(MLIRLinalgAnalysis)
+find_mlir_lib(MLIRLinalgTransforms)
+find_mlir_lib(MLIRLinalgUtils)
 find_mlir_lib(MLIRSCF)
 find_mlir_lib(MLIRLLVMIRTransforms)
 find_mlir_lib(MLIRMlirOptMain)
@@ -165,6 +170,10 @@ find_mlir_lib(MLIRTargetLLVMIR)
 find_mlir_lib(MLIRTransformUtils)
 find_mlir_lib(MLIRTranslation)
 find_mlir_lib(MLIRVector)
+find_mlir_lib(MLIRMlirOptMain)
+find_mlir_lib(MLIRAffineEDSC)
+find_mlir_lib(MLIRLinalgEDSC)
+find_mlir_lib(MLIRViewLikeInterface)
 
 find_mlir_lib(LLVMCore)
 find_mlir_lib(LLVMSupport)
@@ -226,6 +235,14 @@ set(MLIRLibs
         ${MLIRStandardToLLVM}
         ${MLIRSupport}
         ${MLIRTranslation}
+        ${MLIRLinalgOps}
+        ${MLIRLinalgEDSC}
+        ${MLIRLinalgAnalysis}
+        ${MLIRLinalgTransforms}
+        ${MLIRLinalgUtils}
+        ${MLIRAffineEDSC}
+        ${MLIRLinalgEDSC}
+        ${MLIRViewLikeInterface}
         # strict order verified
         ${LLVMBitWriter}
         ${LLVMObject}
diff --git a/src/Tool/ONNXMLIROpt/ONNXMLIROpt.cpp b/src/Tool/ONNXMLIROpt/ONNXMLIROpt.cpp
index d94e2e7..063af72 100644
--- a/src/Tool/ONNXMLIROpt/ONNXMLIROpt.cpp
+++ b/src/Tool/ONNXMLIROpt/ONNXMLIROpt.cpp
@@ -12,8 +12,11 @@
 #include <llvm/Support/InitLLVM.h>
 #include <llvm/Support/ToolOutputFile.h>
 #include <mlir/IR/AsmState.h>
+#include <mlir/IR/Dialect.h>
+#include <mlir/IR/MLIRContext.h>
 #include <mlir/InitAllDialects.h>
 #include <mlir/InitAllPasses.h>
+#include <mlir/Interfaces/ViewLikeInterface.h>
 #include <mlir/Pass/Pass.h>
 #include <mlir/Pass/PassManager.h>
 #include <mlir/Support/FileUtilities.h>
@@ -56,6 +59,7 @@ static llvm::cl::opt<bool> allowUnregisteredDialects(
     llvm::cl::init(false));
 
 int main(int argc, char **argv) {
+  mlir::registerDialect<mlir::linalg::LinalgDialect>();
   mlir::registerDialect<mlir::AffineDialect>();
   mlir::registerDialect<mlir::LLVM::LLVMDialect>();
   mlir::registerDialect<mlir::scf::SCFDialect>();
@@ -65,6 +69,12 @@ int main(int argc, char **argv) {
 #define GEN_PASS_REGISTRATION
 #include "mlir/Transforms/Passes.h.inc"
 
+#define GEN_PASS_REGISTRATION
+#include "mlir/Dialect/Affine/Passes.h.inc"
+
+#define GEN_PASS_REGISTRATION
+#include "mlir/Dialect/Linalg/Passes.h.inc"
+
   llvm::InitLLVM y(argc, argv);
 
   mlir::registerDialect<mlir::ONNXOpsDialect>();