refactor

src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp

@@ -29,100 +29,6 @@ static bool isUsedOnlyAsExplicitHostOperand(Value value) {
   });
 }
 
-static bool isMaterializableExternalTensorOp(Operation* op) {
-  return isa<spatial::SpatChannelReceiveOp,
-             spatial::SpatExtractRowsOp,
-             tensor::ExtractSliceOp,
-             tensor::ExpandShapeOp,
-             tensor::CollapseShapeOp>(op);
-}
-
-//TODO REMOVE THIS UGLY FIX 
-//TODO: Remove this helper once compute_batch external tensor captures are
-// fixed at the producer side.
-//
-// This function is a temporary SpatialToPim repair path. It clones selected
-// external tensor producers, such as channel_receive and tensor view/slice ops,
-// into the new pim.core_batch body when the old spat.compute_batch body refers
-// to tensor values defined outside the batch.
-//
-// The real invariant should be stronger:
-//
-//   A spat.compute_batch body must not capture external tensor values.
-//   Every tensor used inside the body must be either:
-//     - a compute_batch block argument,
-//     - defined inside the compute_batch body,
-//     - or a legal constant-like value.
-//
-// If this invariant is violated, the responsible producer, most likely merge
-// schedule materialization, should emit verifier-clean Spatial IR instead of
-// relying on SpatialToPim to clone external producer chains later.
-//
-// After that producer-side fix:
-//   1. remove isMaterializableExternalTensorOp,
-//   2. remove materializeExternalTensorValue,
-//   3. make lowerComputeBatchOp emit a hard diagnostic for any unmapped external
-//      tensor operand,
-//   4. keep/strengthen the Spatial verifier so the invalid capture is rejected
-//      before SpatialToPim.
-//
-// Be careful not to replace every external tensor capture with a normal
-// compute_batch input blindly: host-backed tensors and explicit inter-core
-// communication have different semantics. In particular, channel_receive-like
-// values should be materialized through the communication model, not silently
-// treated as host inputs.
-static FailureOr<Value> materializeExternalTensorValue(IRRewriter& rewriter,
-                                                       Location loc,
-                                                       Block& oldBlock,
-                                                       Value value,
-                                                       IRMapping& mapper) {
-  if (mapper.contains(value))
-    return mapper.lookup(value);
-
-  if (!isa<TensorType>(value.getType()))
-    return value;
-
-  Operation* definingOp = value.getDefiningOp();
-  if (!definingOp || definingOp->hasTrait<OpTrait::ConstantLike>())
-    return failure();
-
-  if (definingOp->getBlock() == &oldBlock)
-    return failure();
-
-  if (!isMaterializableExternalTensorOp(definingOp))
-    return failure();
-
-  for (Value operand : definingOp->getOperands()) {
-    FailureOr<Value> materializedOperand = materializeExternalTensorValue(rewriter, loc, oldBlock, operand, mapper);
-    if (succeeded(materializedOperand))
-      mapper.map(operand, *materializedOperand);
-  }
-
-  Operation* cloned = rewriter.clone(*definingOp, mapper);
-  for (auto [originalResult, clonedResult] : llvm::zip(definingOp->getResults(), cloned->getResults()))
-    mapper.map(originalResult, clonedResult);
-
-  return mapper.lookup(value);
-}
-
-static FailureOr<SmallVector<int32_t>> getPimCoreIdsForBatchOp(spatial::SpatScheduledComputeBatch computeBatchOp,
-                                                               size_t& fallbackCoreId) {
-  if (auto coreIdsAttr = computeBatchOp->getAttrOfType<DenseI32ArrayAttr>(onnx_mlir::kCoreIdsAttrName))
-    return SmallVector<int32_t>(coreIdsAttr.asArrayRef().begin(), coreIdsAttr.asArrayRef().end());
-
-  SmallVector<int32_t> coreIds;
-  coreIds.reserve(static_cast<size_t>(computeBatchOp.getLaneCount()));
-  for (uint32_t lane = 0; lane < computeBatchOp.getLaneCount(); ++lane) {
-    auto checkedCoreId =
-      pim::checkedI32(static_cast<uint64_t>(fallbackCoreId), computeBatchOp, "fallback spatial compute_batch core id");
-    if (failed(checkedCoreId))
-      return failure();
-    coreIds.push_back(*checkedCoreId);
-    ++fallbackCoreId;
-  }
-  return coreIds;
-}
-
 static FailureOr<unsigned> getDirectReturnOperandIndex(OpResult result) {
   if (!result.hasOneUse())
     return failure();
@@ -386,7 +292,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul
       "resultful compute_batch lowering currently requires a spat.in_parallel terminator");
   }
 
