From af75b4c75e18af43e869df2194fb4d9412fdaabf Mon Sep 17 00:00:00 2001
From: Gheorghe-Teodor Bercea <gt.bercea@gmail.com>
Date: Thu, 9 Jul 2020 18:12:57 -0400
Subject: [PATCH] Add selected MLIR options (#213)

* Reorganize main function.

* Follow review comments.

* Emit constants are globals in Krnl and LLVM dialects.

* Add selected MLIR passes.

* Refactor.

* Fix descriptions.

* Refactor.
---
 MLIR.cmake                           |  6 ++++++
 src/InitMLIRPasses.hpp               | 25 +++++++++++++++++++++++++
 src/Tool/ONNXMLIROpt/ONNXMLIROpt.cpp |  6 ++++++
 3 files changed, 37 insertions(+)
 create mode 100644 src/InitMLIRPasses.hpp

diff --git a/MLIR.cmake b/MLIR.cmake
index 45a3995..41d0050 100644
--- a/MLIR.cmake
+++ b/MLIR.cmake
@@ -152,6 +152,7 @@ find_mlir_lib(MLIRLinalgAnalysis)
 find_mlir_lib(MLIRLinalgTransforms)
 find_mlir_lib(MLIRLinalgUtils)
 find_mlir_lib(MLIRSCF)
+find_mlir_lib(MLIRSCFTransforms)
 find_mlir_lib(MLIRLLVMIRTransforms)
 find_mlir_lib(MLIRMlirOptMain)
 find_mlir_lib(MLIRParser)
@@ -171,6 +172,8 @@ find_mlir_lib(MLIRTargetLLVMIR)
 find_mlir_lib(MLIRTransformUtils)
 find_mlir_lib(MLIRTranslation)
 find_mlir_lib(MLIRVector)
+find_mlir_lib(MLIRVectorToLLVM)
+find_mlir_lib(MLIRVectorToSCF)
 find_mlir_lib(MLIRMlirOptMain)
 find_mlir_lib(MLIRAffineEDSC)
 find_mlir_lib(MLIRLinalgEDSC)
@@ -203,6 +206,8 @@ set(MLIRLibs
         ${MLIRTransforms}
         ${MLIRSCFToStandard}
         ${MLIRVector}
+        ${MLIRVectorToLLVM}
+        ${MLIRVectorToSCF}
         ${MLIRSCF}
         ${MLIRIR}
         ${MLIRLLVMIR}
@@ -226,6 +231,7 @@ set(MLIRLibs
         ${MLIRLLVMIRTransforms}        
         ${MLIRSCFToStandard}
         ${MLIRSCF}
+        ${MLIRSCFTransforms}
         ${MLIRLoopAnalysis}
         ${MLIRLoopLikeInterface}
         ${MLIROpenMP}
diff --git a/src/InitMLIRPasses.hpp b/src/InitMLIRPasses.hpp
new file mode 100644
index 0000000..d1bdaf9
--- /dev/null
+++ b/src/InitMLIRPasses.hpp
@@ -0,0 +1,25 @@
+#include "mlir/Pass/Pass.h"
+
+namespace onnx_mlir {
+
+void initMLIRPasses() {
+  mlir::registerPass("convert-vector-to-scf", "Convert vector to SCF.",
+      []() -> std::unique_ptr<mlir::Pass> {
+        return mlir::createConvertVectorToSCFPass();
+      });
+  mlir::registerPass("lower-affine",
+      "Lower Affine Dialect to Standard Dialect.",
+      []() -> std::unique_ptr<mlir::Pass> {
+        return mlir::createLowerAffinePass();
+      });
+  mlir::registerPass("convert-scf-to-std", "Lower SCF to Standard Dialect.",
+      []() -> std::unique_ptr<mlir::Pass> {
+        return mlir::createLowerToCFGPass();
+      });
+  mlir::registerPass("convert-vector-to-llvm",
+      "Lower Vector Dialect to LLVM IR Dialect.",
+      []() -> std::unique_ptr<mlir::Pass> {
+        return mlir::createConvertVectorToLLVMPass();
+      });
+}
+} // namespace onnx_mlir
\ No newline at end of file
diff --git a/src/Tool/ONNXMLIROpt/ONNXMLIROpt.cpp b/src/Tool/ONNXMLIROpt/ONNXMLIROpt.cpp
index 063af72..f3a4c27 100644
--- a/src/Tool/ONNXMLIROpt/ONNXMLIROpt.cpp
+++ b/src/Tool/ONNXMLIROpt/ONNXMLIROpt.cpp
@@ -24,6 +24,7 @@
 
 #include "src/Dialect/Krnl/KrnlOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
+#include "src/InitMLIRPasses.hpp"
 #include "src/InitOMPasses.hpp"
 #include "src/Pass/Passes.hpp"
 
@@ -64,6 +65,7 @@ int main(int argc, char **argv) {
   mlir::registerDialect<mlir::LLVM::LLVMDialect>();
   mlir::registerDialect<mlir::scf::SCFDialect>();
   mlir::registerDialect<mlir::StandardOpsDialect>();
+  mlir::registerDialect<mlir::vector::VectorDialect>();
 
   // Register transformation passes.
 #define GEN_PASS_REGISTRATION
@@ -75,11 +77,15 @@ int main(int argc, char **argv) {
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/Linalg/Passes.h.inc"
 
+#define GEN_PASS_REGISTRATION
+#include "mlir/Dialect/SCF/Passes.h.inc"
+
   llvm::InitLLVM y(argc, argv);
 
   mlir::registerDialect<mlir::ONNXOpsDialect>();
   mlir::registerDialect<mlir::KrnlOpsDialect>();
   initOMPasses();
+  initMLIRPasses();
 
   mlir::registerAsmPrinterCLOptions();
   mlir::registerMLIRContextCLOptions();