replace greedy pattern rewrites with partial conversions

better failure messages

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

@@ -144,6 +144,7 @@ void SpatialToPimPass::runOnOperation() {
 
   auto entryFunc = getPimEntryFunc(moduleOp);
   if (failed(entryFunc)) {
+    moduleOp.emitError("failed to locate the PIM entry function during Spatial-to-PIM lowering");
     signalPassFailure();
     return;
   }
@@ -169,26 +170,22 @@ void SpatialToPimPass::runOnOperation() {
                     spatial::SpatChannelSendTensorBatchOp,
                     spatial::SpatExtractRowsOp>();
 
-  {
-    RewritePatternSet patterns(ctx);
-    populateWithGenerated(patterns);
-
-    if (failed(applyPartialConversion(moduleOp, target, std::move(patterns)))) {
-      signalPassFailure();
-      return;
-    }
+  RewritePatternSet initialPatterns(ctx);
+  populateWithGenerated(initialPatterns);
+  if (failed(applyPartialConversion(moduleOp, target, std::move(initialPatterns)))) {
+    moduleOp.emitError("failed to lower required Spatial ops to the initial PIM form");
+    signalPassFailure();
+    return;
   }
 
-  {
-    RewritePatternSet patterns(ctx);
-    populateGlobalTensorMaterializationPatterns(patterns);
-
-    walkAndApplyPatterns(moduleOp, std::move(patterns));
-  }
+  RewritePatternSet globalTensorPatterns(ctx);
+  populateGlobalTensorMaterializationPatterns(globalTensorPatterns);
+  walkAndApplyPatterns(moduleOp, std::move(globalTensorPatterns));
   auto returnOp = cast<func::ReturnOp>(funcOp.front().getTerminator());
 
   addReturnOutputBuffers(returnOp, rewriter, outputTensors);
   if (failed(allocateAndInitializeCoreLocalVariables(funcOp, rewriter))) {
+    funcOp.emitOpError("failed to allocate or initialize core-local tensors during Spatial-to-PIM lowering");
     signalPassFailure();
     return;
   }
@@ -197,6 +194,7 @@ void SpatialToPimPass::runOnOperation() {
   for (auto computeOp : funcOp.getOps<spatial::SpatCompute>()) {
     markOpToRemove(computeOp);
     if (failed(lowerComputeOp(computeOp, coreLoweringState, rewriter))) {
+      computeOp.emitOpError("failed to lower spat.compute to pim.core");
       signalPassFailure();
       return;
     }
@@ -205,17 +203,16 @@ void SpatialToPimPass::runOnOperation() {
   for (auto computeBatchOp : funcOp.getOps<spatial::SpatComputeBatch>()) {
     markOpToRemove(computeBatchOp);
     if (failed(lowerComputeBatchOp(computeBatchOp, coreLoweringState, rewriter))) {
+      computeBatchOp.emitOpError("failed to lower spat.compute_batch to pim.core_batch");
       signalPassFailure();
       return;
     }
   }
 
-  {
-    RewritePatternSet patterns(ctx);
-    populateTensorPackingPatterns(patterns);
-    walkAndApplyPatterns(funcOp, std::move(patterns));
-    eraseUnusedTensorPackingOps(funcOp, rewriter);
-  }
+  RewritePatternSet initialTensorPackingPatterns(ctx);
+  populateTensorPackingPatterns(initialTensorPackingPatterns);
+  walkAndApplyPatterns(funcOp, std::move(initialTensorPackingPatterns));
+  eraseUnusedTensorPackingOps(funcOp, rewriter);
 
   SmallVector<spatial::SpatChannelReceiveOp> receiveOps;
   for (auto op : funcOp.getOps<spatial::SpatChannelReceiveOp>())
@@ -229,27 +226,27 @@ void SpatialToPimPass::runOnOperation() {
     }
   }
 
-  {
-    RewritePatternSet coreBodyPatterns(ctx);
-    populateWithGenerated(coreBodyPatterns);
-    FrozenRewritePatternSet frozenCoreBodyPatterns(std::move(coreBodyPatterns));
+  RewritePatternSet coreBodyPatterns(ctx);
+  populateWithGenerated(coreBodyPatterns);
+  FrozenRewritePatternSet frozenCoreBodyPatterns(std::move(coreBodyPatterns));
 
-    SmallVector<pim::PimCoreOp> coreOps;
-    funcOp.walk([&](pim::PimCoreOp coreOp) { coreOps.push_back(coreOp); });
-    for (auto coreOp : coreOps) {
-      if (failed(applyFullConversion(coreOp.getOperation(), target, frozenCoreBodyPatterns))) {
-        signalPassFailure();
-        return;
-      }
+  SmallVector<pim::PimCoreOp> coreOps;
+  funcOp.walk([&](pim::PimCoreOp coreOp) { coreOps.push_back(coreOp); });
+  for (auto coreOp : coreOps) {
+    if (failed(applyFullConversion(coreOp.getOperation(), target, frozenCoreBodyPatterns))) {
+      coreOp.emitOpError("failed to convert nested Spatial ops inside pim.core");
+      signalPassFailure();
+      return;
     }
+  }
 
-    SmallVector<pim::PimCoreBatchOp> coreBatchOps;
-    funcOp.walk([&](pim::PimCoreBatchOp coreBatchOp) { coreBatchOps.push_back(coreBatchOp); });
-    for (auto coreBatchOp : coreBatchOps) {
-      if (failed(applyFullConversion(coreBatchOp.getOperation(), target, frozenCoreBodyPatterns))) {
-        signalPassFailure();
-        return;
-      }
+  SmallVector<pim::PimCoreBatchOp> coreBatchOps;
+  funcOp.walk([&](pim::PimCoreBatchOp coreBatchOp) { coreBatchOps.push_back(coreBatchOp); });
+  for (auto coreBatchOp : coreBatchOps) {
+    if (failed(applyFullConversion(coreBatchOp.getOperation(), target, frozenCoreBodyPatterns))) {
+      coreBatchOp.emitOpError("failed to convert nested Spatial ops inside pim.core_batch");
+      signalPassFailure();
+      return;
     }
   }
 
@@ -259,44 +256,43 @@ void SpatialToPimPass::runOnOperation() {
 
   SmallVector<Operation*> pendingRemovals(operationsToRemove.begin(), operationsToRemove.end());
   if (failed(erasePendingOps(pendingRemovals, rewriter))) {
+    funcOp.emitOpError("failed to erase obsolete Spatial ops after lowering to PIM");
     signalPassFailure();
     return;
   }
 
-  {
-    RewritePatternSet patterns(ctx);
-    populateTensorPackingPatterns(patterns);
-    walkAndApplyPatterns(funcOp, std::move(patterns));
-    eraseUnusedTensorPackingOps(funcOp, rewriter);
-  }
+  RewritePatternSet finalTensorPackingPatterns(ctx);
+  populateTensorPackingPatterns(finalTensorPackingPatterns);
+  walkAndApplyPatterns(funcOp, std::move(finalTensorPackingPatterns));
+  eraseUnusedTensorPackingOps(funcOp, rewriter);
 
-  {
-    ConversionTarget communicationTarget(*ctx);
-    communicationTarget.addLegalDialect<PimDialect,
-                                        tensor::TensorDialect,
-                                        arith::ArithDialect,
-                                        bufferization::BufferizationDialect,
-                                        func::FuncDialect,
-                                        memref::MemRefDialect,
-                                        scf::SCFDialect,
-                                        BuiltinDialect>();
-    communicationTarget.addLegalOp<ModuleOp>();
-    communicationTarget.addIllegalOp<spatial::SpatConcatOp,
-                                     spatial::SpatChannelReceiveOp,
-                                     spatial::SpatChannelReceiveTensorOp,
-                                     spatial::SpatChannelSendOp,
-                                     spatial::SpatChannelSendTensorOp,
-                                     spatial::SpatExtractRowsOp>();
+  ConversionTarget communicationTarget(*ctx);
+  communicationTarget.addLegalDialect<PimDialect,
+                                      tensor::TensorDialect,
+                                      arith::ArithDialect,
+                                      bufferization::BufferizationDialect,
+                                      func::FuncDialect,
+                                      memref::MemRefDialect,
+                                      scf::SCFDialect,
+                                      BuiltinDialect>();
+  communicationTarget.addLegalOp<ModuleOp>();
+  communicationTarget.addIllegalOp<spatial::SpatConcatOp,
+                                   spatial::SpatChannelReceiveOp,
+                                   spatial::SpatChannelReceiveTensorOp,
+                                   spatial::SpatChannelSendOp,
+                                   spatial::SpatChannelSendTensorOp,
+                                   spatial::SpatExtractRowsOp>();
 
-    RewritePatternSet communicationPatterns(ctx);
-    populateChannelLoweringPatterns(communicationPatterns);
-    if (failed(applyFullConversion(funcOp, communicationTarget, std::move(communicationPatterns)))) {
-      signalPassFailure();
-      return;
-    }
+  RewritePatternSet communicationPatterns(ctx);
+  populateChannelLoweringPatterns(communicationPatterns);
+  if (failed(applyFullConversion(funcOp, communicationTarget, std::move(communicationPatterns)))) {
+    funcOp.emitOpError("failed to lower Spatial communication ops to PIM communication ops");
+    signalPassFailure();
+    return;
   }
 
   if (failed(verifySpatialToPimBoundary(moduleOp))) {
+    moduleOp.emitError("Spatial-to-PIM boundary verification failed");
     signalPassFailure();
     return;
   }