add const prop to hybrid pass

Signed-off-by: Soren Lassen <[email protected]>

src/Transform/ONNX/CMakeLists.txt

@@ -72,6 +72,7 @@ add_onnx_mlir_library(OMHybridTransform
   LINK_LIBS PUBLIC
   OMONNXOps
   OMShapeInferenceOpInterface
+  OMONNXRewrite
   MLIRPass
   MLIRTransforms
   OMShapeInference

src/Transform/ONNX/ConstProp.cpp

@@ -13,6 +13,9 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "src/Transform/ONNX/ConstProp.hpp"
+#include "src/Pass/Passes.hpp"
+
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
@@ -29,7 +32,6 @@
 #include "src/Dialect/ONNX/ONNXOps/OpHelper.hpp"
 #include "src/Dialect/ONNX/ONNXOps/ShapeHelper.hpp"
 #include "src/Dialect/ONNX/OnnxElementsAttrBuilder.hpp"
-#include "src/Pass/Passes.hpp"
 #include "src/Support/TypeUtilities.hpp"
 
 #include <math.h>
@@ -1011,6 +1013,12 @@ class SplitOfConst : public OpRewritePattern<ONNXSplitOp> {
   }
 };
 
+void getPatterns(RewritePatternSet &patterns) {
+  populateWithGenerated(patterns);
+  if (isNotDisabled("SplitOfConst"))
+    patterns.insert<SplitOfConst>(patterns.getContext());
+}
+
 //===----------------------------------------------------------------------===//
 // Code to manage the pass.
 //===----------------------------------------------------------------------===//
@@ -1034,15 +1042,17 @@ void ConstPropONNXToONNXPass::runOnOperation() {
   MLIRContext *context = &getContext();
 
   RewritePatternSet patterns(context);
-  populateWithGenerated(patterns);
-  if (isNotDisabled("SplitOfConst"))
-    patterns.insert<SplitOfConst>(context);
+  getPatterns(patterns);
   if (failed(applyPatternsAndFoldGreedily(function, std::move(patterns))))
     signalPassFailure();
 }
 
 } // end anonymous namespace.
 
+void onnx_mlir::getConstPropPatterns(RewritePatternSet &patterns) {
+  getPatterns(patterns);
+}
+
 void onnx_mlir::configureConstPropONNXToONNXPass(
     int expansionBound, ArrayRef<std::string> disabledPatterns) {
   ConstPropONNXToONNXPassConfiguration::expansionBound = expansionBound;

src/Transform/ONNX/ConstProp.hpp

@@ -0,0 +1,13 @@
+/*
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include "mlir/IR/PatternMatch.h"
+
+namespace onnx_mlir {
+
+void getConstPropPatterns(mlir::RewritePatternSet &patterns);
+
+}

src/Transform/ONNX/ONNXHybridTransformPass.cpp

@@ -5,9 +5,9 @@
 //===------------------ ONNXHybridTransformPass.cpp -----------------------===//
 //
 // Hybrid ONNX transformation pass that combines conversion patterns for
-// shape inference and canonicalization.
+// shape inference and canonicalization and constant propagation.
 //
-// TODO: add constant propagation and decomposition
+// TODO: add decomposition
 //
 //===----------------------------------------------------------------------===//
 
@@ -17,6 +17,7 @@
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 #include "src/Interface/ShapeInferenceOpInterface.hpp"
 #include "src/Pass/Passes.hpp"
+#include "src/Transform/ONNX/ConstProp.hpp"
 #include "src/Transform/ONNX/ShapeInference.hpp"
 
 using namespace mlir;
@@ -54,7 +55,9 @@ struct ONNXHybridTransformPass
     for (RegisteredOperationName op : context->getRegisteredOperations())
       op.getCanonicalizationPatterns(cumulativePatterns, context);
 
-    // TODO: constant propagation, decomposition
+    getConstPropPatterns(cumulativePatterns);
+
+    // TODO: decomposition
 
     patterns = FrozenRewritePatternSet(std::move(cumulativePatterns));
     return success();