replace greedy pattern rewrites with partial conversions

better failure messages

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -5,15 +5,12 @@
 #include "mlir/IR/IRMapping.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
-#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/Debug.h"
 
 #include "Common/Common.hpp"
 #include "Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostLegality.hpp"
@@ -91,13 +88,25 @@ void ONNXToSpatialPass::runOnOperation() {
   ModuleOp moduleOp = getOperation();
   MLIRContext* ctx = &getContext();
 
+  ConversionTarget preTarget(*ctx);
+  preTarget.addLegalDialect<spatial::SpatialDialect,
+                            ONNXDialect,
+                            tensor::TensorDialect,
+                            arith::ArithDialect,
+                            scf::SCFDialect>();
+  preTarget.addIllegalOp<ONNXConstantOp, ONNXMatMulOp, ONNXFlattenOp>();
+
   RewritePatternSet prePatterns(ctx);
   populatePrePatterns(prePatterns, ctx);
-  if (failed(applyPatternsGreedily(moduleOp, std::move(prePatterns))))
-    moduleOp.emitWarning("failed to apply ONNX-to-Spatial pre-patterns; continuing");
+  if (failed(applyPartialConversion(moduleOp, preTarget, std::move(prePatterns)))) {
+    moduleOp.emitError("failed to apply ONNX-to-Spatial pre-rewrites");
+    signalPassFailure();
+    return;
+  }
 
   auto entryFunc = getPimEntryFunc(moduleOp);
   if (failed(entryFunc)) {
+    moduleOp.emitError("failed to locate the PIM entry function during ONNX-to-Spatial lowering");
     signalPassFailure();
     return;
   }
@@ -130,13 +139,24 @@ void ONNXToSpatialPass::runOnOperation() {
   RewritePatternSet conversionPatterns(ctx);
   populateConversionPatterns(conversionPatterns, ctx);
   if (failed(applyPartialConversion(moduleOp, target, std::move(conversionPatterns)))) {
+    moduleOp.emitError("failed to convert required ONNX ops to Spatial ops");
     signalPassFailure();
     return;
   }
 
+  ConversionTarget earlyPostTarget(*ctx);
+  earlyPostTarget.addLegalDialect<spatial::SpatialDialect,
+                                  ONNXDialect,
+                                  tensor::TensorDialect,
+                                  arith::ArithDialect,
+                                  scf::SCFDialect>();
+  earlyPostTarget.addDynamicallyLegalOp<spatial::SpatComputeBatch>(
+    [](spatial::SpatComputeBatch batchOp) { return !requiresEarlyPostRewrite(batchOp); });
+
   RewritePatternSet earlyPostPatterns(ctx);
   populateEarlyPostPatterns(earlyPostPatterns, ctx);
-  if (failed(applyPatternsGreedily(*entryFunc, std::move(earlyPostPatterns)))) {
+  if (failed(applyPartialConversion(*entryFunc, earlyPostTarget, std::move(earlyPostPatterns)))) {
+    moduleOp.emitError("failed to normalize single-lane spat.compute_batch ops before core assignment checks");
     signalPassFailure();
     return;
   }
@@ -162,14 +182,27 @@ void ONNXToSpatialPass::runOnOperation() {
 
   annotateWeightsConstants(*entryFunc);
 
+  ConversionTarget postTarget(*ctx);
+  postTarget.addLegalDialect<spatial::SpatialDialect,
+                             ONNXDialect,
+                             tensor::TensorDialect,
+                             arith::ArithDialect,
+                             scf::SCFDialect>();
+  postTarget.addDynamicallyLegalOp<spatial::SpatCompute>(
+    [](spatial::SpatCompute computeOp) { return !requiresPostRewrite(computeOp); });
+  postTarget.addDynamicallyLegalOp<spatial::SpatComputeBatch>(
+    [](spatial::SpatComputeBatch computeOp) { return !requiresPostRewrite(computeOp); });
+
   RewritePatternSet postPatterns(ctx);
   populatePostPatterns(postPatterns, ctx);
-  if (failed(applyPatternsGreedily(*entryFunc, std::move(postPatterns)))) {
+  if (failed(applyPartialConversion(*entryFunc, postTarget, std::move(postPatterns)))) {
+    moduleOp.emitError("failed to normalize weight-like Spatial compute operands before Spatial-to-PIM lowering");
     signalPassFailure();
     return;
   }
 
   if (failed(verifyONNXToSpatialHostLegality(*entryFunc))) {
+    moduleOp.emitError("ONNX-to-Spatial host legality verification failed");
     signalPassFailure();
     return;
   }

src/PIM/Conversion/ONNXToSpatial/PostPatterns.cpp

@@ -31,6 +31,21 @@ static bool isDirectConstantValue(Value value) {
   return isa_and_nonnull<arith::ConstantOp, ONNXConstantOp>(value.getDefiningOp());
 }
 
+template <typename ComputeOpTy>
+static bool hasPromotableWeightLikeInputs(ComputeOpTy compute) {
+  Block& block = compute.getBody().front();
+  for (auto [inputIdx, input] : llvm::enumerate(compute.getInputs())) {
+    if (inputIdx >= block.getNumArguments())
+      continue;
+    if (!isWeightLikeComputeOperand(input))
+      continue;
+    if (isDirectConstantValue(input) && !canPromoteInputBlockArgument(block.getArgument(inputIdx)))
+      continue;
+    return true;
+  }
+  return false;
+}
+
 // Collapses one-lane batches so later phases do not carry batch-only structure unnecessarily.
 struct FoldSingleLaneComputeBatchPattern : OpRewritePattern<spatial::SpatComputeBatch> {
   using OpRewritePattern<spatial::SpatComputeBatch>::OpRewritePattern;
@@ -262,4 +277,10 @@ void annotateWeightsConstants(func::FuncOp funcOp) {
   });
 }
 
+bool requiresEarlyPostRewrite(spatial::SpatComputeBatch batchOp) { return batchOp.getLaneCount() == 1; }
+
+bool requiresPostRewrite(spatial::SpatCompute computeOp) { return hasPromotableWeightLikeInputs(computeOp); }
+
+bool requiresPostRewrite(spatial::SpatComputeBatch computeOp) { return hasPromotableWeightLikeInputs(computeOp); }
+
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/PostPatterns.hpp

@@ -3,8 +3,16 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/IR/MLIRContext.h"
 
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
 namespace onnx_mlir {
 
+bool requiresEarlyPostRewrite(spatial::SpatComputeBatch batchOp);
+
+bool requiresPostRewrite(spatial::SpatCompute computeOp);
+
+bool requiresPostRewrite(spatial::SpatComputeBatch computeOp);
+
 void populateEarlyPostPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 
 void populatePostPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);