add constant folding and verification pass for pim host operations

better validation scripts output
big refactors

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -1,3 +1,4 @@
+#include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/Pass/Pass.h"
@@ -10,19 +11,39 @@
 
 #include "Common/PIMCommon.hpp"
 #include "Conversion/ONNXToSpatial/Utils/AnnotateReplication.hpp"
-#include "Math/Conv.hpp"
-#include "ONNXToSpatialPass.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPatterns.hpp"
-#include "src/Accelerators/PIM/Dialect/PIM/PimOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Pass/PimPasses.hpp"
 #include "src/Compiler/CompilerOptions.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
 
 using namespace mlir;
 
 namespace onnx_mlir {
 
-namespace spatial {
+bool haveSameStaticShape(Value lhs, Value rhs);
+
+namespace {
+
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatial.hpp.inc"
+
+struct ONNXToSpatialPass : PassWrapper<ONNXToSpatialPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ONNXToSpatialPass)
+  StringRef getArgument() const override { return "convert-onnx-to-spatial"; }
+  StringRef getDescription() const override { return "Lower ONNX ops to Spatial ops."; }
+
+  ONNXToSpatialPass() = default;
+  ONNXToSpatialPass(const ONNXToSpatialPass& pass) {}
+
+  void runOnOperation() override;
+
+private:
+  void annotateWeightsConstants(func::FuncOp funcOp) const;
+};
+
+} // namespace
 
 void ONNXToSpatialPass::runOnOperation() {
   ModuleOp moduleOp = getOperation();
@@ -40,15 +61,19 @@ void ONNXToSpatialPass::runOnOperation() {
     llvm::dbgs() << "Failed to merge activation patterns, continuing...\n";
 
   IRRewriter rewriter(moduleOp);
-  func::FuncOp funcOp = *moduleOp.getOps<func::FuncOp>().begin();
-  if (annotateReplication(funcOp, rewriter).failed()) {
+  auto entryFunc = getPimEntryFunc(moduleOp);
+  if (failed(entryFunc)) {
+    signalPassFailure();
+    return;
+  }
+  if (annotateReplication(*entryFunc, rewriter).failed()) {
     llvm::dbgs() << "Failed during annotation for replication analysis\n";
     signalPassFailure();
     return;
   }
 
   ConversionTarget target(*ctx);
-  target.addLegalDialect<ONNXDialect, SpatialDialect, tensor::TensorDialect, arith::ArithDialect, tosa::TosaDialect>();
+  target.addLegalDialect<spatial::SpatialDialect, ONNXDialect, tensor::TensorDialect, arith::ArithDialect>();
   target.addIllegalOp<ONNXMatMulOp>();
   target.addIllegalOp<ONNXGemmOp>();
   target.addIllegalOp<ONNXConvOp>();
@@ -62,16 +87,9 @@ void ONNXToSpatialPass::runOnOperation() {
   RewritePatternSet patterns(ctx);
   patterns.add<removeLRNPattern>(ctx);
 
-  if (useExperimentalConvImpl) {
-    populateExperimentalTilingConvOpPattern(patterns, ctx);
-    populateExperimentalPoolingTilingPattern(patterns, ctx);
-    populateGemmToConvConversionPattern(patterns, ctx);
-  }
-  else {
-    populateTilingConvOpPattern(patterns, ctx);
-    populatePoolingTilingPattern(patterns, ctx);
-    populateOnnxGemmOpPatterns(patterns, ctx);
-  }
+  populateConvOpPatterns(patterns, ctx);
+  populatePoolingTilingPattern(patterns, ctx);
+  populateOnnxGemmOpPatterns(patterns, ctx);
 
   populateONNXConcatToTensorConcatPattern(patterns, ctx);
   populateReduceMeanConversionPattern(patterns, ctx);
@@ -84,8 +102,8 @@ void ONNXToSpatialPass::runOnOperation() {
   // Count the number of compute ops and check they do not exceed the core count
   if (coresCount != -1) {
     int computeOpsCount = 0;
-    for (auto& op : funcOp.getFunctionBody().front().getOperations())
-      if (isa<SpatWeightedCompute>(op))
+    for (auto& op : entryFunc->getFunctionBody().front().getOperations())
+      if (isa<spatial::SpatWeightedCompute>(op))
         computeOpsCount++;
 
     if (computeOpsCount > coresCount) {
@@ -102,22 +120,21 @@ void ONNXToSpatialPass::runOnOperation() {
   if (failed(applyPatternsGreedily(moduleOp, std::move(removeUnusedHelperOpsPatterns))))
     llvm::dbgs() << "Failed to remove unused helper ops, continuing...\n";
 
-  annotateWeightsConstants(funcOp);
+  annotateWeightsConstants(*entryFunc);
 
   // Dump to file for debug
   dumpModule(moduleOp, "spatial");
 }
 
 void ONNXToSpatialPass::annotateWeightsConstants(func::FuncOp funcOp) const {
-  MLIRContext* ctx = funcOp.getContext();
   funcOp.walk([&](arith::ConstantOp constantOp) {
     bool isAlwaysWeight =
-      llvm::all_of(constantOp->getUsers(), [](auto user) -> bool { return isa<SpatWeightedCompute>(user); });
+      llvm::all_of(constantOp->getUsers(), [](auto user) -> bool { return isa<spatial::SpatWeightedCompute>(user); });
     if (isAlwaysWeight)
-      constantOp->setAttr("weightAlways", UnitAttr::get(ctx));
+      markWeightAlways(constantOp);
   });
 }
 
-} // namespace spatial
+std::unique_ptr<Pass> createONNXToSpatialPass() { return std::make_unique<ONNXToSpatialPass>(); }
 
 } // namespace onnx_mlir