faster pim VerificationPass.cpp and pim code emission

src/PIM/Pass/PimCodegen/VerificationPass.cpp

@@ -6,10 +6,10 @@
 
 #include "llvm/ADT/STLExtras.h"
 
+#include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/SubviewUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
-#include "src/Accelerators/PIM/Compiler/PimBatchEmission.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
@@ -173,6 +173,106 @@ static bool isSupportedCoreInstructionOp(Operation* op) {
              memref::GetGlobalOp>(op);
 }
 
+static FailureOr<ShapedType> getStaticByteSizedShapedType(Type type) {
+  auto shapedType = dyn_cast<ShapedType>(type);
+  if (!shapedType || !shapedType.hasStaticShape())
+    return failure();
+
+  int64_t elementBits = shapedType.getElementTypeBitWidth();
+  if (elementBits <= 0 || elementBits % 8 != 0)
+    return failure();
+
+  return shapedType;
+}
+
+static LogicalResult verifyBatchOpSemantics(Operation& op,
+                                            const StaticValueKnowledge& knowledge,
+                                            pim::CappedDiagnosticReporter& diagnostics) {
+  bool hasFailure = false;
+  auto reportFailure = [&](auto emitDiagnostic) {
+    diagnostics.report(&op, [&](Operation* illegalOp) { emitDiagnostic(illegalOp); });
+    hasFailure = true;
+  };
+
+  if (auto memcpHdBatchOp = dyn_cast<pim::PimMemCopyHostToDevBatchOp>(op)) {
+    if (!isCodegenAddressableValue(memcpHdBatchOp.getHostSource(), knowledge)) {
+      reportFailure([](Operation* illegalOp) {
+        illegalOp->emitOpError("host operand #1 is not backed by contiguous addressable storage");
+      });
+    }
+    return success(!hasFailure);
+  }
+
+  if (auto sendBatchOp = dyn_cast<pim::PimSendBatchOp>(op)) {
+    if (sendBatchOp.getTargetCoreIds().size() != static_cast<size_t>(sendBatchOp->getParentOfType<pim::PimCoreBatchOp>()
+                                                                        .getLaneCount())) {
+      reportFailure([](Operation* illegalOp) {
+        illegalOp->emitOpError("targetCoreIds size must match parent laneCount");
+      });
+    }
+    return success(!hasFailure);
+  }
+
+  if (auto receiveBatchOp = dyn_cast<pim::PimReceiveBatchOp>(op)) {
+    if (receiveBatchOp.getSourceCoreIds().size()
+        != static_cast<size_t>(receiveBatchOp->getParentOfType<pim::PimCoreBatchOp>().getLaneCount())) {
+      reportFailure([](Operation* illegalOp) {
+        illegalOp->emitOpError("sourceCoreIds size must match parent laneCount");
+      });
+    }
+    return success(!hasFailure);
+  }
+
+  auto verifyTensorBatchCommunication = [&](Value tensorValue, ArrayRef<int32_t> coreIds, StringRef kind) {
+    if (coreIds.empty()) {
+      reportFailure([&](Operation* illegalOp) { illegalOp->emitOpError() << kind << " must carry at least one chunk"; });
+      return;
+    }
+
+    auto parentBatchOp = op.getParentOfType<pim::PimCoreBatchOp>();
+    int32_t laneCount = parentBatchOp.getLaneCount();
+    if (laneCount <= 0) {
+      reportFailure([&](Operation* illegalOp) {
+        illegalOp->emitOpError() << kind << " requires a positive parent laneCount";
+      });
+      return;
+    }
+    if (coreIds.size() % static_cast<size_t>(laneCount) != 0) {
+      reportFailure([&](Operation* illegalOp) {
+        illegalOp->emitOpError() << kind << " core id count must be divisible by the parent laneCount";
+      });
+      return;
+    }
+
+    auto shapedType = getStaticByteSizedShapedType(tensorValue.getType());
+    if (failed(shapedType)) {
+      reportFailure([&](Operation* illegalOp) {
+        illegalOp->emitOpError() << kind << " requires a static shaped tensor or memref with byte-sized elements";
+      });
+      return;
+    }
+
+    int64_t chunkCount = static_cast<int64_t>(coreIds.size()) / laneCount;
+    int64_t totalBytes = (*shapedType).getNumElements() * (*shapedType).getElementTypeBitWidth() / 8;
+    if (totalBytes % chunkCount != 0) {
+      reportFailure([&](Operation* illegalOp) {
+        illegalOp->emitOpError() << kind << " tensor byte size must be divisible by the chunk count per lane";
+      });
+    }
+  };
+
+  if (auto sendTensorBatchOp = dyn_cast<pim::PimSendTensorBatchOp>(op))
+    verifyTensorBatchCommunication(sendTensorBatchOp.getInput(),
+                                   sendTensorBatchOp.getTargetCoreIds(),
+                                   "send_tensor_batch");
+  else if (auto receiveTensorBatchOp = dyn_cast<pim::PimReceiveTensorBatchOp>(op))
+    verifyTensorBatchCommunication(receiveTensorBatchOp.getOutput(),
+                                   receiveTensorBatchOp.getSourceCoreIds(),
+                                   "receive_tensor_batch");
+
+  return success(!hasFailure);
+}
+
 struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(VerificationPass)
 
