fix matmul rewriting/lowering

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

src/PIM/Pass/PimCodegen/VerificationPass.cpp

@@ -7,6 +7,7 @@
 #include "llvm/ADT/STLExtras.h"
 
 #include "src/Accelerators/PIM/Common/IR/SubviewUtils.hpp"
+#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -152,14 +153,15 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
 
   void runOnOperation() override {
     ModuleOp moduleOp = getOperation();
-    bool hasFailure = false;
+    pim::CappedDiagnosticReporter diagnostics;
 
     moduleOp.walk([&](Operation* op) {
       if (op->getDialect()->getNamespace() != "spat")
         return;
 
-      op->emitError("illegal Spatial operation reached PIM codegen verification");
-      hasFailure = true;
+      diagnostics.report(op, [](Operation* illegalOp) {
+        illegalOp->emitError("illegal Spatial operation reached PIM codegen verification");
+      });
     });
 
     for (func::FuncOp funcOp : moduleOp.getOps<func::FuncOp>()) {
@@ -168,36 +170,36 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
 
       for (Operation& op : funcOp.getBody().front().getOperations()) {
         if (auto coreOp = dyn_cast<pim::PimCoreOp>(&op)) {
-          if (failed(verifyCoreWeights(moduleOp, coreOp)) || failed(verifyCoreOperands(coreOp)))
-            hasFailure = true;
+          (void) verifyCoreWeights(moduleOp, coreOp, diagnostics);
+          (void) verifyCoreOperands(coreOp, diagnostics);
           continue;
         }
 
         if (auto coreBatchOp = dyn_cast<pim::PimCoreBatchOp>(&op)) {
-          if (failed(verifyCoreWeights(moduleOp, coreBatchOp)) || failed(verifyCoreOperands(coreBatchOp)))
-            hasFailure = true;
+          (void) verifyCoreWeights(moduleOp, coreBatchOp, diagnostics);
+          (void) verifyCoreOperands(coreBatchOp, diagnostics);
           continue;
         }
 
         if (auto returnOp = dyn_cast<func::ReturnOp>(&op)) {
-          if (failed(verifyReturnOp(returnOp)))
-            hasFailure = true;
+          (void) verifyReturnOp(returnOp, diagnostics);
           continue;
         }
 
         if (!isAddressOnlyHostOp(&op)) {
-          op.emitOpError("illegal host-side runtime op remains after PIM bufferization; "
-                         "fold it to constants or lower it into pim.core");
-          hasFailure = true;
+          diagnostics.report(&op, [](Operation* illegalOp) {
+            illegalOp->emitOpError("illegal host-side runtime op remains after PIM bufferization; "
+                                   "fold it to constants or lower it into pim.core");
+          });
           continue;
         }
 
-        if (failed(verifyAddressOnlyHostOp(&op)))
-          hasFailure = true;
+        (void) verifyAddressOnlyHostOp(&op, diagnostics);
       }
     }
 
-    if (hasFailure) {
+    if (diagnostics.hasFailure()) {
+      diagnostics.emitSuppressedSummary(moduleOp, "verification failures");
       moduleOp.emitError("PIM codegen verification failed; see diagnostics above");
       signalPassFailure();
     }
@@ -205,14 +207,19 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
 
 private:
   template <typename CoreOpTy>
-  static LogicalResult verifyCoreWeights(ModuleOp moduleOp, CoreOpTy coreOp) {
+  static LogicalResult
+  verifyCoreWeights(ModuleOp moduleOp, CoreOpTy coreOp, pim::CappedDiagnosticReporter& diagnostics) {
     bool hasFailure = false;
-    for (auto [weightIndex, weight] : llvm::enumerate(coreOp.getWeights())) {
+    for (auto it : llvm::enumerate(coreOp.getWeights())) {
+      size_t weightIndex = it.index();
+      Value weight = it.value();
       auto getGlobalOp = weight.template getDefiningOp<memref::GetGlobalOp>();
       if (!getGlobalOp && !isConstantGlobalView(weight)) {
-        coreOp.emitOpError() << "weight #" << weightIndex
-                             << " must be materialized as a constant memref.global or a static view of one before JSON "
-                                "codegen";
+        diagnostics.report(coreOp.getOperation(), [&](Operation*) {
+          coreOp.emitOpError() << "weight #" << weightIndex
+                               << " must be materialized as a constant memref.global or a static view of one before "
+                                  "JSON codegen";
+        });
         hasFailure = true;
         continue;
       }
@@ -222,14 +229,18 @@ private:
 
       auto globalOp = lookupGlobalForGetGlobal(moduleOp, getGlobalOp);
       if (!globalOp) {
-        coreOp.emitOpError() << "weight #" << weightIndex << " references an unknown memref.global";
+        diagnostics.report(coreOp.getOperation(), [&](Operation*) {
+          coreOp.emitOpError() << "weight #" << weightIndex << " references an unknown memref.global";
+        });
         hasFailure = true;
         continue;
       }
 
       if (!globalOp.getConstant() || !globalOp.getInitialValue()) {
-        coreOp.emitOpError() << "weight #" << weightIndex
-                             << " must come from a constant memref.global with an initial value";
+        diagnostics.report(coreOp.getOperation(), [&](Operation*) {
+          coreOp.emitOpError() << "weight #" << weightIndex
+                               << " must come from a constant memref.global with an initial value";
+        });
         hasFailure = true;
       }
     }
@@ -237,11 +248,15 @@ private:
     return success(!hasFailure);
   }
 
-  static LogicalResult verifyReturnOp(func::ReturnOp returnOp) {
+  static LogicalResult verifyReturnOp(func::ReturnOp returnOp, pim::CappedDiagnosticReporter& diagnostics) {
     bool hasFailure = false;
-    for (auto [resultIndex, operand] : llvm::enumerate(returnOp.getOperands())) {
+    for (auto it : llvm::enumerate(returnOp.getOperands())) {
+      size_t resultIndex = it.index();
+      Value operand = it.value();
       if (!isCodegenAddressableValue(operand)) {
-        returnOp.emitOpError() << "result #" << resultIndex << " is not backed by contiguous addressable storage";
+        diagnostics.report(returnOp.getOperation(), [&](Operation*) {
+          returnOp.emitOpError() << "result #" << resultIndex << " is not backed by contiguous addressable storage";
+        });
         hasFailure = true;
       }
     }
@@ -249,38 +264,50 @@ private:
   }
 
   template <typename CoreOpTy>
-  static LogicalResult verifyCoreOperands(CoreOpTy coreOp) {
+  static LogicalResult verifyCoreOperands(CoreOpTy coreOp, pim::CappedDiagnosticReporter& diagnostics) {
     return walkPimCoreBlock(
-      coreOp.getBody().front(), StaticValueKnowledge {}, [](Operation& op, const StaticValueKnowledge& knowledge) {
+      coreOp.getBody().front(), StaticValueKnowledge {}, [&](Operation& op, const StaticValueKnowledge& knowledge) {
         bool hasFailure = false;
         if (!isSupportedCoreInstructionOp(&op)) {
-          op.emitOpError("unsupported executable op reached PIM codegen verification");
+          diagnostics.report(&op, [](Operation* illegalOp) {
+            illegalOp->emitOpError("unsupported executable op reached PIM codegen verification");
+          });
           hasFailure = true;
         }
 
-        for (auto [operandIndex, operand] : llvm::enumerate(op.getOperands())) {
+        for (auto it : llvm::enumerate(op.getOperands())) {
+          size_t operandIndex = it.index();
+          Value operand = it.value();
           if (!isa<BaseMemRefType>(operand.getType()))
             continue;
 
           auto resolvedAddress = resolveContiguousAddress(operand, knowledge);
           if (failed(resolvedAddress)) {
-            op.emitOpError() << "operand #" << operandIndex << " is not backed by contiguous addressable storage";
+            diagnostics.report(&op, [&](Operation* illegalOp) {
+              illegalOp->emitOpError() << "operand #" << operandIndex
+                                       << " is not backed by contiguous addressable storage";
+            });
             hasFailure = true;
             continue;
           }
 
           if (isExplicitHostOperand(&op, operandIndex)) {
             if (!isCodegenAddressableValue(operand, knowledge)) {
-              op.emitOpError() << "host operand #" << operandIndex
-                               << " is not backed by contiguous addressable storage";
+              diagnostics.report(&op, [&](Operation* illegalOp) {
+                illegalOp->emitOpError() << "host operand #" << operandIndex
+                                         << " is not backed by contiguous addressable storage";
+              });
               hasFailure = true;
             }
             continue;
           }
 
           if (!isa<memref::AllocOp>(resolvedAddress->base.getDefiningOp())) {
-            op.emitOpError() << "operand #" << operandIndex
-                             << " must be backed by device-local memory; materialize host values with pim.memcp_hd";
+            diagnostics.report(&op, [&](Operation* illegalOp) {
+              illegalOp->emitOpError() << "operand #" << operandIndex
+                                       << " must be backed by device-local memory; materialize host values with "
+                                          "pim.memcp_hd";
+            });
             hasFailure = true;
           }
         }
@@ -288,18 +315,20 @@ private:
       });
   }
 
-  static LogicalResult verifyAddressOnlyHostOp(Operation* op) {
+  static LogicalResult verifyAddressOnlyHostOp(Operation* op, pim::CappedDiagnosticReporter& diagnostics) {
     if (auto subviewOp = dyn_cast<memref::SubViewOp>(op))
-      return verifyAddressOnlyBase(op, subviewOp.getSource());
+      return verifyAddressOnlyBase(op, subviewOp.getSource(), diagnostics);
     if (auto castOp = dyn_cast<memref::CastOp>(op))
-      return verifyAddressOnlySource(op, castOp.getSource());
+      return verifyAddressOnlySource(op, castOp.getSource(), diagnostics);
     if (auto collapseOp = dyn_cast<memref::CollapseShapeOp>(op))
-      return verifyAddressOnlySource(op, collapseOp.getSrc());
+      return verifyAddressOnlySource(op, collapseOp.getSrc(), diagnostics);
     if (auto expandOp = dyn_cast<memref::ExpandShapeOp>(op))
-      return verifyAddressOnlySource(op, expandOp.getSrc());
+      return verifyAddressOnlySource(op, expandOp.getSrc(), diagnostics);
     if (auto copyOp = dyn_cast<memref::CopyOp>(op)) {
       if (!isBaseAddressableValue(copyOp.getSource()) || !isBaseAddressableValue(copyOp.getTarget())) {
-        op->emitOpError("depends on a value that is not backed by addressable storage");
+        diagnostics.report(op, [](Operation* illegalOp) {
+          illegalOp->emitOpError("depends on a value that is not backed by addressable storage");
+        });
         return failure();
       }
       return success();
@@ -307,19 +336,24 @@ private:
     return success();
   }
 
-  static LogicalResult verifyAddressOnlySource(Operation* op, Value source) {
+  static LogicalResult
+  verifyAddressOnlySource(Operation* op, Value source, pim::CappedDiagnosticReporter& diagnostics) {
     if (isCodegenAddressableValue(source))
       return success();
 
-    op->emitOpError("depends on a value that is not backed by contiguous addressable storage");
+    diagnostics.report(op, [](Operation* illegalOp) {
+      illegalOp->emitOpError("depends on a value that is not backed by contiguous addressable storage");
+    });
     return failure();
   }
 
-  static LogicalResult verifyAddressOnlyBase(Operation* op, Value source) {
+  static LogicalResult verifyAddressOnlyBase(Operation* op, Value source, pim::CappedDiagnosticReporter& diagnostics) {
     if (isBaseAddressableValue(source))
       return success();
 
-    op->emitOpError("depends on a value that is not backed by addressable storage");
+    diagnostics.report(op, [](Operation* illegalOp) {
+      illegalOp->emitOpError("depends on a value that is not backed by addressable storage");
+    });
     return failure();
   }
 };