-  auto coreIds = getPimCoreIdsForBatchOp(computeBatchOp, coreId);
+  auto coreIds = getRequiredScheduledBatchCoreIds(computeBatchOp, "spatial compute_batch core id");
   if (failed(coreIds))
     return failure();
   SmallVector<Value> batchWeights(computeBatchOp.getWeights().begin(), computeBatchOp.getWeights().end());
@@ -638,9 +544,6 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul
       if (definingOp && definingOp->hasTrait<OpTrait::ConstantLike>())
         continue;
 
-      if (succeeded(materializeExternalTensorValue(rewriter, loc, oldBlock, operand, mapper)))
-        continue;
-
       InFlightDiagnostic diagnostic =
         computeBatchOp.emitOpError("expected external tensor communication to be materialized in Spatial before batch lowering");
       diagnostic << " while cloning nested op '" << op.getName() << "' tensor operand #" << operandIndex;

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp

@@ -9,6 +9,7 @@
 
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Conversion/SpatialToPim/SpatialToPimPass.hpp"
+#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
@@ -141,17 +142,6 @@ cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewriter& rewrite
   }
 }
 
-static FailureOr<int32_t> getPimCoreIdForComputeOp(spatial::SpatScheduledCompute computeOp, size_t& fallbackCoreId) {
-  if (auto spatialCoreIdAttr = computeOp->getAttrOfType<IntegerAttr>(onnx_mlir::kCoreIdAttrName))
-    return pim::checkedI32(spatialCoreIdAttr.getInt(), computeOp, "spatial compute core id");
-  auto checkedCoreId =
-    pim::checkedI32(static_cast<uint64_t>(fallbackCoreId), computeOp, "fallback spatial compute core id");
-  if (failed(checkedCoreId))
-    return failure();
-  ++fallbackCoreId;
-  return *checkedCoreId;
-}
-
 static LogicalResult collectHelperComputeChain(spatial::SpatScheduledCompute computeOp,
                                                SmallVectorImpl<Operation*>& helperChain,
                                                bool requireReturnUse = true) {
@@ -311,7 +301,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCom
   if (!computeOp.getWeights().empty())
     computeWeights.append(computeOp.getWeights().begin(), computeOp.getWeights().end());
   rewriter.setInsertionPointAfter(computeOp);
-  auto checkedCoreId = getPimCoreIdForComputeOp(computeOp, coreId);
+  auto checkedCoreId = getRequiredScheduledCoreId(computeOp, "spatial compute core id");
   if (failed(checkedCoreId))
     return failure();
   auto coreIdAttr = pim::getCheckedI32Attr(rewriter, computeOp, static_cast<int64_t>(*checkedCoreId), "pim core id");

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

@@ -44,121 +44,29 @@ using namespace pim;
 
 namespace onnx_mlir {
 
-static memref::GlobalOp getOrCreateZeroGlobal(IRRewriter& rewriter, Location loc, RankedTensorType tensorType) {
-  auto moduleOp = rewriter.getBlock()->getParentOp()->getParentOfType<ModuleOp>();
-  auto memRefType = MemRefType::get(tensorType.getShape(), tensorType.getElementType());
-  auto zeroAttr = DenseElementsAttr::get(tensorType, rewriter.getZeroAttr(tensorType.getElementType()));
-
-  for (auto globalOp : moduleOp.getOps<memref::GlobalOp>()) {
-    if (!globalOp.getConstant() || globalOp.getType() != memRefType || !globalOp.getInitialValue())
-      continue;
-    if (dyn_cast<DenseElementsAttr>(*globalOp.getInitialValue()) == zeroAttr)
-      return globalOp;
-  }
-
-  std::string nameStem;
-  llvm::raw_string_ostream nameStream(nameStem);
-  nameStream << "__pim_zero_" << tensorType.getRank() << "d_" << tensorType.getNumElements();
-  nameStream.flush();
-
-  std::string symbolName = nameStem;
-  unsigned suffix = 0;
-  while (SymbolTable::lookupSymbolIn(moduleOp, symbolName))
-    symbolName = (nameStem + "_" + Twine(suffix++)).str();
-
-  OpBuilder::InsertionGuard guard(rewriter);
-  rewriter.setInsertionPointToStart(moduleOp.getBody());
-  return memref::GlobalOp::create(rewriter,
-                                  loc,
-                                  rewriter.getStringAttr(symbolName),
-                                  rewriter.getStringAttr("private"),
-                                  TypeAttr::get(memRefType),
-                                  zeroAttr,
-                                  rewriter.getUnitAttr(),
-                                  IntegerAttr {});
-}
-
-static FailureOr<Value> createZeroedDeviceHVector(IRRewriter& rewriter,
-                                                  Location loc,
-                                                  RankedTensorType tensorType,
-                                                  OperationFolder& constantFolder) {
-  auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, tensorType);
-  auto zeroGlobal = getOrCreateZeroGlobal(rewriter, loc, tensorType);
-  auto zeroValue = memref::GetGlobalOp::create(rewriter, loc, zeroGlobal.getType(), zeroGlobal.getName());
-  auto zeroIndex = getOrCreateIndexConstant(constantFolder, outputBuffer.getOperation(), 0);
-  auto byteSize =
-    pim::getCheckedShapedTypeSizeInBytes(tensorType, outputBuffer.getOperation(), "host-to-device zero copy byte size");
-  if (failed(byteSize))
-    return failure();
-  auto sizeAttr =
-    pim::getCheckedI32Attr(rewriter, outputBuffer.getOperation(), *byteSize, "host-to-device zero copy byte size");
-  if (failed(sizeAttr))
-    return failure();
-  return PimMemCopyHostToDevOp::create(
-           rewriter, loc, tensorType, zeroIndex, zeroIndex, outputBuffer, zeroValue, *sizeAttr)
-    .getOutput();
-}
-
-static bool isHostBackedMemRefValue(Value value) {
-  while (Operation* definingOp = value.getDefiningOp()) {
-    if (auto subviewOp = dyn_cast<memref::SubViewOp>(definingOp)) {
-      value = subviewOp.getSource();
-      continue;
-    }
-    if (auto castOp = dyn_cast<memref::CastOp>(definingOp)) {
-      value = castOp.getSource();
-      continue;
-    }
-    if (auto collapseOp = dyn_cast<memref::CollapseShapeOp>(definingOp)) {
-      value = collapseOp.getSrc();
-      continue;
-    }
-    if (auto expandOp = dyn_cast<memref::ExpandShapeOp>(definingOp)) {
-      value = expandOp.getSrc();
-      continue;
-    }
-    return isa<memref::GetGlobalOp>(definingOp);
-  }
-  return false;
-}
-
-static bool isHostBackedTensorValue(Value value) {
-  while (Operation* definingOp = value.getDefiningOp()) {
-    if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(definingOp)) {
-      auto sourceType = dyn_cast<RankedTensorType>(extractSliceOp.getSource().getType());
-      auto resultType = dyn_cast<RankedTensorType>(extractSliceOp.getResult().getType());
-      if (!sourceType || !resultType || !sourceType.hasStaticShape() || !resultType.hasStaticShape())
-        return false;
-      if (!onnx_mlir::isContiguousSubviewWithDynamicOffsets(sourceType.getShape(),
-                                                            extractSliceOp.getMixedOffsets(),
-                                                            extractSliceOp.getStaticSizes(),
-                                                            extractSliceOp.getStaticStrides())) {
-        return false;
-      }
-      value = extractSliceOp.getSource();
-      continue;
-    }
-    if (auto collapseOp = dyn_cast<tensor::CollapseShapeOp>(definingOp)) {
-      value = collapseOp.getSrc();
-      continue;
-    }
-    if (auto expandOp = dyn_cast<tensor::ExpandShapeOp>(definingOp)) {
-      value = expandOp.getSrc();
-      continue;
-    }
-    if (auto castOp = dyn_cast<tensor::CastOp>(definingOp)) {
-      value = castOp.getSource();
-      continue;
-    }
-    if (auto toTensorOp = dyn_cast<bufferization::ToTensorOp>(definingOp))
-      return isHostBackedMemRefValue(toTensorOp.getBuffer());
-    return false;
-  }
-  return false;
-}
-
 static FailureOr<Value>
-padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector, OperationFolder& constantFolder) {
+createZeroPaddedTensor(IRRewriter& rewriter, Location loc, Value value, RankedTensorType resultType) {
+  auto sourceType = cast<RankedTensorType>(value.getType());
+  SmallVector<OpFoldResult> lowPads(sourceType.getRank(), rewriter.getIndexAttr(0));
+  SmallVector<OpFoldResult> highPads;
+  highPads.reserve(sourceType.getRank());
+  for (auto [sourceDim, resultDim] : llvm::zip(sourceType.getShape(), resultType.getShape()))
+    highPads.push_back(rewriter.getIndexAttr(resultDim - sourceDim));
+
+  auto padOp = tensor::PadOp::create(rewriter, loc, resultType, value, lowPads, highPads);
+  auto* padBlock = new Block();
+  for (int64_t i = 0; i < sourceType.getRank(); ++i)
+    padBlock->addArgument(rewriter.getIndexType(), loc);
+  padOp.getRegion().push_back(padBlock);
+  rewriter.setInsertionPointToStart(padBlock);
+  auto zero = getOrCreateConstant(
+    rewriter, padOp.getOperation(), rewriter.getZeroAttr(sourceType.getElementType()), sourceType.getElementType());
+  tensor::YieldOp::create(rewriter, loc, zero);
+  rewriter.setInsertionPointAfter(padOp);
+  return padOp.getResult();
+}
+
+static FailureOr<Value> padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector) {
   auto vectorType = cast<RankedTensorType>(vector.getType());
   ArrayRef<int64_t> shape = vectorType.getShape();
   assert(isHVectorShape(shape) && "expected a horizontal vector");
@@ -169,26 +77,10 @@ padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector,
 
   auto paddedType = RankedTensorType::get(
     {shape[0], static_cast<int64_t>(crossbarSize)}, vectorType.getElementType(), vectorType.getEncoding());
-  auto zeroed = createZeroedDeviceHVector(rewriter, loc, paddedType, constantFolder);
-  if (failed(zeroed))
-    return failure();
-  Value zeroIndex = getOrCreateIndexConstant(constantFolder, zeroed->getDefiningOp(), 0);
-  auto byteSize =
-    pim::getCheckedShapedTypeSizeInBytes(vectorType, zeroed->getDefiningOp(), "device padding copy byte size");
-  if (failed(byteSize))
-    return failure();
-  auto sizeAttr = pim::getCheckedI32Attr(rewriter, zeroed->getDefiningOp(), *byteSize, "device padding copy byte size");
-  if (failed(sizeAttr))
-    return failure();
-  if (isHostBackedTensorValue(vector)) {
-    return PimMemCopyHostToDevOp::create(rewriter, loc, paddedType, zeroIndex, zeroIndex, *zeroed, vector, *sizeAttr)
-      .getOutput();
-  }
-  return PimMemCopyOp::create(rewriter, loc, paddedType, zeroIndex, zeroIndex, *zeroed, vector, *sizeAttr).getOutput();
+  return createZeroPaddedTensor(rewriter, loc, vector, paddedType);
 }
 
 void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
