refactorone

src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp

@@ -2,6 +2,7 @@
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/IRMapping.h"
+#include "mlir/IR/Matchers.h"
 
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
@@ -17,6 +18,37 @@ namespace {
 
 static int32_t translateSpatialCoreIdToPimCoreId(size_t spatialCoreId) { return static_cast<int32_t>(spatialCoreId); }
 
+static FailureOr<int32_t> getConstantI32Value(Value value) {
+  APInt constantValue;
+  if (!matchPattern(value, m_ConstantInt(&constantValue)))
+    return failure();
+  return static_cast<int32_t>(constantValue.getSExtValue());
+}
+
+static FailureOr<SmallVector<int32_t>> getConstantI32Values(ValueRange values) {
+  SmallVector<int32_t> constants;
+  constants.reserve(values.size());
+  for (Value value : values) {
+    FailureOr<int32_t> constantValue = getConstantI32Value(value);
+    if (failed(constantValue))
+      return failure();
+    constants.push_back(*constantValue);
+  }
+  return constants;
+}
+
+static bool isExplicitHostOperand(Operation* op, unsigned operandIndex) {
+  if (isa<pim::PimMemCopyDevToHostOp>(op))
+    return operandIndex == 2;
+  return false;
+}
+
+static bool isUsedOnlyAsExplicitHostOperand(Value value) {
+  return !value.use_empty() && llvm::all_of(value.getUses(), [](OpOperand& use) {
+    return isExplicitHostOperand(use.getOwner(), use.getOperandNumber());
+  });
+}
+
 static SmallVector<int32_t> getPimCoreIdsForBatchOp(spatial::SpatComputeBatch computeBatchOp, size_t& fallbackCoreId) {
   if (auto coreIdsAttr = computeBatchOp->getAttrOfType<DenseI32ArrayAttr>(onnx_mlir::kCoreIdsAttrName))
     return SmallVector<int32_t>(coreIdsAttr.asArrayRef().begin(), coreIdsAttr.asArrayRef().end());
@@ -28,27 +60,30 @@ static SmallVector<int32_t> getPimCoreIdsForBatchOp(spatial::SpatComputeBatch co
   return coreIds;
 }
 
-static void lowerChannelSendTensorBatch(spatial::SpatChannelSendTensorBatchOp sendTensorBatchOp,
-                                        IRMapping& mapper,
-                                        IRRewriter& rewriter) {
-  SmallVector<int32_t> targetCoreIds;
-  targetCoreIds.reserve(sendTensorBatchOp.getTargetCoreIds().size());
-  for (int32_t targetCoreId : sendTensorBatchOp.getTargetCoreIds())
-    targetCoreIds.push_back(translateSpatialCoreIdToPimCoreId(targetCoreId));
+static LogicalResult lowerChannelSendTensorBatch(spatial::SpatChannelSendTensorBatchOp sendTensorBatchOp,
+                                                 IRMapping& mapper,
+                                                 IRRewriter& rewriter) {
+  FailureOr<SmallVector<int32_t>> targetCoreIds = getConstantI32Values(sendTensorBatchOp.getTargetCoreIds());
+  if (failed(targetCoreIds))
+    return sendTensorBatchOp.emitOpError("expected constant targetCoreIds");
+  for (int32_t& targetCoreId : *targetCoreIds)
+    targetCoreId = translateSpatialCoreIdToPimCoreId(targetCoreId);
 
   pim::PimSendTensorBatchOp::create(rewriter,
                                     sendTensorBatchOp.getLoc(),
                                     mapper.lookup(sendTensorBatchOp.getInput()),
-                                    rewriter.getDenseI32ArrayAttr(targetCoreIds));
+                                    rewriter.getDenseI32ArrayAttr(*targetCoreIds));
+  return success();
 }
 
-static void lowerChannelReceiveTensorBatch(spatial::SpatChannelReceiveTensorBatchOp receiveTensorBatchOp,
-                                           IRMapping& mapper,
-                                           IRRewriter& rewriter) {
-  SmallVector<int32_t> sourceCoreIds;
-  sourceCoreIds.reserve(receiveTensorBatchOp.getSourceCoreIds().size());
-  for (int32_t sourceCoreId : receiveTensorBatchOp.getSourceCoreIds())
-    sourceCoreIds.push_back(translateSpatialCoreIdToPimCoreId(sourceCoreId));
+static LogicalResult lowerChannelReceiveTensorBatch(spatial::SpatChannelReceiveTensorBatchOp receiveTensorBatchOp,
+                                                    IRMapping& mapper,
+                                                    IRRewriter& rewriter) {
+  FailureOr<SmallVector<int32_t>> sourceCoreIds = getConstantI32Values(receiveTensorBatchOp.getSourceCoreIds());
+  if (failed(sourceCoreIds))
+    return receiveTensorBatchOp.emitOpError("expected constant sourceCoreIds");
+  for (int32_t& sourceCoreId : *sourceCoreIds)
+    sourceCoreId = translateSpatialCoreIdToPimCoreId(sourceCoreId);
 
   auto outputType = cast<ShapedType>(receiveTensorBatchOp.getOutput().getType());
   auto outputBuffer = createEmptyTensorFromShaped(rewriter, receiveTensorBatchOp.getLoc(), outputType);
@@ -56,24 +91,26 @@ static void lowerChannelReceiveTensorBatch(spatial::SpatChannelReceiveTensorBatc
                                                         receiveTensorBatchOp.getLoc(),
                                                         outputBuffer.getType(),
                                                         outputBuffer,
-                                                        rewriter.getDenseI32ArrayAttr(sourceCoreIds))
+                                                        rewriter.getDenseI32ArrayAttr(*sourceCoreIds))
                      .getOutput();
   mapper.map(receiveTensorBatchOp.getOutput(), received);
+  return success();
 }
 
 } // namespace
 
 LogicalResult
 lowerComputeBatchOp(spatial::SpatComputeBatch computeBatchOp, CoreLoweringState& state, IRRewriter& rewriter) {
-  if (computeBatchOp.getNumResults() != 0)
-    return computeBatchOp.emitOpError(
-      "batched Spatial-to-PIM lowering currently requires channelized compute_batch with no results");
-
   Location loc = computeBatchOp.getLoc();
   Block& oldBlock = computeBatchOp.getBody().front();
-  auto oldYield = cast<spatial::SpatYieldOp>(oldBlock.getTerminator());
-  if (oldYield.getNumOperands() != 0)
-    return computeBatchOp.emitOpError("batched Spatial-to-PIM lowering currently requires empty spat.yield");
+  if (computeBatchOp.getNumResults() != 0)
+    return computeBatchOp.emitOpError(
+      "batched Spatial-to-PIM lowering currently requires channelized compute_batch with no results; "
+      "materialize explicit communication before lowering to PIM");
+
+  auto oldYield = dyn_cast<spatial::SpatYieldOp>(oldBlock.getTerminator());
+  if (!oldYield || oldYield.getNumOperands() != 0)
+    return computeBatchOp.emitOpError("resultless compute_batch lowering requires empty spat.yield");
 
   SmallVector<int32_t> coreIds = getPimCoreIdsForBatchOp(computeBatchOp, state.nextCoreId);
   SmallVector<Value> batchWeights(computeBatchOp.getWeights().begin(), computeBatchOp.getWeights().end());
@@ -102,7 +139,12 @@ lowerComputeBatchOp(spatial::SpatComputeBatch computeBatchOp, CoreLoweringState&
 
   IRMapping mapper;
   rewriter.setInsertionPointToStart(newBlock);
-  for (auto [oldArg, newArg] : llvm::zip(oldBlock.getArguments(), newBlock->getArguments())) {
+  mapper.map(computeBatchOp.getLaneArgument(), coreBatchOp.getLaneArgument());
+  for (unsigned weightIndex = 0; weightIndex < computeBatchOp.getWeights().size(); ++weightIndex)
+    mapper.map(computeBatchOp.getWeightArgument(weightIndex), coreBatchOp.getWeightArgument(weightIndex));
+  for (unsigned inputIndex = 0; inputIndex < computeBatchOp.getInputs().size(); ++inputIndex) {
+    BlockArgument oldArg = computeBatchOp.getInputArgument(inputIndex);
+    BlockArgument newArg = coreBatchOp.getInputArgument(inputIndex);
     auto newArgType = cast<ShapedType>(newArg.getType());
     auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, newArgType);
     auto copied = pim::PimMemCopyHostToDevBatchOp::create(rewriter,
@@ -142,20 +184,31 @@ lowerComputeBatchOp(spatial::SpatComputeBatch computeBatchOp, CoreLoweringState&
       continue;
 
     if (auto sendBatchOp = dyn_cast<spatial::SpatChannelSendBatchOp>(op)) {
+      FailureOr<SmallVector<int32_t>> targetCoreIds = getConstantI32Values(sendBatchOp.getTargetCoreIds());
+      if (failed(targetCoreIds))
+        return sendBatchOp.emitOpError("expected constant targetCoreIds");
+      for (int32_t& targetCoreId : *targetCoreIds)
+        targetCoreId = translateSpatialCoreIdToPimCoreId(targetCoreId);
       pim::PimSendBatchOp::create(rewriter,
                                   loc,
                                   mapper.lookup(sendBatchOp.getInput()),
                                   getTensorSizeInBytesAttr(rewriter, mapper.lookup(sendBatchOp.getInput())),
-                                  sendBatchOp.getTargetCoreIdsAttr());
+                                  rewriter.getDenseI32ArrayAttr(*targetCoreIds));
       continue;
     }
 
     if (auto sendTensorBatchOp = dyn_cast<spatial::SpatChannelSendTensorBatchOp>(op)) {
-      lowerChannelSendTensorBatch(sendTensorBatchOp, mapper, rewriter);
+      if (failed(lowerChannelSendTensorBatch(sendTensorBatchOp, mapper, rewriter)))
+        return failure();
       continue;
     }
 
     if (auto receiveBatchOp = dyn_cast<spatial::SpatChannelReceiveBatchOp>(op)) {
+      FailureOr<SmallVector<int32_t>> sourceCoreIds = getConstantI32Values(receiveBatchOp.getSourceCoreIds());
+      if (failed(sourceCoreIds))
+        return receiveBatchOp.emitOpError("expected constant sourceCoreIds");
+      for (int32_t& sourceCoreId : *sourceCoreIds)
+        sourceCoreId = translateSpatialCoreIdToPimCoreId(sourceCoreId);
       auto outputType = cast<ShapedType>(receiveBatchOp.getOutput().getType());
       auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, outputType);
       auto received = pim::PimReceiveBatchOp::create(rewriter,
@@ -163,14 +216,15 @@ lowerComputeBatchOp(spatial::SpatComputeBatch computeBatchOp, CoreLoweringState&
                                                      outputBuffer.getType(),
                                                      outputBuffer,
                                                      getTensorSizeInBytesAttr(rewriter, receiveBatchOp.getOutput()),
-                                                     receiveBatchOp.getSourceCoreIdsAttr())
+                                                     rewriter.getDenseI32ArrayAttr(*sourceCoreIds))
                         .getOutput();
       mapper.map(receiveBatchOp.getOutput(), received);
       continue;
     }
 
     if (auto receiveTensorBatchOp = dyn_cast<spatial::SpatChannelReceiveTensorBatchOp>(op)) {
-      lowerChannelReceiveTensorBatch(receiveTensorBatchOp, mapper, rewriter);
+      if (failed(lowerChannelReceiveTensorBatch(receiveTensorBatchOp, mapper, rewriter)))
+        return failure();
       continue;
     }
 
@@ -178,6 +232,10 @@ lowerComputeBatchOp(spatial::SpatComputeBatch computeBatchOp, CoreLoweringState&
       if (isa_and_present<memref::GetGlobalOp>(toTensorOp.getBuffer().getDefiningOp())) {
         Operation* cloned = rewriter.clone(op, mapper);
         auto clonedTensor = cloned->getResult(0);
+        if (isUsedOnlyAsExplicitHostOperand(toTensorOp.getResult())) {
+          mapper.map(toTensorOp.getResult(), clonedTensor);
+          continue;
+        }
         auto clonedType = cast<ShapedType>(clonedTensor.getType());
         auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, clonedType);
         auto copied = pim::PimMemCopyHostToDevBatchOp::create(rewriter,
@@ -194,9 +252,11 @@ lowerComputeBatchOp(spatial::SpatComputeBatch computeBatchOp, CoreLoweringState&
       }
     }
 
-    for (Value operand : op.getOperands()) {
+    for (auto [operandIndex, operand] : llvm::enumerate(op.getOperands())) {
       if (!isa<TensorType>(operand.getType()) || mapper.contains(operand))
         continue;
+      if (isExplicitHostOperand(&op, operandIndex))
+        continue;
 
       Operation* definingOp = operand.getDefiningOp();
       if (definingOp && definingOp->getBlock() == &oldBlock)

src/PIM/Conversion/SpatialToPim/CMakeLists.txt

@@ -22,6 +22,8 @@ add_pim_library(OMSpatialToPim
 
   LINK_LIBS PUBLIC
   MLIRSCFDialect
+  MLIRSCFUtils
+  MLIRTransformUtils
   MLIRTosaDialect
   OMCompilerOptions
   OMPimCommon

src/PIM/Conversion/SpatialToPim/ChannelLoweringPatterns.cpp

@@ -1,4 +1,5 @@
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/Matchers.h"
 
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/ChannelLoweringPatterns.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
@@ -12,15 +13,24 @@ namespace {
 
 static int32_t toPimCoreId(int32_t spatialCoreId) { return spatialCoreId; }
 
+static FailureOr<SmallVector<int32_t>> getConstantI32Values(ValueRange values) {
+  SmallVector<int32_t> constants;
+  constants.reserve(values.size());
+  for (Value value : values) {
+    APInt constantValue;
+    if (!matchPattern(value, m_ConstantInt(&constantValue)))
+      return failure();
+    constants.push_back(static_cast<int32_t>(constantValue.getSExtValue()));
+  }
+  return constants;
+}
+
 struct ChannelSendLowering : OpRewritePattern<spatial::SpatChannelSendOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(spatial::SpatChannelSendOp op, PatternRewriter& rewriter) const override {
-    pim::PimSendOp::create(rewriter,
-                           op.getLoc(),
-                           op.getInput(),
-                           getTensorSizeInBytesAttr(rewriter, op.getInput()),
-                           rewriter.getI32IntegerAttr(toPimCoreId(op.getTargetCoreId())));
+    pim::PimSendOp::create(
+      rewriter, op.getLoc(), op.getInput(), getTensorSizeInBytesAttr(rewriter, op.getInput()), op.getTargetCoreId());
     rewriter.eraseOp(op);
     return success();
   }
@@ -42,7 +52,7 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
                                                op.getResult().getType(),
                                                outputBuffer,
                                                getTensorSizeInBytesAttr(rewriter, op.getResult()),
-                                               rewriter.getI32IntegerAttr(toPimCoreId(op.getSourceCoreId())))
+                                               op.getSourceCoreId())
                        .getOutput();
     rewriter.replaceOp(op, received);
     return success();
@@ -53,11 +63,12 @@ struct ChannelSendTensorLowering : OpRewritePattern<spatial::SpatChannelSendTens
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(spatial::SpatChannelSendTensorOp op, PatternRewriter& rewriter) const override {
-    SmallVector<int32_t> targetCoreIds;
-    targetCoreIds.reserve(op.getTargetCoreIds().size());
-    for (int32_t targetCoreId : op.getTargetCoreIds())
-      targetCoreIds.push_back(toPimCoreId(targetCoreId));
-    pim::PimSendTensorOp::create(rewriter, op.getLoc(), op.getInput(), rewriter.getDenseI32ArrayAttr(targetCoreIds));
+    FailureOr<SmallVector<int32_t>> targetCoreIds = getConstantI32Values(op.getTargetCoreIds());
+    if (failed(targetCoreIds))
+      return rewriter.notifyMatchFailure(op, "expected constant targetCoreIds");
+    for (int32_t& targetCoreId : *targetCoreIds)
+      targetCoreId = toPimCoreId(targetCoreId);
+    pim::PimSendTensorOp::create(rewriter, op.getLoc(), op.getInput(), rewriter.getDenseI32ArrayAttr(*targetCoreIds));
     rewriter.eraseOp(op);
     return success();
   }
@@ -67,16 +78,17 @@ struct ChannelReceiveTensorLowering : OpRewritePattern<spatial::SpatChannelRecei
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(spatial::SpatChannelReceiveTensorOp op, PatternRewriter& rewriter) const override {
-    SmallVector<int32_t> sourceCoreIds;
-    sourceCoreIds.reserve(op.getSourceCoreIds().size());
-    for (int32_t sourceCoreId : op.getSourceCoreIds())
-      sourceCoreIds.push_back(toPimCoreId(sourceCoreId));
+    FailureOr<SmallVector<int32_t>> sourceCoreIds = getConstantI32Values(op.getSourceCoreIds());
+    if (failed(sourceCoreIds))
+      return rewriter.notifyMatchFailure(op, "expected constant sourceCoreIds");
+    for (int32_t& sourceCoreId : *sourceCoreIds)
+      sourceCoreId = toPimCoreId(sourceCoreId);
     auto outputType = cast<ShapedType>(op.getOutput().getType());
     Value outputBuffer =
       tensor::EmptyOp::create(rewriter, op.getLoc(), outputType.getShape(), outputType.getElementType()).getResult();
     Value received =
       pim::PimReceiveTensorOp::create(
-        rewriter, op.getLoc(), op.getOutput().getType(), outputBuffer, rewriter.getDenseI32ArrayAttr(sourceCoreIds))
+        rewriter, op.getLoc(), op.getOutput().getType(), outputBuffer, rewriter.getDenseI32ArrayAttr(*sourceCoreIds))
         .getOutput();
     rewriter.replaceOp(op, received);
     return success();

src/PIM/Conversion/SpatialToPim/ComputeLikeRegionUtils.cpp

@@ -29,7 +29,10 @@ void replaceAndEraseDirectComputeLikeInput(PatternRewriter& rewriter,
                                            unsigned inputIndex,
                                            Value replacement) {
   Block& body = owner->getRegion(0).front();
-  BlockArgument bodyArgument = body.getArgument(inputIndex);
+  BlockArgument bodyArgument = isa<spatial::SpatCompute>(owner)
+                                 ? cast<spatial::SpatCompute>(owner).getInputArgument(inputIndex)
+                                 : cast<spatial::SpatComputeBatch>(owner).getInputArgument(inputIndex);
+  unsigned bodyArgIndex = bodyArgument.getArgNumber();
 
   rewriter.startOpModification(owner);
   bodyArgument.replaceAllUsesWith(replacement);
@@ -37,7 +40,7 @@ void replaceAndEraseDirectComputeLikeInput(PatternRewriter& rewriter,
     compute.getInputsMutable().erase(inputIndex);
   else
     cast<spatial::SpatComputeBatch>(owner).getInputsMutable().erase(inputIndex);
-  body.eraseArgument(inputIndex);
+  body.eraseArgument(bodyArgIndex);
   rewriter.finalizeOpModification(owner);
 }
 

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp

@@ -3,6 +3,7 @@
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/IR/IRMapping.h"
+#include "mlir/IR/Matchers.h"
 
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
@@ -27,7 +28,8 @@ static bool isChannelUseChainOp(Operation* op) {
              pim::PimTransposeOp>(op);
 }
 
-static void cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewriter& rewriter) {
+static void
+cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewriter& rewriter, OperationFolder& constantFolder) {
   for (Value operand : op->getOperands()) {
     if (mapping.lookupOrNull(operand))
       continue;
@@ -36,7 +38,12 @@ static void cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewri
     if (!definingOp)
       continue;
 
-    if (!isa<tensor::EmptyOp, arith::ConstantOp>(definingOp))
+    if (auto constantOp = dyn_cast<arith::ConstantOp>(definingOp)) {
+      mapping.map(operand, getOrCreateHostConstantLike(constantOp, constantFolder));
+      continue;
+    }
+
+    if (!isa<tensor::EmptyOp>(definingOp))
       continue;
 
     Operation* clonedOp = rewriter.clone(*definingOp, mapping);
@@ -48,6 +55,18 @@ static void cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewri
 
 static int32_t translateSpatialCoreIdToPimCoreId(size_t spatialCoreId) { return static_cast<int32_t>(spatialCoreId); }
 
+static FailureOr<SmallVector<int32_t>> getConstantI32Values(ValueRange values) {
+  SmallVector<int32_t> constants;
+  constants.reserve(values.size());
+  for (Value value : values) {
+    APInt constantValue;
+    if (!matchPattern(value, m_ConstantInt(&constantValue)))
+      return failure();
+    constants.push_back(static_cast<int32_t>(constantValue.getSExtValue()));
+  }
+  return constants;
+}
+
 static int32_t getPimCoreIdForComputeOp(spatial::SpatCompute computeOp, size_t& fallbackCoreId) {
   if (auto spatialCoreIdAttr = computeOp->getAttrOfType<IntegerAttr>(onnx_mlir::kCoreIdAttrName))
     return static_cast<int32_t>(spatialCoreIdAttr.getInt());
@@ -92,7 +111,9 @@ static LogicalResult collectHelperComputeChain(spatial::SpatCompute computeOp,
   return success();
 }
 
-static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatCompute computeOp, IRRewriter& rewriter) {
+static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatCompute computeOp,
+                                                           IRRewriter& rewriter,
+                                                           OperationFolder& constantFolder) {
   if (!computeOp.getInputs().empty() || computeOp.getNumResults() != 1)
     return false;
   if (!llvm::all_of(computeOp.getResult(0).getUsers(), [](Operation* user) {
@@ -101,7 +122,7 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatCompute
     return false;
 
   Block& block = computeOp.getBody().front();
-  if (block.getNumArguments() != 0)
+  if (block.getNumArguments() != computeOp.getWeights().size())
     return false;
 
   auto yieldOp = dyn_cast<spatial::SpatYieldOp>(block.getTerminator());
@@ -110,8 +131,10 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatCompute
 
   rewriter.setInsertionPoint(computeOp);
   IRMapping mapping;
+  for (auto [weightIndex, weight] : llvm::enumerate(computeOp.getWeights()))
+    mapping.map(computeOp.getWeightArgument(weightIndex), weight);
   for (Operation& op : block.without_terminator()) {
-    cloneMappedHelperOperands(&op, mapping, rewriter);
+    cloneMappedHelperOperands(&op, mapping, rewriter, constantFolder);
     Operation* clonedOp = rewriter.clone(op, mapping);
     for (auto [originalResult, newResult] : llvm::zip(op.getResults(), clonedOp->getResults()))
       mapping.map(originalResult, newResult);
@@ -133,7 +156,7 @@ void markOpToRemove(CoreLoweringState& state, Operation* op) {
 LogicalResult lowerComputeOp(spatial::SpatCompute computeOp, CoreLoweringState& state, IRRewriter& rewriter) {
   Location loc = computeOp->getLoc();
 
-  if (inlineInputlessHelperComputeForWeightLikeUsers(computeOp, rewriter))
+  if (inlineInputlessHelperComputeForWeightLikeUsers(computeOp, rewriter, state.constantFolder))
     return success();
 
   SmallVector<Operation*> helperChain;
@@ -143,21 +166,42 @@ LogicalResult lowerComputeOp(spatial::SpatCompute computeOp, CoreLoweringState&
   auto& block = computeOp.getRegion().front();
   auto yieldOp = cast<spatial::SpatYieldOp>(block.getTerminator());
 
-  for (auto [argIndex, blockArg] : llvm::enumerate(block.getArguments())) {
-    auto receiveOp = dyn_cast_or_null<spatial::SpatChannelReceiveOp>(computeOp.getInputs()[argIndex].getDefiningOp());
-    if (!receiveOp || blockArg.use_empty())
+  for (auto [inputIndex, input] : llvm::enumerate(computeOp.getInputs())) {
+    BlockArgument blockArg = computeOp.getInputArgument(inputIndex);
+    auto receiveOp = dyn_cast_or_null<spatial::SpatChannelReceiveOp>(input.getDefiningOp());
+    if (receiveOp && !blockArg.use_empty()) {
+      rewriter.setInsertionPoint(getEarliestUserWithinBlock(blockArg));
+      auto outputType = cast<ShapedType>(blockArg.getType());
+      auto outputBuffer = createEmptyTensorFromShaped(rewriter, receiveOp.getLoc(), outputType);
+      auto sizeAttr = getTensorSizeInBytesAttr(rewriter, blockArg);
+      Value received =
+        PimReceiveOp::create(
+          rewriter, receiveOp.getLoc(), outputBuffer.getType(), outputBuffer, sizeAttr, receiveOp.getSourceCoreId())
+          .getOutput();
+      blockArg.replaceAllUsesWith(received);
+      markOpToRemove(state, receiveOp);
       continue;
+    }
 
-    rewriter.setInsertionPoint(getEarliestUserWithinBlock(blockArg));
-    auto outputType = cast<ShapedType>(blockArg.getType());
-    auto outputBuffer = createEmptyTensorFromShaped(rewriter, receiveOp.getLoc(), outputType);
-    auto sizeAttr = getTensorSizeInBytesAttr(rewriter, blockArg);
-    auto sourceCoreIdAttr = rewriter.getI32IntegerAttr(translateSpatialCoreIdToPimCoreId(receiveOp.getSourceCoreId()));
-    Value received = PimReceiveOp::create(
-                       rewriter, receiveOp.getLoc(), outputBuffer.getType(), outputBuffer, sizeAttr, sourceCoreIdAttr)
-                       .getOutput();
-    blockArg.replaceAllUsesWith(received);
-    markOpToRemove(state, receiveOp);
+    auto receiveTensorOp = dyn_cast_or_null<spatial::SpatChannelReceiveTensorOp>(input.getDefiningOp());
+    if (receiveTensorOp && !blockArg.use_empty()) {
+      FailureOr<SmallVector<int32_t>> sourceCoreIds = getConstantI32Values(receiveTensorOp.getSourceCoreIds());
+      if (failed(sourceCoreIds))
+        return receiveTensorOp.emitOpError("expected constant sourceCoreIds");
+      for (int32_t& sourceCoreId : *sourceCoreIds)
+        sourceCoreId = translateSpatialCoreIdToPimCoreId(sourceCoreId);
+      rewriter.setInsertionPoint(getEarliestUserWithinBlock(blockArg));
+      auto outputType = cast<ShapedType>(blockArg.getType());
+      auto outputBuffer = createEmptyTensorFromShaped(rewriter, receiveTensorOp.getLoc(), outputType);
+      Value received = PimReceiveTensorOp::create(rewriter,
+                                                  receiveTensorOp.getLoc(),
+                                                  outputBuffer.getType(),
+                                                  outputBuffer,
+                                                  rewriter.getDenseI32ArrayAttr(*sourceCoreIds))
+                         .getOutput();
+      blockArg.replaceAllUsesWith(received);
+      markOpToRemove(state, receiveTensorOp);
+    }
   }
 
   if (computeOp.getNumResults() != yieldOp.getNumOperands())
@@ -197,11 +241,36 @@ LogicalResult lowerComputeOp(spatial::SpatCompute computeOp, CoreLoweringState&
                                   loc,
                                   ValueRange(computeWeights),
                                   rewriter.getI32IntegerAttr(getPimCoreIdForComputeOp(computeOp, state.nextCoreId)));
+  rewriter.setInsertionPointToStart(&block);
   auto& coreOpBlocks = coreOp.getBody().getBlocks();
-  for (auto [argIndex, blockArg] : llvm::enumerate(block.getArguments()))
-    if (!blockArg.use_empty())
-      blockArg.replaceAllUsesWith(computeOp.getInputs()[argIndex]);
-  block.eraseArguments(0, block.getNumArguments());
+  for (auto [inputIndex, input] : llvm::enumerate(computeOp.getInputs())) {
+    BlockArgument blockArg = computeOp.getInputArgument(inputIndex);
+    if (blockArg.use_empty())
+      continue;
+
+    if (auto constantOp = input.getDefiningOp<arith::ConstantOp>()) {
+      blockArg.replaceAllUsesWith(getOrCreateHostConstantLike(constantOp, state.constantFolder));
+      continue;
+    }
+
+    auto inputType = dyn_cast<ShapedType>(input.getType());
+    if (!inputType)
+      return computeOp.emitOpError("expected shaped compute input during pim.core lowering");
+    auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, inputType);
+    auto copied =
+      PimMemCopyHostToDevOp::create(rewriter,
+                                    loc,
+                                    outputBuffer.getType(),
+                                    getOrCreateHostIndexConstant(outputBuffer.getOperation(), 0, state.constantFolder),
+                                    getOrCreateHostIndexConstant(outputBuffer.getOperation(), 0, state.constantFolder),
+                                    outputBuffer,
+                                    input,
+                                    getTensorSizeInBytesAttr(rewriter, input))
+        .getOutput();
+    blockArg.replaceAllUsesWith(copied);
+  }
+  if (!computeOp.getInputs().empty())
+    block.eraseArguments(computeOp.getWeights().size(), computeOp.getInputs().size());
   coreOpBlocks.splice(coreOpBlocks.begin(), computeOp.getBody().getBlocks());
   Block* tempComputeBlock = new Block();
   computeOp.getBody().push_back(tempComputeBlock);

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.hpp

@@ -1,6 +1,7 @@
 #pragma once
 
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/Transforms/FoldUtils.h"
 
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/ReturnPathNormalization.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -11,6 +12,7 @@ struct CoreLoweringState {
   size_t& nextCoreId;
   llvm::SmallVectorImpl<OutputTensorFactory>& outputTensors;
   llvm::SmallVectorImpl<mlir::Operation*>& operationsToRemove;
+  mlir::OperationFolder& constantFolder;
 };
 
 void markOpToRemove(CoreLoweringState& state, mlir::Operation* op);

src/PIM/Conversion/SpatialToPim/GlobalTensorMaterialization.cpp

@@ -76,8 +76,7 @@ struct MoveExtractSliceIntoCompute final : OpRewritePattern<mlir::tensor::Extrac
         auto inputIndex = getDirectComputeLikeInputIndex(spatCompute, uses.getOperandNumber());
         if (!inputIndex)
           return failure();
-        auto BBArgIndex = *inputIndex;
-        auto BBArgValue = spatCompute.getBody().front().getArgument(BBArgIndex);
+        auto BBArgValue = spatCompute.getInputArgument(*inputIndex);
 
         if (BBArgValue.use_empty())
           continue;
@@ -89,14 +88,13 @@ struct MoveExtractSliceIntoCompute final : OpRewritePattern<mlir::tensor::Extrac
         }
 
         replaceAndEraseDirectComputeLikeInput(
-          rewriter, spatCompute.getOperation(), BBArgIndex, mapSpatToExtract[spatCompute.getOperation()]);
+          rewriter, spatCompute.getOperation(), *inputIndex, mapSpatToExtract[spatCompute.getOperation()]);
       }
       else if (auto spatComputeBatch = dyn_cast<spatial::SpatComputeBatch>(uses.getOwner())) {
         auto inputIndex = getDirectComputeLikeInputIndex(spatComputeBatch, uses.getOperandNumber());
         if (!inputIndex)
           return failure();
-        auto BBArgIndex = *inputIndex;
-        auto BBArgValue = spatComputeBatch.getBody().front().getArgument(BBArgIndex);
+        auto BBArgValue = spatComputeBatch.getInputArgument(*inputIndex);
 
         if (BBArgValue.use_empty())
           continue;
@@ -108,7 +106,7 @@ struct MoveExtractSliceIntoCompute final : OpRewritePattern<mlir::tensor::Extrac
         }
 
         replaceAndEraseDirectComputeLikeInput(
-          rewriter, spatComputeBatch.getOperation(), BBArgIndex, mapSpatToExtract[spatComputeBatch.getOperation()]);
+          rewriter, spatComputeBatch.getOperation(), *inputIndex, mapSpatToExtract[spatComputeBatch.getOperation()]);
       }
       else {
         {
@@ -254,7 +252,7 @@ struct ArithConstToGlobalMemoryPattern final : OpRewritePattern<mlir::arith::Con
       }
     }
     else if (constantOp.getType().isIntOrIndexOrFloat()) {
-      llvm::DenseMap<Operation*, Value> mapSpatComputeToConst;
+      Value hostConstant = constantOp.getResult();
 
       for (auto& constUses : llvm::make_early_inc_range(constantOp->getUses())) {
         auto constUsers = constUses.getOwner();
@@ -264,40 +262,22 @@ struct ArithConstToGlobalMemoryPattern final : OpRewritePattern<mlir::arith::Con
           if (!inputIndex)
             return failure();
           auto BBArgIndex = *inputIndex;
-          rewriter.setInsertionPoint(&spatCompute.getBody().front().front());
-          auto newConst = rewriter.clone(*constantOp);
-
-          replaceAndEraseDirectComputeLikeInput(
-            rewriter, spatCompute.getOperation(), BBArgIndex, newConst->getResult(0));
+          replaceAndEraseDirectComputeLikeInput(rewriter, spatCompute.getOperation(), BBArgIndex, hostConstant);
         }
         else if (auto spatComputeBatch = llvm::dyn_cast<spatial::SpatComputeBatch>(constUsers)) {
           auto inputIndex = getDirectComputeLikeInputIndex(spatComputeBatch, constUses.getOperandNumber());
           if (!inputIndex)
             return failure();
           auto BBArgIndex = *inputIndex;
-          rewriter.setInsertionPoint(&spatComputeBatch.getBody().front().front());
-          auto newConst = rewriter.clone(*constantOp);
-
-          replaceAndEraseDirectComputeLikeInput(
-            rewriter, spatComputeBatch.getOperation(), BBArgIndex, newConst->getResult(0));
+          replaceAndEraseDirectComputeLikeInput(rewriter, spatComputeBatch.getOperation(), BBArgIndex, hostConstant);
         }
-        else if (auto parent = constUsers->getParentOfType<spatial::SpatCompute>()) {
-          if (!mapSpatComputeToConst.contains(parent)) {
-            rewriter.setInsertionPoint(&parent.getBody().front().front());
-            auto newConst = rewriter.clone(*constantOp);
-            mapSpatComputeToConst.insert({parent.getOperation(), newConst->getResult(0)});
-          }
-          constUses.set(mapSpatComputeToConst[parent.getOperation()]);
+        else if (constUsers->getParentOfType<spatial::SpatCompute>()) {
+          constUses.set(hostConstant);
         }
         else {
           auto batchParent = constUsers->getParentOfType<spatial::SpatComputeBatch>();
           assert(batchParent && "Global Constant used direcly not within a compute");
-          if (!mapSpatComputeToConst.contains(batchParent.getOperation())) {
-            rewriter.setInsertionPoint(&batchParent.getBody().front().front());
-            auto newConst = rewriter.clone(*constantOp);
-            mapSpatComputeToConst.insert({batchParent.getOperation(), newConst->getResult(0)});
-          }
-          constUses.set(mapSpatComputeToConst[batchParent.getOperation()]);
+          constUses.set(hostConstant);
         }
       }
     }

src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp

@@ -6,8 +6,10 @@
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/SymbolTable.h"
+#include "mlir/Transforms/FoldUtils.h"
 
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/ReturnPathNormalization.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
@@ -318,7 +320,8 @@ static LogicalResult mapIndicesThroughHelperChain(ArrayRef<int64_t> sourceIndice
   return success();
 }
 
-static void cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewriter& rewriter) {
+static void
+cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewriter& rewriter, OperationFolder& constantFolder) {
   for (Value operand : op->getOperands()) {
     if (mapping.lookupOrNull(operand))
       continue;
@@ -327,7 +330,12 @@ static void cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewri
     if (!definingOp)
       continue;
 
-    if (!isa<tensor::EmptyOp, arith::ConstantOp>(definingOp))
+    if (auto constantOp = dyn_cast<arith::ConstantOp>(definingOp)) {
+      mapping.map(operand, getOrCreateHostConstantLike(constantOp, constantFolder));
+      continue;
+    }
+
+    if (!isa<tensor::EmptyOp>(definingOp))
       continue;
 
     Operation* clonedOp = rewriter.clone(*definingOp, mapping);
@@ -337,15 +345,18 @@ static void cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewri
   }
 }
 
-static void
-cloneHelperChain(Value sourceValue, ArrayRef<Operation*> helperChain, IRRewriter& rewriter, Value& clonedValue) {
+static void cloneHelperChain(Value sourceValue,
+                             ArrayRef<Operation*> helperChain,
+                             IRRewriter& rewriter,
+                             OperationFolder& constantFolder,
+                             Value& clonedValue) {
   IRMapping mapping;
   mapping.map(sourceValue, sourceValue);
   clonedValue = sourceValue;
 
   rewriter.setInsertionPointAfterValue(sourceValue);
   for (Operation* op : helperChain) {
-    cloneMappedHelperOperands(op, mapping, rewriter);
+    cloneMappedHelperOperands(op, mapping, rewriter, constantFolder);
     Operation* clonedOp = rewriter.clone(*op, mapping);
     for (auto [originalResult, newResult] : llvm::zip(op->getResults(), clonedOp->getResults()))
       mapping.map(originalResult, newResult);
@@ -360,14 +371,19 @@ static Value emitHostCopy(IRRewriter& rewriter,
                           Value sourceValue,
                           int32_t hostTargetOffset,
                           int32_t deviceSourceOffset,
-                          int32_t sizeInBytes) {
+                          int32_t sizeInBytes,
+                          OperationFolder& constantFolder) {
+  Operation* anchorOp = sourceValue.getDefiningOp() ? sourceValue.getDefiningOp() : outputTensor.getDefiningOp();
+  assert(anchorOp && "expected a concrete op anchor for return-path host copy constants");
+  Value hostTargetOffsetValue = getOrCreateHostIndexConstant(anchorOp, hostTargetOffset, constantFolder);
+  Value deviceSourceOffsetValue = getOrCreateHostIndexConstant(anchorOp, deviceSourceOffset, constantFolder);
   return PimMemCopyDevToHostOp::create(rewriter,
                                        loc,
                                        outputTensor.getType(),
+                                       hostTargetOffsetValue,
+                                       deviceSourceOffsetValue,
                                        outputTensor,
                                        sourceValue,
-                                       rewriter.getI32IntegerAttr(hostTargetOffset),
-                                       rewriter.getI32IntegerAttr(deviceSourceOffset),
                                        rewriter.getI32IntegerAttr(sizeInBytes))
     .getOutput();
 }
@@ -411,69 +427,84 @@ void addReturnOutputBuffers(func::ReturnOp returnOp,
   }
 }
 
-ReturnPathLoweringResult lowerComputeResultReturnPath(
-  spatial::SpatCompute computeOp, OpResult result, Value yieldValue, ReturnPathState& state, IRRewriter& rewriter) {
-  Location loc = computeOp->getLoc();
-  auto yieldType = cast<TensorType>(yieldValue.getType());
+ReturnPathLoweringResult lowerProducedValueReturnPath(
+  Operation* producerOp, Value producedValue, Value storedValue, ReturnPathState& state, IRRewriter& rewriter) {
+  Location loc = producerOp->getLoc();
+  OperationFolder constantFolder(producerOp->getContext());
+  auto storedTensorType = cast<TensorType>(storedValue.getType());
 
-  if (auto returnUse = analyzeReturnUse(result)) {
-    Value storedValue = yieldValue;
-    cloneHelperChain(yieldValue, returnUse->helperChain, rewriter, storedValue);
+  if (auto returnUse = analyzeReturnUse(producedValue)) {
+    Value currentStoredValue = storedValue;
+    cloneHelperChain(storedValue, returnUse->helperChain, rewriter, constantFolder, currentStoredValue);
     for (Operation* op : returnUse->helperChain)
       markOpToRemove(state, op);
 
-    auto storedType = cast<ShapedType>(storedValue.getType());
+    auto storedType = cast<ShapedType>(currentStoredValue.getType());
     size_t elementSize = storedType.getElementTypeBitWidth() / 8;
-    if (auto storedOp = storedValue.getDefiningOp())
+    if (auto storedOp = currentStoredValue.getDefiningOp())
       rewriter.setInsertionPointAfter(storedOp);
     Value outputTensor = state.outputTensors[returnUse->returnIndex](rewriter, loc);
-    emitHostCopy(
-      rewriter, loc, outputTensor, storedValue, 0, 0, static_cast<int32_t>(storedType.getNumElements() * elementSize));
+    emitHostCopy(rewriter,
+                 loc,
+                 outputTensor,
+                 currentStoredValue,
+                 0,
+                 0,
+                 static_cast<int32_t>(storedType.getNumElements() * elementSize),
+                 constantFolder);
     return ReturnPathLoweringResult::Handled;
   }
 
-  auto resultUses = result.getUses();
+  auto resultUses = producedValue.getUses();
   if (rangeLength(resultUses) == 1) {
     OpOperand& resultUse = *resultUses.begin();
     Operation* resultUser = resultUse.getOwner();
 
     if (isa<func::ReturnOp>(resultUser)) {
       size_t resultIndexInReturn = resultUse.getOperandNumber();
-      size_t elementSize = yieldType.getElementType().getIntOrFloatBitWidth() / 8;
-      rewriter.setInsertionPointAfterValue(yieldValue);
+      size_t elementSize = storedTensorType.getElementType().getIntOrFloatBitWidth() / 8;
+      rewriter.setInsertionPointAfterValue(storedValue);
       Value outputTensor = state.outputTensors[resultIndexInReturn](rewriter, loc);
-      emitHostCopy(
-        rewriter, loc, outputTensor, yieldValue, 0, 0, static_cast<int32_t>(yieldType.getNumElements() * elementSize));
+      emitHostCopy(rewriter,
+                   loc,
+                   outputTensor,
+                   storedValue,
+                   0,
+                   0,
+                   static_cast<int32_t>(storedTensorType.getNumElements() * elementSize),
+                   constantFolder);
       return ReturnPathLoweringResult::Handled;
     }
   }
 
-  if (auto concatReturnUse = analyzeConcatReturnUse(result)) {
-    size_t elementSize = yieldType.getElementTypeBitWidth() / 8;
+  if (auto concatReturnUse = analyzeConcatReturnUse(producedValue)) {
+    size_t elementSize = storedTensorType.getElementTypeBitWidth() / 8;
     for (Operation* concatOp : concatReturnUse->concatChain)
       markOpToRemove(state, concatOp);
 
     if (concatReturnUse->helperChain.empty()) {
-      rewriter.setInsertionPointAfterValue(yieldValue);
+      rewriter.setInsertionPointAfterValue(storedValue);
       Value outputTensor = state.outputTensors[concatReturnUse->returnIndex](rewriter, loc);
       auto outputType = cast<ShapedType>(outputTensor.getType());
       int64_t flatOffset = computeFlatElementIndex(concatReturnUse->sliceOffsets, outputType.getShape());
       emitHostCopy(rewriter,
                    loc,
                    outputTensor,
-                   yieldValue,
+                   storedValue,
                    static_cast<int32_t>(flatOffset * elementSize),
                    0,
-                   static_cast<int32_t>(yieldType.getNumElements() * elementSize));
+                   static_cast<int32_t>(storedTensorType.getNumElements() * elementSize),
+                   constantFolder);
       return ReturnPathLoweringResult::Handled;
     }
 
-    auto storedType = dyn_cast<RankedTensorType>(yieldValue.getType());
+    auto storedType = dyn_cast<RankedTensorType>(storedValue.getType());
     if (!storedType) {
-      computeOp.emitOpError("has an unsupported non-ranked concat-return helper yield during Spatial-to-PIM lowering");
+      producerOp->emitOpError(
+        "has an unsupported non-ranked concat-return helper yield during Spatial-to-PIM lowering");
       return ReturnPathLoweringResult::Failure;
     }
-    rewriter.setInsertionPointAfterValue(yieldValue);
+    rewriter.setInsertionPointAfterValue(storedValue);
     Value outputTensor = state.outputTensors[concatReturnUse->returnIndex](rewriter, loc);
     auto outputType = cast<ShapedType>(outputTensor.getType());
     for (int64_t linearIndex = 0; linearIndex < storedType.getNumElements(); ++linearIndex) {
@@ -484,7 +515,7 @@ ReturnPathLoweringResult lowerComputeResultReturnPath(
       SmallVector<int64_t> destinationIndices;
       if (failed(mapIndicesThroughHelperChain(
             sourceIndices, concatReturnUse->concatShape, concatReturnUse->helperChain, destinationIndices))) {
-        computeOp.emitOpError("has an unsupported concat-return helper chain during Spatial-to-PIM lowering");
+        producerOp->emitOpError("has an unsupported concat-return helper chain during Spatial-to-PIM lowering");
         return ReturnPathLoweringResult::Failure;
       }
 
@@ -503,7 +534,7 @@ ReturnPathLoweringResult lowerComputeResultReturnPath(
       auto scalarTensorType =
         RankedTensorType::get(SmallVector<int64_t>(storedType.getRank(), 1), storedType.getElementType());
       auto elementSlice = tensor::ExtractSliceOp::create(
-        rewriter, loc, scalarTensorType, yieldValue, extractOffsets, extractSizes, extractStrides);
+        rewriter, loc, scalarTensorType, storedValue, extractOffsets, extractSizes, extractStrides);
       rewriter.setInsertionPointAfter(elementSlice);
 
       int64_t destinationFlatOffset = computeFlatElementIndex(destinationIndices, outputType.getShape());
@@ -513,7 +544,8 @@ ReturnPathLoweringResult lowerComputeResultReturnPath(
                                   elementSlice.getResult(),
                                   static_cast<int32_t>(destinationFlatOffset * elementSize),
                                   0,
-                                  static_cast<int32_t>(elementSize));
+                                  static_cast<int32_t>(elementSize),
+                                  constantFolder);
     }
     return ReturnPathLoweringResult::Handled;
   }
@@ -521,6 +553,11 @@ ReturnPathLoweringResult lowerComputeResultReturnPath(
   return ReturnPathLoweringResult::NotReturnPath;
 }
 
+ReturnPathLoweringResult lowerComputeResultReturnPath(
+  spatial::SpatCompute computeOp, OpResult result, Value yieldValue, ReturnPathState& state, IRRewriter& rewriter) {
+  return lowerProducedValueReturnPath(computeOp.getOperation(), result, yieldValue, state, rewriter);
+}
+
 void replaceReturnWithOutputBuffers(func::ReturnOp returnOp, IRRewriter& rewriter, ReturnPathState& state) {
   auto markOwnedReturnChain = [&](Operation* op, auto&& markOwnedReturnChain) -> void {
     if (!op)
@@ -569,7 +606,16 @@ void replaceReturnWithOutputBuffers(func::ReturnOp returnOp, IRRewriter& rewrite
       markOpToRemove(state, concatOp);
       for (Value operand : concatOp.getInputs())
         markOwnedReturnChain(operand.getDefiningOp(), markOwnedReturnChain);
+      return;
     }
+
+    if (auto receiveOp = dyn_cast<spatial::SpatChannelReceiveOp>(op)) {
+      markOpToRemove(state, receiveOp);
+      return;
+    }
+
+    if (auto receiveTensorOp = dyn_cast<spatial::SpatChannelReceiveTensorOp>(op))
+      markOpToRemove(state, receiveTensorOp);
   };
 
   SmallVector<Value> originalOperands(returnOp.getOperands().begin(), returnOp.getOperands().end());

src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.hpp

@@ -32,6 +32,12 @@ ReturnPathLoweringResult lowerComputeResultReturnPath(spatial::SpatCompute compu
                                                       ReturnPathState& state,
                                                       mlir::IRRewriter& rewriter);
 
+ReturnPathLoweringResult lowerProducedValueReturnPath(mlir::Operation* producerOp,
+                                                      mlir::Value producedValue,
+                                                      mlir::Value storedValue,
+                                                      ReturnPathState& state,
+                                                      mlir::IRRewriter& rewriter);
+
 void replaceReturnWithOutputBuffers(mlir::func::ReturnOp returnOp, mlir::IRRewriter& rewriter, ReturnPathState& state);
 
 } // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/SpatialToPim.td

@@ -16,8 +16,8 @@ def onnxToPimTranspose : Pat<
 >;
 
 def spatToPimVMM : Pat<
-  (SpatVMMOp:$srcOpRes $weightIndex, $vector),
-  (PimVMMOp $weightIndex, $vector,
+  (SpatVMMOp:$srcOpRes $weight, $vector),
+  (PimVMMOp $weight, $vector,
     (NativeCodeCall<"onnx_mlir::getBestOutputTensorFromOperandsOrAllocate($_builder, $0.getDefiningOp())"> $srcOpRes))
 >;
 

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

@@ -3,6 +3,7 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/SCF/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/BuiltinOps.h"
@@ -12,6 +13,8 @@
 #include "mlir/IR/SymbolTable.h"
 #include "mlir/IR/Value.h"
 #include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/FoldUtils.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/WalkPatternRewriteDriver.h"
 
 #include "llvm/ADT/StringRef.h"
@@ -104,23 +107,34 @@ static memref::GlobalOp getOrCreateZeroGlobal(IRRewriter& rewriter, Location loc
                                   IntegerAttr {});
 }
 
-static Value createZeroedDeviceHVector(IRRewriter& rewriter, Location loc, RankedTensorType tensorType) {
+static 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 zeroAttr = rewriter.getI32IntegerAttr(0);
+  auto zeroIndex = getOrCreateHostIndexConstant(outputBuffer.getOperation(), 0, constantFolder);
   auto sizeAttr = rewriter.getI32IntegerAttr(static_cast<int32_t>(getShapedTypeSizeInBytes(tensorType)));
 
   if (outputBuffer->getParentOfType<PimCoreBatchOp>())
-    return PimMemCopyHostToDevBatchOp::create(
-             rewriter, loc, tensorType, outputBuffer, zeroValue, zeroAttr, zeroAttr, sizeAttr)
+    return PimMemCopyHostToDevBatchOp::create(rewriter,
+                                              loc,
+                                              tensorType,
+                                              outputBuffer,
+                                              zeroValue,
+                                              rewriter.getI32IntegerAttr(0),
+                                              rewriter.getI32IntegerAttr(0),
+                                              sizeAttr)
       .getOutput();
 
-  return PimMemCopyHostToDevOp::create(rewriter, loc, tensorType, outputBuffer, zeroValue, zeroAttr, zeroAttr, sizeAttr)
+  return PimMemCopyHostToDevOp::create(
+           rewriter, loc, tensorType, zeroIndex, zeroIndex, outputBuffer, zeroValue, sizeAttr)
     .getOutput();
 }
 
-static Value padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector) {
+static Value
+padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector, OperationFolder& constantFolder) {
   auto vectorType = cast<RankedTensorType>(vector.getType());
   ArrayRef<int64_t> shape = vectorType.getShape();
   assert(isHVectorShape(shape) && "expected a horizontal vector");
@@ -131,7 +145,7 @@ static Value padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, V
 
   auto paddedType = RankedTensorType::get(
     {shape[0], static_cast<int64_t>(crossbarSize)}, vectorType.getElementType(), vectorType.getEncoding());
-  Value zeroed = createZeroedDeviceHVector(rewriter, loc, paddedType);
+  Value zeroed = createZeroedDeviceHVector(rewriter, loc, paddedType, constantFolder);
   auto zeroAttr = rewriter.getI32IntegerAttr(0);
   auto sizeAttr = rewriter.getI32IntegerAttr(static_cast<int32_t>(getShapedTypeSizeInBytes(vectorType)));
   return PimMemCopyOp::create(rewriter, loc, paddedType, zeroed, vector, zeroAttr, zeroAttr, sizeAttr).getOutput();
@@ -151,6 +165,7 @@ void SpatialToPimPass::runOnOperation() {
   func::FuncOp funcOp = *entryFunc;
 
   IRRewriter rewriter(&getContext());
+  OperationFolder constantFolder(&getContext());
 
   ConversionTarget target(*ctx);
   target.addLegalDialect<PimDialect,
@@ -181,16 +196,17 @@ void SpatialToPimPass::runOnOperation() {
   RewritePatternSet globalTensorPatterns(ctx);
   populateGlobalTensorMaterializationPatterns(globalTensorPatterns);
   walkAndApplyPatterns(moduleOp, std::move(globalTensorPatterns));
-  auto returnOp = cast<func::ReturnOp>(funcOp.front().getTerminator());
 
+  auto returnOp = cast<func::ReturnOp>(funcOp.front().getTerminator());
   addReturnOutputBuffers(returnOp, rewriter, outputTensors);
+  ReturnPathState returnPathState {outputTensors, operationsToRemove};
   if (failed(allocateAndInitializeCoreLocalVariables(funcOp, rewriter))) {
     funcOp.emitOpError("failed to allocate or initialize core-local tensors during Spatial-to-PIM lowering");
     signalPassFailure();
     return;
   }
 
-  CoreLoweringState coreLoweringState {coreId, outputTensors, operationsToRemove};
+  CoreLoweringState coreLoweringState {coreId, outputTensors, operationsToRemove, constantFolder};
   for (auto computeOp : funcOp.getOps<spatial::SpatCompute>()) {
     markOpToRemove(computeOp);
     if (failed(lowerComputeOp(computeOp, coreLoweringState, rewriter))) {
@@ -251,7 +267,6 @@ void SpatialToPimPass::runOnOperation() {
   }
 
   enlargeVMMOutTensorsToCrossbarSize(funcOp, rewriter);
-  ReturnPathState returnPathState {outputTensors, operationsToRemove};
   replaceReturnWithOutputBuffers(returnOp, rewriter, returnPathState);
 
   SmallVector<Operation*> pendingRemovals(operationsToRemove.begin(), operationsToRemove.end());
@@ -302,6 +317,7 @@ void SpatialToPimPass::runOnOperation() {
 }
 
 void SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func::FuncOp funcOp, IRRewriter& rewriter) {
+  OperationFolder constantFolder(funcOp.getContext());
   funcOp.walk([&](PimVMMOp vmmOp) {
     auto outputType = cast<RankedTensorType>(vmmOp.getOutput().getType());
     ArrayRef<int64_t> outputShape = outputType.getShape();
@@ -309,7 +325,7 @@ void SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func::FuncOp funcOp, I
     assert(outputShape[1] <= static_cast<int64_t>(crossbarSize) && "output width must fit in one crossbar");
 
     rewriter.setInsertionPoint(vmmOp);
-    Value paddedInput = padHVectorInputToCrossbarSize(rewriter, vmmOp.getLoc(), vmmOp.getInput());
+    Value paddedInput = padHVectorInputToCrossbarSize(rewriter, vmmOp.getLoc(), vmmOp.getInput(), constantFolder);
     auto paddedOutputType = RankedTensorType::get(
       {outputShape[0], static_cast<int64_t>(crossbarSize)}, outputType.getElementType(), outputType.getEncoding());
     Value paddedOutputBuffer = outputShape[1] == static_cast<int64_t>(crossbarSize)
@@ -336,10 +352,13 @@ void SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func::FuncOp funcOp, I
 
 LogicalResult SpatialToPimPass::allocateAndInitializeCoreLocalVariables(func::FuncOp funcOp, IRRewriter& rewriter) {
   Location loc = funcOp.getLoc();
+  OperationFolder constantFolder(funcOp.getContext());
 
   auto insertMemCopyHostToDev = [&](Value inputTensor, int64_t elementsOffset) {
     auto tensorType = cast<ShapedType>(inputTensor.getType());
     Type elementType = tensorType.getElementType();
+    if (!elementType.isIntOrFloat())
+      return;
     size_t elementByteSize = elementType.getIntOrFloatBitWidth() / 8;
     rewriter.setInsertionPointAfter(inputTensor.getDefiningOp());
 
@@ -349,10 +368,11 @@ LogicalResult SpatialToPimPass::allocateAndInitializeCoreLocalVariables(func::Fu
       rewriter,
       loc,
       tensorType,
+      getOrCreateHostIndexConstant(deviceTensor.getOperation(), 0, constantFolder),
+      getOrCreateHostIndexConstant(
+        deviceTensor.getOperation(), static_cast<int64_t>(elementsOffset * elementByteSize), constantFolder),
       deviceTensor,
       inputTensor,
-      rewriter.getI32IntegerAttr(0),
-      rewriter.getI32IntegerAttr(static_cast<int32_t>(elementsOffset * elementByteSize)),
       rewriter.getI32IntegerAttr(static_cast<int32_t>(tensorType.getNumElements() * elementByteSize)));
 
     rewriter.replaceAllUsesExcept(inputTensor, memCopyHostToDevOp.getResult(), {memCopyHostToDevOp});