@@ -204,16 +304,24 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
       for (Operation& op : funcOp.getBody().front().getOperations()) {
         if (auto coreOp = dyn_cast<pim::PimCoreOp>(&op)) {
           (void) verifyCoreWeights(moduleOp, coreOp, diagnostics);
-          (void) verifyCoreOperands(coreOp, diagnostics);
+          StaticValueKnowledge knowledge;
+          (void) verifyCoreLikeOperands(coreOp, knowledge, diagnostics);
           continue;
         }
 
         if (auto coreBatchOp = dyn_cast<pim::PimCoreBatchOp>(&op)) {
           (void) verifyCoreWeights(moduleOp, coreBatchOp, diagnostics);
-          for (unsigned lane = 0; lane < static_cast<unsigned>(coreBatchOp.getLaneCount()); ++lane)
-            (void) withScalarCoreFromBatchLane(coreBatchOp, lane, [&](pim::PimCoreOp scalarCore) {
-              return verifyCoreOperands(scalarCore, diagnostics);
-            });
+          llvm::SmallVector<unsigned, 2> lanes;
+          lanes.push_back(0);
+          if (coreBatchOp.getLaneCount() > 1)
+            lanes.push_back(static_cast<unsigned>(coreBatchOp.getLaneCount() - 1));
+          for (unsigned lane : lanes) {
+            StaticValueKnowledge knowledge;
+            knowledge.indexValues[coreBatchOp.getLaneArgument()] = lane;
+            for (unsigned i = 0; i < coreBatchOp.getInputs().size(); ++i)
+              knowledge.aliases[coreBatchOp.getInputArgument(i)] = coreBatchOp.getInputs()[i];
+            (void) verifyCoreLikeOperands(coreBatchOp, knowledge, diagnostics);
+          }
           continue;
         }
 
@@ -299,10 +407,13 @@ private:
     return success(!hasFailure);
   }
 
-  template <typename CoreOpTy>
-  static LogicalResult verifyCoreOperands(CoreOpTy coreOp, pim::CappedDiagnosticReporter& diagnostics) {
-    return walkPimCoreBlock(
-      coreOp.getBody().front(), StaticValueKnowledge {}, [&](Operation& op, const StaticValueKnowledge& knowledge) {
+  template <typename CoreLikeOpTy>
+  static LogicalResult verifyCoreLikeOperands(CoreLikeOpTy coreLikeOp,
+                                              const StaticValueKnowledge& initialKnowledge,
+                                              pim::CappedDiagnosticReporter& diagnostics) {
+    return walkPimCoreBlockStructurally(coreLikeOp.getBody().front(),
+                                        initialKnowledge,
+                                        [&](Operation& op, const StaticValueKnowledge& knowledge) {
         bool hasFailure = false;
         if (!isSupportedCoreInstructionOp(&op)) {
           diagnostics.report(&op, [](Operation* illegalOp) {
@@ -370,6 +481,9 @@ private:
             hasFailure = true;
           }
         }
+
+        if (failed(verifyBatchOpSemantics(op, knowledge, diagnostics)))
+          hasFailure = true;
         return success(!hasFailure);
       });
   }