-  coreId = 0;
   outputTensors.clear();
   operationsToRemove.clear();
   ModuleOp moduleOp = getOperation();
@@ -362,7 +254,6 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
 }
 
 LogicalResult raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func::FuncOp funcOp, IRRewriter& rewriter) {
-  OperationFolder constantFolder(funcOp.getContext());
   bool hasFailure = false;
   funcOp.walk([&](PimVMMOp vmmOp) {
     auto outputType = cast<RankedTensorType>(vmmOp.getOutput().getType());
@@ -371,7 +262,7 @@ LogicalResult raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func:
     assert(outputShape[1] <= static_cast<int64_t>(crossbarSize) && "output width must fit in one crossbar");
 
     rewriter.setInsertionPoint(vmmOp);
-    auto paddedInput = padHVectorInputToCrossbarSize(rewriter, vmmOp.getLoc(), vmmOp.getInput(), constantFolder);
+    auto paddedInput = padHVectorInputToCrossbarSize(rewriter, vmmOp.getLoc(), vmmOp.getInput());
     if (failed(paddedInput)) {
       hasFailure = true;
       return WalkResult::interrupt();

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.hpp

@@ -36,7 +36,6 @@ private:
   using OutputTensorFactory = std::function<mlir::Value(mlir::IRRewriter& rewriter, mlir::Location loc)>;
 
   llvm::SmallVector<OutputTensorFactory> outputTensors;
-  size_t coreId = 0;
   llvm::SmallVector<mlir::Operation*> operationsToRemove;
 
   mlir::LogicalResult allocateAndInitializeCoreLocalVariables(mlir::func::FuncOp funcOp, mlir::IRRewriter& rewriter);