fix matmul rewriting/lowering

fix reshape lowering
add support for grouped-convolution lowering
quieter verifier with capped error messages

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -6,13 +6,14 @@
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Transforms/Passes.h"
+#include "mlir/Transforms/WalkPatternRewriteDriver.h"
 
 #include "llvm/ADT/SmallVector.h"
 
 #include "Common/Common.hpp"
 #include "Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostFoldability.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostLegality.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PostPatterns.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PrePatterns.hpp"
@@ -84,6 +85,30 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
     returnOp.setOperand(index, computeResult);
 }
 
+static void wrapTopLevelRuntimeTransposes(func::FuncOp funcOp) {
+  IRRewriter rewriter(funcOp.getContext());
+  Block& entryBlock = funcOp.getFunctionBody().front();
+
+  for (Operation& op : llvm::make_early_inc_range(entryBlock)) {
+    auto transposeOp = dyn_cast<ONNXTransposeOp>(&op);
+    if (!transposeOp || isHostFoldableOp(transposeOp))
+      continue;
+
+    // Transpose stays globally legal because constant/view-only cases are
+    // allowed on the host. Any residual runtime transpose must be sunk into
+    // spat.compute before the host legality check.
+    auto resultType = transposeOp.getResult().getType();
+    rewriter.setInsertionPoint(transposeOp);
+    auto computeOp = createSpatCompute<1>(
+      rewriter, transposeOp.getLoc(), TypeRange {resultType}, {}, ValueRange {transposeOp.getData()}, [&](Value input) {
+        Value transposed =
+          ONNXTransposeOp::create(rewriter, transposeOp.getLoc(), resultType, input, transposeOp.getPermAttr());
+        spatial::SpatYieldOp::create(rewriter, transposeOp.getLoc(), transposed);
+      });
+    rewriter.replaceOp(transposeOp, computeOp.getResult(0));
+  }
+}
+
 void ONNXToSpatialPass::runOnOperation() {
   ModuleOp moduleOp = getOperation();
   MLIRContext* ctx = &getContext();
@@ -94,7 +119,7 @@ void ONNXToSpatialPass::runOnOperation() {
                             tensor::TensorDialect,
                             arith::ArithDialect,
                             scf::SCFDialect>();
-  preTarget.addIllegalOp<ONNXConstantOp, ONNXMatMulOp, ONNXFlattenOp>();
+  preTarget.addIllegalOp<ONNXConstantOp, ONNXFlattenOp>();
 
   RewritePatternSet prePatterns(ctx);
   populatePrePatterns(prePatterns, ctx);
@@ -111,6 +136,21 @@ void ONNXToSpatialPass::runOnOperation() {
     return;
   }
 
+  RewritePatternSet matmulPatterns(ctx);
+  populateMatMulRewritePatterns(matmulPatterns, ctx);
+  walkAndApplyPatterns(moduleOp, std::move(matmulPatterns));
+
+  bool hasUnloweredMatMul = false;
+  moduleOp.walk([&](ONNXMatMulOp matmulOp) {
+    hasUnloweredMatMul = true;
+    matmulOp.emitOpError("remaining ONNX MatMul before the required ONNX-to-Spatial conversion");
+  });
+  if (hasUnloweredMatMul) {
+    moduleOp.emitError("failed to lower all ONNX MatMul ops before ONNX-to-Spatial conversion");
+    signalPassFailure();
+    return;
+  }
+
   ConversionTarget target(*ctx);
   target.addLegalDialect<spatial::SpatialDialect,
                          ONNXDialect,
@@ -161,20 +201,6 @@ void ONNXToSpatialPass::runOnOperation() {
     return;
   }
 
-  if (coresCount != -1) {
-    int computeOpsCount = 0;
-    for (Operation& op : entryFunc->getFunctionBody().front().getOperations())
-      if (isa<spatial::SpatCompute>(op))
-        computeOpsCount++;
-
-    if (computeOpsCount > coresCount) {
-      entryFunc->emitError() << "number of compute ops (" << computeOpsCount << ") exceeds the core count ("
-                             << coresCount << ")";
-      signalPassFailure();
-      return;
-    }
-  }
-
   PassManager cleanupPM(ctx);
   cleanupPM.addPass(createCanonicalizerPass());
   if (failed(cleanupPM.run(moduleOp)))
@@ -201,6 +227,8 @@ void ONNXToSpatialPass::runOnOperation() {
     return;
   }
 
+  wrapTopLevelRuntimeTransposes(*entryFunc);
+
   if (failed(verifyONNXToSpatialHostLegality(*entryFunc))) {
     moduleOp.emitError("ONNX-to-Spatial host legality verification failed");
     signalPassFailure();