huge refactor for high RewritePatterns usage and less ad-hoc cpp code

remove Spatial many ops in favor of tensor ops like in pim

src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt

@@ -3,6 +3,11 @@ mlir_tablegen(ONNXToSpatial.hpp.inc -gen-rewriters "-I${ONNX_MLIR_SRC_ROOT}")
 add_public_tablegen_target(ONNXToSpatialIncGen)
 
 add_pim_library(OMONNXToSpatial
+  ConversionPatterns.cpp
+  HostFoldability.cpp
+  HostLegality.cpp
+  PrePatterns.cpp
+  PostPatterns.cpp
   Patterns/Math/Conv.cpp
   Patterns/Math/Elementwise.cpp
   Patterns/Math/Gemm.cpp

src/PIM/Conversion/ONNXToSpatial/Common/Common.hpp

@@ -1,8 +1,7 @@
 #pragma once
 
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Dialect/ONNX/ONNXOps.hpp"
-
 #include "ComputeRegionBuilder.hpp"
 #include "ShapeTilingUtils.hpp"
 #include "WeightMaterialization.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"

src/PIM/Conversion/ONNXToSpatial/Common/ShapeTilingUtils.cpp

@@ -5,6 +5,8 @@
 
 #include "ShapeTilingUtils.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostFoldability.hpp"
 
 using namespace mlir;
 
@@ -30,10 +32,29 @@ SmallVector<Value> sliceTensor(
 
   for (int64_t i = 0; i < numSlices; i++) {
     offsets[axis] = rewriter.getIndexAttr(i * sliceSize);
-    if (i == numSlices - 1 && lastSliceSize != 0)
+    int64_t currentSliceSize = sliceSize;
+    if (i == numSlices - 1 && lastSliceSize != 0) {
+      currentSliceSize = lastSliceSize;
       sizes[axis] = rewriter.getIndexAttr(lastSliceSize);
+    }
 
-    Value slice = tensor::ExtractSliceOp::create(rewriter, loc, tensorToSlice, offsets, sizes, strides);
+    SmallVector<int64_t> sliceShape(shape.begin(), shape.end());
+    sliceShape[axis] = currentSliceSize;
+    auto sliceType =
+      RankedTensorType::get(sliceShape, cast<RankedTensorType>(tensorToSlice.getType()).getElementType());
+
+    Value slice;
+    if (isHostFoldableValue(tensorToSlice)) {
+      slice = tensor::ExtractSliceOp::create(rewriter, loc, tensorToSlice, offsets, sizes, strides);
+    }
+    else {
+      auto sliceCompute =
+        createSpatCompute<1>(rewriter, loc, TypeRange {sliceType}, {}, ValueRange {tensorToSlice}, [&](Value input) {
+          Value computedSlice = tensor::ExtractSliceOp::create(rewriter, loc, input, offsets, sizes, strides);
+          spatial::SpatYieldOp::create(rewriter, loc, computedSlice);
+        });
+      slice = sliceCompute.getResult(0);
+    }
     slices.push_back(slice);
   }
 

src/PIM/Conversion/ONNXToSpatial/Common/ShapeTilingUtils.hpp

@@ -5,15 +5,15 @@
 #include "mlir/IR/Value.h"
 #include "mlir/Transforms/DialectConversion.h"
 
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+
 #include <cassert>
 #include <cstddef>
 #include <type_traits>
 #include <utility>
 
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/SmallVector.h"
-
 namespace onnx_mlir {
 
 template <class ShapedType>
@@ -105,7 +105,8 @@ inline auto getTensorShape(mlir::Value tensor) {
 inline bool haveSameStaticShape(mlir::Value lhs, mlir::Value rhs) {
   auto lhsType = mlir::dyn_cast<mlir::RankedTensorType>(lhs.getType());
   auto rhsType = mlir::dyn_cast<mlir::RankedTensorType>(rhs.getType());
-  return lhsType && rhsType && lhsType.hasStaticShape() && rhsType.hasStaticShape() && lhsType.getShape() == rhsType.getShape();
+  return lhsType && rhsType && lhsType.hasStaticShape() && rhsType.hasStaticShape()
+      && lhsType.getShape() == rhsType.getShape();
 }
 
 /// Slices a statically shaped tensor along one axis into contiguous pieces of

src/PIM/Conversion/ONNXToSpatial/Common/WeightMaterialization.cpp

@@ -5,12 +5,12 @@
 #include "mlir/IR/Value.h"
 #include "mlir/Support/LogicalResult.h"
 
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/ADT/STLExtras.h"
 
-#include "WeightMaterialization.hpp"
 #include "ShapeTilingUtils.hpp"
+#include "WeightMaterialization.hpp"
 #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -28,7 +28,7 @@ bool isWeightLikeComputeOperand(Value value) {
   while (auto* definingOp = value.getDefiningOp()) {
     if (!visited.insert(definingOp).second)
       return false;
-    if (hasWeightAlways(definingOp))
+    if (isa<arith::ConstantOp, ONNXConstantOp>(definingOp) || hasWeightAlways(definingOp))
       return true;
 
     if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(definingOp)) {

src/PIM/Conversion/ONNXToSpatial/ConversionPatterns.cpp

@@ -0,0 +1,32 @@
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+namespace {
+
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatial.hpp.inc"
+
+} // namespace
+
+void populateConversionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx) {
+  patterns.add<removeLRN>(ctx);
+
+  populateElementwisePatterns(patterns, ctx);
+  populateGemmPatterns(patterns, ctx);
+  populateConvPatterns(patterns, ctx);
+  populatePoolPatterns(patterns, ctx);
+  populateReduceMeanPatterns(patterns, ctx);
+  populateReluPatterns(patterns, ctx);
+  populateSigmoidPatterns(patterns, ctx);
+  populateSoftmaxPatterns(patterns, ctx);
+  populateConcatPatterns(patterns, ctx);
+  populateGatherPatterns(patterns, ctx);
+  populateResizePatterns(patterns, ctx);
+  populateReshapePatterns(patterns, ctx);
+  populateSplitPatterns(patterns, ctx);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp

@@ -5,6 +5,8 @@
 
 namespace onnx_mlir {
 
+void populateConversionPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+
 void populateConvPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 
 void populateElementwisePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);

src/PIM/Conversion/ONNXToSpatial/HostFoldability.cpp

@@ -0,0 +1,75 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/BuiltinTypes.h"
+
+#include "llvm/ADT/SmallPtrSet.h"
+
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostFoldability.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+namespace {
+
+static bool hasStaticUnitStrides(tensor::ExtractSliceOp extractSliceOp) {
+  return llvm::all_of(extractSliceOp.getStaticStrides(), [](int64_t stride) { return stride == 1; });
+}
+
+static bool isStaticTensorResult(Operation* op) {
+  return llvm::all_of(op->getResultTypes(), [](Type type) {
+    auto shapedType = dyn_cast<ShapedType>(type);
+    return shapedType && shapedType.hasStaticShape();
+  });
+}
+
+static bool isHostFoldableOpImpl(Operation* op, llvm::SmallPtrSetImpl<Operation*>& visited) {
+  if (!op || !visited.insert(op).second)
+    return false;
+
+  if (isa<arith::ConstantOp, ONNXConstantOp, ONNXNoneOp>(op))
+    return true;
+
+  if (!isStaticTensorResult(op))
+    return false;
+
+  if (auto transposeOp = dyn_cast<ONNXTransposeOp>(op))
+    return isHostFoldableValue(transposeOp.getData());
+
+  if (auto collapseShapeOp = dyn_cast<tensor::CollapseShapeOp>(op))
+    return isHostFoldableValue(collapseShapeOp.getSrc());
+
+  if (auto expandShapeOp = dyn_cast<tensor::ExpandShapeOp>(op))
+    return isHostFoldableValue(expandShapeOp.getSrc());
+
+  if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(op))
+    return hasStaticUnitStrides(extractSliceOp) && isHostFoldableValue(extractSliceOp.getSource());
+
+  if (auto extractRowsOp = dyn_cast<spatial::SpatExtractRowsOp>(op))
+    return isHostFoldableValue(extractRowsOp.getInput());
+
+  if (auto concatOp = dyn_cast<spatial::SpatConcatOp>(op))
+    return llvm::all_of(concatOp.getInputs(), isHostFoldableValue);
+
+  return false;
+}
+
+} // namespace
+
+bool isHostFoldableValue(Value value) {
+  auto* definingOp = value.getDefiningOp();
+  if (!definingOp)
+    return false;
+
+  llvm::SmallPtrSet<Operation*, 8> visited;
+  return isHostFoldableOpImpl(definingOp, visited);
+}
+
+bool isHostFoldableOp(Operation* op) {
+  llvm::SmallPtrSet<Operation*, 8> visited;
+  return isHostFoldableOpImpl(op, visited);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/HostFoldability.hpp

@@ -0,0 +1,12 @@
+#pragma once
+
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/Value.h"
+
+namespace onnx_mlir {
+
+bool isHostFoldableValue(mlir::Value value);
+
+bool isHostFoldableOp(mlir::Operation* op);
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/HostLegality.cpp

@@ -0,0 +1,29 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostFoldability.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostLegality.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+LogicalResult verifyONNXToSpatialHostLegality(func::FuncOp funcOp) {
+  bool hasFailure = false;
+
+  for (Operation& op : funcOp.getFunctionBody().front()) {
+    if (isa<func::ReturnOp, spatial::SpatCompute, spatial::SpatComputeBatch>(&op))
+      continue;
+    if (isHostFoldableOp(&op))
+      continue;
+
+    op.emitOpError("non-foldable top-level runtime op remains after ONNX-to-Spatial; lower it inside spat.compute");
+    hasFailure = true;
+  }
+
+  return success(!hasFailure);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/HostLegality.hpp

@@ -0,0 +1,10 @@
+#pragma once
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Support/LogicalResult.h"
+
+namespace onnx_mlir {
+
+mlir::LogicalResult verifyONNXToSpatialHostLegality(mlir::func::FuncOp funcOp);
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -8,21 +8,17 @@
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "mlir/Transforms/Passes.h"
 
-#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/Casting.h"
 #include "llvm/Support/Debug.h"
-#include "llvm/Support/raw_os_ostream.h"
-
-#include <fstream>
-#include <iterator>
-#include <utility>
 
 #include "Common/Common.hpp"
 #include "Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostLegality.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PostPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PrePatterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Compiler/CompilerOptions.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -33,8 +29,6 @@ namespace onnx_mlir {
 
 namespace {
 
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatial.hpp.inc"
-
 struct ONNXToSpatialPass : PassWrapper<ONNXToSpatialPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ONNXToSpatialPass)
   StringRef getArgument() const override { return "convert-onnx-to-spatial"; }
@@ -44,71 +38,64 @@ struct ONNXToSpatialPass : PassWrapper<ONNXToSpatialPass, OperationPass<ModuleOp
   ONNXToSpatialPass(const ONNXToSpatialPass& pass) {}
 
   void runOnOperation() override;
-
-private:
-  void annotateWeightsConstants(func::FuncOp funcOp) const;
-  LogicalResult encapsulateGlobalInstruction(func::FuncOp funcOp);
-  LogicalResult promoteConstantInputsToWeights(func::FuncOp funcOp);
-  void populateEmptyFunction(func::FuncOp funcOp);
 };
 
 } // namespace
 
-static void foldSingleLaneComputeBatches(func::FuncOp funcOp) {
+static void populateEmptyFunction(func::FuncOp funcOp) {
   IRRewriter rewriter(funcOp.getContext());
-  SmallVector<spatial::SpatComputeBatch> batchOps;
-  funcOp.walk([&](spatial::SpatComputeBatch batchOp) { batchOps.push_back(batchOp); });
+  IRMapping mapper;
+  SmallVector<spatial::SpatCompute> computes(funcOp.getOps<spatial::SpatCompute>());
+  if (!computes.empty())
+    return;
 
-  for (auto batchOp : batchOps) {
-    if (batchOp.getLaneCount() != 1)
-      continue;
+  auto returnOp = cast<func::ReturnOp>(funcOp.getFunctionBody().front().getTerminator());
+  rewriter.setInsertionPoint(returnOp);
 
-    auto loc = batchOp.getLoc();
-    rewriter.setInsertionPoint(batchOp);
-    auto computeOp =
-      spatial::SpatCompute::create(rewriter, loc, batchOp.getResultTypes(), batchOp.getWeights(), batchOp.getInputs());
-    computeOp.getProperties().setOperandSegmentSizes(
-      {static_cast<int>(batchOp.getWeights().size()), static_cast<int>(batchOp.getInputs().size())});
+  SmallVector<Type> sourceTypes;
+  SmallVector<Location> sourceLocs;
+  sourceTypes.reserve(funcOp.getNumArguments());
+  sourceLocs.reserve(funcOp.getNumArguments());
+  for (Value source : funcOp.getArguments()) {
+    sourceTypes.push_back(source.getType());
+    sourceLocs.push_back(source.getLoc());
+  }
 
-    Block& templateBlock = batchOp.getBody().front();
-    SmallVector<Type> blockArgTypes;
-    SmallVector<Location> blockArgLocs;
-    for (BlockArgument arg : templateBlock.getArguments()) {
-      blockArgTypes.push_back(arg.getType());
-      blockArgLocs.push_back(loc);
-    }
-    auto* newBlock =
-      rewriter.createBlock(&computeOp.getBody(), computeOp.getBody().end(), TypeRange(blockArgTypes), blockArgLocs);
+  auto newCompute = spatial::SpatCompute::create(
+    rewriter, returnOp.getLoc(), returnOp.getOperandTypes(), funcOp.getArguments(), {}, {});
+  auto* newBlock = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), sourceTypes, sourceLocs);
+  for (auto [blockArg, computeArg] : llvm::zip(newBlock->getArguments(), newCompute.getOperands()))
+    mapper.map(computeArg, blockArg);
+  newCompute.getProperties().setOperandSegmentSizes({0, static_cast<int>(sourceTypes.size())});
 
-    IRMapping mapper;
-    for (auto [oldArg, newArg] : llvm::zip(templateBlock.getArguments(), newBlock->getArguments()))
-      mapper.map(oldArg, newArg);
-    rewriter.setInsertionPointToEnd(newBlock);
-    for (Operation& op : templateBlock)
+  rewriter.setInsertionPointToEnd(newBlock);
+  for (Operation& op : funcOp.getOps())
+    if (!isa<spatial::SpatCompute, func::ReturnOp>(&op))
       rewriter.clone(op, mapper);
 
-    batchOp.replaceAllUsesWith(computeOp.getResults());
-    rewriter.eraseOp(batchOp);
-  }
+  auto yield = spatial::SpatYieldOp::create(rewriter, funcOp.getLoc(), returnOp.getOperands());
+  for (size_t i = 0; i < yield.getNumOperands(); ++i)
+    yield.setOperand(i, mapper.lookupOrDefault(yield.getOperand(i)));
+
+  for (Operation& op : llvm::make_early_inc_range(funcOp.getOps()))
+    if (!isa<spatial::SpatCompute, func::ReturnOp>(&op)) {
+      op.dropAllUses();
+      rewriter.eraseOp(&op);
+    }
+
+  for (auto [index, computeResult] : llvm::enumerate(newCompute.getResults()))
+    returnOp.setOperand(index, computeResult);
 }
 
 void ONNXToSpatialPass::runOnOperation() {
   ModuleOp moduleOp = getOperation();
   MLIRContext* ctx = &getContext();
 
-  RewritePatternSet mergeActivationPatterns(ctx);
-  mergeActivationPatterns.add<onnxToArithConstant>(ctx);
-  mergeActivationPatterns.add<convAddToConvWithBiasLeft>(ctx);
-  mergeActivationPatterns.add<convAddToConvWithBiasRight>(ctx);
-  mergeActivationPatterns.add<matMulAddToGemm>(ctx);
-  mergeActivationPatterns.add<matMulToGemm>(ctx);
-  mergeActivationPatterns.add<removeFlattenSameShape>(ctx);
-  populateMatMulRewritePatterns(mergeActivationPatterns, ctx);
+  RewritePatternSet prePatterns(ctx);
+  populatePrePatterns(prePatterns, ctx);
+  if (failed(applyPatternsGreedily(moduleOp, std::move(prePatterns))))
+    llvm::dbgs() << "Failed to apply pre-patterns, continuing...\n";
 
-  if (failed(applyPatternsGreedily(moduleOp, std::move(mergeActivationPatterns))))
-    llvm::dbgs() << "Failed to merge activation patterns, continuing...\n";
-
-  IRRewriter rewriter(moduleOp);
   auto entryFunc = getPimEntryFunc(moduleOp);
   if (failed(entryFunc)) {
     signalPassFailure();
@@ -140,34 +127,23 @@ void ONNXToSpatialPass::runOnOperation() {
   target.addIllegalOp<ONNXReduceMeanV13Op>();
   target.addIllegalOp<ONNXSplitOp>();
 
-  RewritePatternSet patterns(ctx);
-  patterns.add<removeLRN>(ctx);
-
-  populateElementwisePatterns(patterns, ctx);
-  populateGemmPatterns(patterns, ctx);
-  populateConvPatterns(patterns, ctx);
-  populatePoolPatterns(patterns, ctx);
-  populateReduceMeanPatterns(patterns, ctx);
-  populateReluPatterns(patterns, ctx);
-  populateSigmoidPatterns(patterns, ctx);
-  populateSoftmaxPatterns(patterns, ctx);
-  populateConcatPatterns(patterns, ctx);
-  populateGatherPatterns(patterns, ctx);
-  populateResizePatterns(patterns, ctx);
-  populateReshapePatterns(patterns, ctx);
-  populateSplitPatterns(patterns, ctx);
-
-  if (failed(applyPartialConversion(moduleOp, target, std::move(patterns)))) {
+  RewritePatternSet conversionPatterns(ctx);
+  populateConversionPatterns(conversionPatterns, ctx);
+  if (failed(applyPartialConversion(moduleOp, target, std::move(conversionPatterns)))) {
     signalPassFailure();
     return;
   }
 
-  foldSingleLaneComputeBatches(*entryFunc);
+  RewritePatternSet earlyPostPatterns(ctx);
+  populateEarlyPostPatterns(earlyPostPatterns, ctx);
+  if (failed(applyPatternsGreedily(*entryFunc, std::move(earlyPostPatterns)))) {
+    signalPassFailure();
+    return;
+  }
 
-  // Count the number of compute ops and check they do not exceed the core count
   if (coresCount != -1) {
     int computeOpsCount = 0;
-    for (auto& op : entryFunc->getFunctionBody().front().getOperations())
+    for (Operation& op : entryFunc->getFunctionBody().front().getOperations())
       if (isa<spatial::SpatCompute>(op))
         computeOpsCount++;
 
@@ -185,355 +161,23 @@ void ONNXToSpatialPass::runOnOperation() {
 
   annotateWeightsConstants(*entryFunc);
 
+  RewritePatternSet postPatterns(ctx);
+  populatePostPatterns(postPatterns, ctx);
+  if (failed(applyPatternsGreedily(*entryFunc, std::move(postPatterns)))) {
+    signalPassFailure();
+    return;
+  }
+
+  if (failed(verifyONNXToSpatialHostLegality(*entryFunc))) {
+    signalPassFailure();
+    return;
+  }
+
   populateEmptyFunction(*entryFunc);
 
-  if (failed(encapsulateGlobalInstruction(*entryFunc))) {
-    signalPassFailure();
-    return;
-  }
-
-  if (failed(promoteConstantInputsToWeights(*entryFunc))) {
-    signalPassFailure();
-    return;
-  }
-
-  // Dump to file for debug
   dumpModule(moduleOp, "spatial0");
 }
 
-template <typename T>
-bool encapsulator(IRRewriter& rewriter, Location loc, Operation* inst, std::function<Value(T)> funcSource) {
-  if (T toRemoveOp = llvm::dyn_cast_if_present<T>(inst)) {
-    Value source = funcSource(toRemoveOp);
-    rewriter.setInsertionPointAfter(toRemoveOp);
-    auto newCompute = spatial::SpatCompute::create(rewriter, loc, inst->getResultTypes(), source);
-    auto BB = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), {source.getType()}, {loc});
-    newCompute.getProperties().setOperandSegmentSizes({(int) 0, (int) 1});
-    rewriter.setInsertionPointToEnd(BB);
-    IRMapping mapper;
-    mapper.map(source, BB->getArgument(0));
-    auto newInst = rewriter.clone(*inst, mapper);
-    spatial::SpatYieldOp::create(rewriter, loc, newInst->getResults());
-    inst->replaceAllUsesWith(newCompute->getResults());
-    inst->erase();
-    return true;
-  }
-  return false;
-}
-
-bool encapsulateSlice(IRRewriter& rewriter, Location loc, Operation* inst) {
-  if (tensor::ExtractSliceOp toRemoveOp = llvm::dyn_cast_if_present<tensor::ExtractSliceOp>(inst)) {
-    auto source = toRemoveOp.getSource();
-    rewriter.setInsertionPointAfter(toRemoveOp);
-    auto newCompute = spatial::SpatCompute::create(rewriter, loc, inst->getResultTypes(), source);
-    auto BB = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), {source.getType()}, {loc});
-    newCompute.getProperties().setOperandSegmentSizes({(int) 0, (int) 1});
-    rewriter.setInsertionPointToEnd(BB);
-    IRMapping mapper;
-    mapper.map(source, BB->getArgument(0));
-    auto newInst = rewriter.clone(*inst, mapper);
-    spatial::SpatYieldOp::create(rewriter, loc, newInst->getResults());
-    inst->replaceAllUsesWith(newCompute->getResults());
-    inst->erase();
-    return true;
-  }
-  return false;
-}
-
-bool encapsulateConcat(IRRewriter& rewriter, Location loc, Operation* inst) {
-  if (auto toRemoveOp = llvm::dyn_cast_if_present<tensor::ConcatOp>(inst)) {
-    auto sources = toRemoveOp.getInputs();
-    rewriter.setInsertionPointAfter(toRemoveOp);
-    if (llvm::any_of(sources,
-                     [](auto source) { return isa_and_present<spatial::SpatCompute>(source.getDefiningOp()); })) {
-      auto newCompute = spatial::SpatCompute::create(rewriter, loc, inst->getResultTypes(), sources);
-      SmallVector<Type> sourceTypes;
-      SmallVector<Location> sourceLoc;
-      for (auto source : sources) {
-        sourceTypes.push_back(source.getType());
-        sourceLoc.push_back(loc);
-      }
-      auto BB = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), sourceTypes, sourceLoc);
-      newCompute.getProperties().setOperandSegmentSizes({(int) 0, (int) sources.size()});
-      rewriter.setInsertionPointToEnd(BB);
-      IRMapping mapper;
-      for (auto [source, bbArg] : llvm::zip(sources, BB->getArguments()))
-        mapper.map(source, bbArg);
-      auto newConcat = spatial::SpatConcatOp::create(rewriter,
-                                                     loc,
-                                                     toRemoveOp.getType(),
-                                                     rewriter.getI64IntegerAttr(toRemoveOp.getDim()),
-                                                     ValueRange(BB->getArguments()));
-      spatial::SpatYieldOp::create(rewriter, loc, newConcat.getOutput());
-      inst->replaceAllUsesWith(newCompute->getResults());
-      inst->erase();
-      return true;
-    }
-    auto newCompute = spatial::SpatCompute::create(rewriter, loc, inst->getResultTypes(), sources);
-    SmallVector<Type> sourceTypes;
-    SmallVector<Location> sourceLoc;
-    for (auto source : sources) {
-      sourceTypes.push_back(source.getType());
-      sourceLoc.push_back(loc);
-    }
-    auto BB = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), sourceTypes, sourceLoc);
-    newCompute.getProperties().setOperandSegmentSizes({(int) 0, (int) sources.size()});
-    rewriter.setInsertionPointToEnd(BB);
-    IRMapping mapper;
-    for (auto [source, bbArg] : llvm::zip(sources, BB->getArguments()))
-      mapper.map(source, bbArg);
-    auto newConcat = rewriter.clone(*inst, mapper);
-    spatial::SpatYieldOp::create(rewriter, loc, newConcat->getResults());
-    inst->replaceAllUsesWith(newCompute->getResults());
-    inst->erase();
-    return true;
-  }
-  return false;
-}
-
-static FailureOr<bool> sourceOperandHasWeightAlways(Operation* op) {
-  if (op == nullptr)
-    return false;
-
-  Operation* source = nullptr;
-  do {
-
-    if (isa<spatial::SpatCompute, spatial::SpatComputeBatch>(*op)) {
-      return false;
-    }
-    else if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(*op)) {
-      auto tmpSource = extractSliceOp.getSource();
-      auto definingOp = tmpSource.getDefiningOp();
-      if (definingOp)
-        op = definingOp;
-      else
-        return false;
-    }
-    else if (auto extractRowsOp = dyn_cast<spatial::SpatExtractRowsOp>(*op)) {
-      auto tmpSource = extractRowsOp.getInput();
-      auto definingOp = tmpSource.getDefiningOp();
-      if (definingOp)
-        op = definingOp;
-      else
-        return false;
-    }
-    else if (auto expandShapeOp = dyn_cast<tensor::ExpandShapeOp>(*op)) {
-      auto tmpSource = expandShapeOp.getSrc();
-      auto definingOp = tmpSource.getDefiningOp();
-      if (definingOp)
-        op = definingOp;
-      else
-        return false;
-    }
-    else if (auto transposeOp = dyn_cast<ONNXTransposeOp>(*op)) {
-      auto tmpSource = transposeOp.getData();
-      auto definingOp = tmpSource.getDefiningOp();
-      if (definingOp)
-        op = definingOp;
-      else
-        return false;
-    }
-    else if (auto collapseShapeOp = dyn_cast<tensor::CollapseShapeOp>(*op)) {
-      auto tmpSource = collapseShapeOp.getSrc();
-      auto definingOp = tmpSource.getDefiningOp();
-      if (definingOp)
-        op = definingOp;
-      else
-        return false;
-    }
-    else if (auto constantOp = dyn_cast<arith::ConstantOp>(*op)) {
-      source = constantOp;
-    }
-    else if (auto concatOp = dyn_cast<tensor::ConcatOp>(*op)) {
-      bool res = false;
-      for (auto operand : concatOp.getOperands()) {
-        res |= hasWeightAlways(operand.getDefiningOp());
-        if (res)
-          return res;
-      }
-      return res;
-    }
-    else if (auto concatOp = dyn_cast<spatial::SpatConcatOp>(*op)) {
-      bool res = false;
-      for (auto operand : concatOp.getOperands()) {
-        res |= hasWeightAlways(operand.getDefiningOp());
-        if (res)
-          return res;
-      }
-      return res;
-    }
-    else {
-      op->emitOpError("unsupported global instruction while promoting weight-backed operands into Spatial computes");
-      return failure();
-    }
-  }
-  while (source == nullptr);
-
-  return hasWeightAlways(source);
-}
-
-// TODO what we want to keep in global?
-LogicalResult ONNXToSpatialPass::encapsulateGlobalInstruction(func::FuncOp funcOp) {
-  Location loc = funcOp.getLoc();
-  IRRewriter rewriter(&getContext());
-  bool keep = true;
-  while (keep) {
-    keep = false;
-    for (auto& instruction : llvm::make_early_inc_range(funcOp.getOps())) {
-      if (isa<spatial::SpatCompute, spatial::SpatComputeBatch, spatial::SpatExtractRowsOp>(instruction)
-          || isa<func::ReturnOp>(instruction))
-        continue;
-
-      auto weightBacked = sourceOperandHasWeightAlways(&instruction);
-      if (failed(weightBacked))
-        return failure();
-      if (*weightBacked)
-        continue;
-
-      keep |= encapsulateSlice(rewriter, loc, &instruction);
-
-      keep |= encapsulator<tensor::ExpandShapeOp>(
-        rewriter, loc, &instruction, [](tensor::ExpandShapeOp expand) { return expand.getSrc(); });
-
-      keep |= encapsulator<ONNXTransposeOp>(
-        rewriter, loc, &instruction, [](ONNXTransposeOp transpose) { return transpose.getData(); });
-
-      keep |= encapsulator<tensor::CollapseShapeOp>(
-        rewriter, loc, &instruction, [](tensor::CollapseShapeOp collapse) { return collapse.getSrc(); });
-
-      keep |= encapsulateConcat(rewriter, loc, &instruction);
-    }
-  }
-  return success();
-}
-
-void ONNXToSpatialPass::annotateWeightsConstants(func::FuncOp funcOp) const {
-  funcOp.walk([&](arith::ConstantOp constantOp) {
-    if (hasOnlySpatialMvmVmmWeightUses(constantOp.getResult()))
-      markWeightAlways(constantOp);
-  });
-}
-
-LogicalResult ONNXToSpatialPass::promoteConstantInputsToWeights(func::FuncOp funcOp) {
-  IRRewriter rewriter(&getContext());
-  SmallVector<spatial::SpatCompute> computes(funcOp.getOps<spatial::SpatCompute>());
-
-  for (auto compute : computes) {
-    SmallVector<bool> promoteInput(compute.getInputs().size(), false);
-    bool needsRewrite = false;
-    for (auto [inputIdx, input] : llvm::enumerate(compute.getInputs())) {
-      if (!isWeightLikeComputeOperand(input))
-        continue;
-      promoteInput[inputIdx] = true;
-      needsRewrite = true;
-    }
-    if (!needsRewrite)
-      continue;
-
-    rewriter.setInsertionPointAfter(compute);
-
-    SmallVector<Value> newWeights(compute.getWeights().begin(), compute.getWeights().end());
-    SmallVector<Value> newInputs;
-    SmallVector<Type> newInputTypes;
-    SmallVector<Location> newInputLocs;
-    newWeights.reserve(compute.getWeights().size() + compute.getInputs().size());
-    newInputs.reserve(compute.getInputs().size());
-    newInputTypes.reserve(compute.getInputs().size());
-    newInputLocs.reserve(compute.getInputs().size());
-
-    for (auto [inputIdx, input] : llvm::enumerate(compute.getInputs())) {
-      if (promoteInput[inputIdx]) {
-        newWeights.push_back(input);
-        continue;
-      }
-      newInputs.push_back(input);
-      newInputTypes.push_back(input.getType());
-      newInputLocs.push_back(input.getLoc());
-    }
-
-    auto newCompute =
-      spatial::SpatCompute::create(rewriter, compute.getLoc(), compute.getResultTypes(), newWeights, newInputs);
-    auto* newBlock =
-      rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), newInputTypes, newInputLocs);
-    newCompute.getProperties().setOperandSegmentSizes(
-      {static_cast<int>(newWeights.size()), static_cast<int>(newInputs.size())});
-    rewriter.setInsertionPointToStart(newBlock);
-
-    IRMapping mapper;
-    auto& oldBlock = compute.getBody().front();
-    size_t newInputIdx = 0;
-    for (auto [oldInputIdx, oldArg] : llvm::enumerate(oldBlock.getArguments())) {
-      if (!promoteInput[oldInputIdx]) {
-        mapper.map(oldArg, newBlock->getArgument(newInputIdx++));
-        continue;
-      }
-
-      auto clonedValue = materializeWeightLikeValueInBlock(compute.getInputs()[oldInputIdx], rewriter, mapper);
-      if (failed(clonedValue))
-        return compute.emitError("failed to materialize promoted weight-like operand inside compute body");
-      mapper.map(oldArg, *clonedValue);
-    }
-
-    for (auto& op : oldBlock.without_terminator())
-      rewriter.clone(op, mapper);
-
-    auto oldYield = cast<spatial::SpatYieldOp>(oldBlock.getTerminator());
-    SmallVector<Value> newYieldOperands;
-    newYieldOperands.reserve(oldYield.getOutputs().size());
-    for (Value operand : oldYield.getOutputs()) {
-      auto mapped = mapper.lookupOrNull(operand);
-      newYieldOperands.push_back(mapped ? cast<Value>(mapped) : operand);
-    }
-    spatial::SpatYieldOp::create(rewriter, oldYield.getLoc(), newYieldOperands);
-
-    compute.replaceAllUsesWith(newCompute);
-    compute.erase();
-  }
-
-  return success();
-}
-
-void ONNXToSpatialPass::populateEmptyFunction(func::FuncOp funcOp) {
-  IRRewriter rewriter(&getContext());
-  IRMapping mapper;
-  SmallVector<spatial::SpatCompute> computes(funcOp.getOps<spatial::SpatCompute>());
-  if (!computes.empty())
-    return;
-  auto returnOp = llvm::cast<func::ReturnOp>(funcOp.getRegion().front().getTerminator());
-  rewriter.setInsertionPoint(returnOp);
-
-  SmallVector<Type> sourceTypes;
-  SmallVector<Location> sourceLoc;
-  for (auto source : funcOp.getArguments()) {
-    sourceTypes.push_back(source.getType());
-    sourceLoc.push_back(source.getLoc());
-  }
-
-  auto newCompute = spatial::SpatCompute::create(
-    rewriter, returnOp.getLoc(), returnOp.getOperandTypes(), funcOp.getArguments(), {}, {});
-  auto BB = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), sourceTypes, sourceLoc);
-  for (auto [bbArg, computeArg] : llvm::zip(BB->getArguments(), newCompute.getOperands()))
-    mapper.map(computeArg, bbArg);
-  newCompute.getProperties().setOperandSegmentSizes({(int) 0, (int) sourceTypes.size()});
-  rewriter.setInsertionPointToEnd(BB);
-  for (Operation& inst : funcOp.getOps())
-    if (!isa<spatial::SpatCompute, func::ReturnOp>(&inst))
-      rewriter.clone(inst, mapper);
-
-  auto yield = spatial::SpatYieldOp::create(rewriter, funcOp.getLoc(), returnOp.getOperands());
-  for (size_t i = 0; i < yield.getNumOperands(); ++i)
-    yield.setOperand(i, mapper.lookupOrDefault(yield.getOperand(i)));
-
-  for (Operation& inst : llvm::make_early_inc_range(funcOp.getOps()))
-    if (!isa<spatial::SpatCompute, func::ReturnOp>(&inst)){
-      inst.dropAllUses();
-      rewriter.eraseOp(&inst);
-  }
-
-  for (auto [index, computeResult] : llvm::enumerate(newCompute.getResults()))
-    returnOp.setOperand(index, computeResult);
-}
-
 std::unique_ptr<Pass> createONNXToSpatialPass() { return std::make_unique<ONNXToSpatialPass>(); }
 
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Elementwise.cpp

@@ -5,9 +5,9 @@
 
 #include "llvm/ADT/SmallVector.h"
 
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/Gemm.cpp

@@ -11,6 +11,7 @@
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostFoldability.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -49,6 +50,45 @@ materializeScaledConstantTensor(Value value, float factor, ConversionPatternRewr
   return arith::ConstantOp::create(rewriter, loc, denseAttr.getType(), scaledAttr).getResult();
 }
 
+static Value transposeForSpatial(Value value,
+                                 RankedTensorType resultType,
+                                 ArrayRef<int64_t> permutation,
+                                 ConversionPatternRewriter& rewriter,
+                                 Location loc) {
+  if (isHostFoldableValue(value))
+    return ONNXTransposeOp::create(rewriter, loc, resultType, value, rewriter.getI64ArrayAttr(permutation));
+
+  auto computeOp = createSpatCompute<1>(rewriter, loc, TypeRange {resultType}, {}, value, [&](Value input) {
+    Value transposed = ONNXTransposeOp::create(rewriter, loc, resultType, input, rewriter.getI64ArrayAttr(permutation));
+    spatial::SpatYieldOp::create(rewriter, loc, transposed);
+  });
+  return computeOp.getResult(0);
+}
+
+static Value
+expandRankOneBias(Value value, RankedTensorType resultType, ConversionPatternRewriter& rewriter, Location loc) {
+  if (isHostFoldableValue(value))
+    return tensor::ExpandShapeOp::create(rewriter,
+                                         loc,
+                                         resultType,
+                                         value,
+                                         SmallVector<ReassociationIndices> {
+                                           {0, 1}
+    });
+
+  auto computeOp = createSpatCompute<1>(rewriter, loc, TypeRange {resultType}, {}, value, [&](Value input) {
+    Value expanded = tensor::ExpandShapeOp::create(rewriter,
+                                                   loc,
+                                                   resultType,
+                                                   input,
+                                                   SmallVector<ReassociationIndices> {
+                                                     {0, 1}
+    });
+    spatial::SpatYieldOp::create(rewriter, loc, expanded);
+  });
+  return computeOp.getResult(0);
+}
+
 struct GemmToManyGemv : OpConversionPattern<ONNXGemmOp> {
   using OpConversionPattern::OpConversionPattern;
 
@@ -81,6 +121,11 @@ static SmallVector<Value> materializeBatchRowSlices(Value matrix,
   auto rowType = RankedTensorType::get({1, matrixType.getDimSize(1)}, matrixType.getElementType());
   SmallVector<Type> resultTypes(static_cast<size_t>(numRows), rowType);
 
+  if (isHostFoldableValue(matrix)) {
+    auto extractRowsOp = spatial::SpatExtractRowsOp::create(rewriter, loc, TypeRange(resultTypes), matrix);
+    return SmallVector<Value>(extractRowsOp->result_begin(), extractRowsOp->result_end());
+  }
+
   auto buildRowSlices = [&](Value matrixArg) {
     auto extractRowsOp = spatial::SpatExtractRowsOp::create(rewriter, loc, TypeRange(resultTypes), matrixArg);
     return SmallVector<Value>(extractRowsOp->result_begin(), extractRowsOp->result_end());
@@ -122,7 +167,8 @@ static SmallVector<Value> materializeBatchRowSlices(Value matrix,
     rootValue = definingOp->getOperand(0);
   }
 
-  return buildRowSlices(matrix);
+  SmallVector<Operation*> reversedChainOps(chainOps.rbegin(), chainOps.rend());
+  return cloneBatchInputChainIntoSliceCompute(rootValue, reversedChainOps, rootValue);
 }
 
 } // namespace
@@ -175,13 +221,7 @@ LogicalResult GemmToManyGemv::matchAndRewrite(ONNXGemmOp gemmOp,
     // Expand rank-1 bias [N] to rank-2 [1, N] for uniform handling
     if (cType.getRank() == 1) {
       auto expandedType = RankedTensorType::get({1, cType.getDimSize(0)}, cType.getElementType());
-      c = tensor::ExpandShapeOp::create(rewriter,
-                                        loc,
-                                        expandedType,
-                                        c,
-                                        SmallVector<ReassociationIndices> {
-                                          {0, 1}
-      });
+      c = expandRankOneBias(c, expandedType, rewriter, loc);
       cType = expandedType;
     }
     if (!cType.hasStaticShape()) {
@@ -196,25 +236,18 @@ LogicalResult GemmToManyGemv::matchAndRewrite(ONNXGemmOp gemmOp,
   }
 
   auto outRowType = RankedTensorType::get({1, outType.getDimSize(1)}, outType.getElementType());
+  SmallVector<Value> aSlices = materializeBatchRowSlices(a, aType, rewriter, loc);
+  SmallVector<Value> cSlices;
+  if (hasC && cHasNumOutRows)
+    cSlices = materializeBatchRowSlices(c, cType, rewriter, loc);
 
   SmallVector<Value> gemvOps;
-  gemvOps.reserve(numOutRows);
+  gemvOps.reserve(static_cast<size_t>(numOutRows));
   for (int64_t rowIdx = 0; rowIdx < numOutRows; rowIdx++) {
-    SmallVector<OpFoldResult> offsets = {rewriter.getIndexAttr(rowIdx), rewriter.getIndexAttr(0)};
-    SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(aType.getDimSize(1))};
-    SmallVector<OpFoldResult> strides = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-    auto aSliceType = RankedTensorType::get({1, aType.getDimSize(1)}, aType.getElementType());
-    auto aSlice = tensor::ExtractSliceOp::create(rewriter, loc, aSliceType, a, offsets, sizes, strides).getResult();
-
     Value cSlice = c;
     if (hasC) {
-      if (cHasNumOutRows) {
-        SmallVector<OpFoldResult> offsets = {rewriter.getIndexAttr(rowIdx), rewriter.getIndexAttr(0)};
-        SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(cType.getDimSize(1))};
-        SmallVector<OpFoldResult> strides = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-        auto cSliceType = RankedTensorType::get({1, cType.getDimSize(1)}, cType.getElementType());
-        cSlice = tensor::ExtractSliceOp::create(rewriter, loc, cSliceType, c, offsets, sizes, strides).getResult();
-      }
+      if (cHasNumOutRows)
+        cSlice = cSlices[static_cast<size_t>(rowIdx)];
       else if (!isVectorShape(getTensorShape(c))) {
         gemmOp.emitOpError("requires Gemm bias C to be vector-like when shared across decomposed rows");
         return failure();
@@ -224,7 +257,7 @@ LogicalResult GemmToManyGemv::matchAndRewrite(ONNXGemmOp gemmOp,
     auto gemvOp = ONNXGemmOp::create(rewriter,
                                      loc,
                                      outRowType,
-                                     aSlice,
+                                     aSlices[static_cast<size_t>(rowIdx)],
                                      b,
                                      cSlice,
                                      rewriter.getF32FloatAttr(1.0f),
@@ -267,13 +300,7 @@ LogicalResult GemvToSpatialCompute::matchAndRewrite(ONNXGemmOp gemmOp,
     // Expand rank-1 bias [N] to rank-2 [1, N] for uniform handling
     if (cType.getRank() == 1) {
       auto expandedType = RankedTensorType::get({1, cType.getDimSize(0)}, cType.getElementType());
-      c = tensor::ExpandShapeOp::create(rewriter,
-                                        gemmLoc,
-                                        expandedType,
-                                        c,
-                                        SmallVector<ReassociationIndices> {
-                                          {0, 1}
-      });
+      c = expandRankOneBias(c, expandedType, rewriter, gemmLoc);
       cType = expandedType;
     }
     if (!cType.hasStaticShape()) {
@@ -305,13 +332,14 @@ LogicalResult GemvToSpatialCompute::matchAndRewrite(ONNXGemmOp gemmOp,
 
   if (transA) {
     auto aShape = aType.getShape();
-    auto transposedType = aType.cloneWith(ArrayRef({aShape[1], aShape[0]}), aType.getElementType());
-    a = ONNXTransposeOp::create(rewriter, gemmLoc, transposedType, a, rewriter.getI64ArrayAttr({1, 0}));
+    auto transposedType = RankedTensorType::get({aShape[1], aShape[0]}, aType.getElementType());
+    a = transposeForSpatial(a, transposedType, {1, 0}, rewriter, gemmLoc);
+    aType = cast<RankedTensorType>(a.getType());
   }
   if (transB) {
     auto bShape = bType.getShape();
-    auto transposedType = bType.cloneWith(ArrayRef({bShape[1], bShape[0]}), bType.getElementType());
-    b = ONNXTransposeOp::create(rewriter, gemmLoc, transposedType, b, rewriter.getI64ArrayAttr({1, 0}));
+    auto transposedType = RankedTensorType::get({bShape[1], bShape[0]}, bType.getElementType());
+    b = transposeForSpatial(b, transposedType, {1, 0}, rewriter, gemmLoc);
     bType = cast<RankedTensorType>(b.getType());
   }
 
@@ -335,7 +363,6 @@ LogicalResult GemvToSpatialCompute::matchAndRewrite(ONNXGemmOp gemmOp,
   auto [aNumHSlices, aLastHSliceSize] = ceilIntegerDivideWithRemainder(aType.getDimSize(1), crossbarSize.getValue());
   auto [bNumHSlices, bLastHSliceSize] = ceilIntegerDivideWithRemainder(bType.getDimSize(1), crossbarSize.getValue());
   auto bNumVSlices = aNumHSlices;
-  auto bLastVSliceSize = aLastHSliceSize;
   auto cNumHSlices = bNumHSlices;
   auto cLastHSliceSize = bLastHSliceSize;
   auto outNumHSlices = cNumHSlices;
@@ -469,12 +496,15 @@ LogicalResult GemmToSpatialComputeBatch::matchAndRewrite(ONNXGemmOp gemmOp,
 
   if (gemmOpAdaptor.getTransB()) {
     auto bShape = bType.getShape();
-    auto transposedType = bType.cloneWith(ArrayRef({bShape[1], bShape[0]}), bType.getElementType());
-    b = ONNXTransposeOp::create(rewriter, loc, transposedType, b, rewriter.getI64ArrayAttr({1, 0}));
+    auto transposedType = RankedTensorType::get({bShape[1], bShape[0]}, bType.getElementType());
+    b = transposeForSpatial(b, transposedType, {1, 0}, rewriter, loc);
     bType = cast<RankedTensorType>(b.getType());
   }
   (void) bType;
 
+  if (!isHostFoldableValue(b))
+    return failure();
+
   Value sharedBias;
   if (hasC) {
     auto scaledC = materializeScaledConstantTensor(c, gemmOpAdaptor.getBeta().convertToFloat(), rewriter, loc);
@@ -484,13 +514,7 @@ LogicalResult GemmToSpatialComputeBatch::matchAndRewrite(ONNXGemmOp gemmOp,
     auto cType = cast<RankedTensorType>(c.getType());
     if (cType.getRank() == 1) {
       auto expandedType = RankedTensorType::get({1, cType.getDimSize(0)}, cType.getElementType());
-      c = tensor::ExpandShapeOp::create(rewriter,
-                                        loc,
-                                        expandedType,
-                                        c,
-                                        SmallVector<ReassociationIndices> {
-                                          {0, 1}
-      });
+      c = expandRankOneBias(c, expandedType, rewriter, loc);
       cType = cast<RankedTensorType>(c.getType());
     }
     if (!cType.hasStaticShape()) {

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/MatMul.cpp

@@ -2,11 +2,11 @@
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/PatternMatch.h"
 
-#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostFoldability.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -36,49 +36,27 @@ static Value extractBatchMatrix(Value value,
   SmallVector<OpFoldResult> sizes = {
     rewriter.getIndexAttr(1), rewriter.getIndexAttr(rows), rewriter.getIndexAttr(cols)};
   SmallVector<OpFoldResult> strides = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-  Value slice = tensor::ExtractSliceOp::create(rewriter, loc, sliceType, value, offsets, sizes, strides);
-
   auto matrixType = RankedTensorType::get({rows, cols}, type.getElementType());
-  return tensor::CollapseShapeOp::create(rewriter,
-                                         loc,
-                                         matrixType,
-                                         slice,
-                                         SmallVector<ReassociationIndices> {
-                                           {0, 1},
-                                           {2}
-  });
-}
+  auto buildMatrix = [&](Value input) -> Value {
+    Value slice = tensor::ExtractSliceOp::create(rewriter, loc, sliceType, input, offsets, sizes, strides);
+    return tensor::CollapseShapeOp::create(rewriter,
+                                           loc,
+                                           matrixType,
+                                           slice,
+                                           SmallVector<ReassociationIndices> {
+                                             {0, 1},
+                                             {2}
+    });
+  };
 
-static bool isConstantLikeOperand(Value value) {
-  llvm::SmallPtrSet<Operation*, 8> visited;
+  if (isHostFoldableValue(value))
+    return buildMatrix(value);
 
-  while (auto* definingOp = value.getDefiningOp()) {
-    if (!visited.insert(definingOp).second)
-      return false;
-    if (definingOp->hasTrait<OpTrait::ConstantLike>())
-      return true;
-
-    if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(definingOp)) {
-      value = extractSliceOp.getSource();
-      continue;
-    }
-    if (auto expandShapeOp = dyn_cast<tensor::ExpandShapeOp>(definingOp)) {
-      value = expandShapeOp.getSrc();
-      continue;
-    }
-    if (auto collapseShapeOp = dyn_cast<tensor::CollapseShapeOp>(definingOp)) {
-      value = collapseShapeOp.getSrc();
-      continue;
-    }
-    if (auto transposeOp = dyn_cast<ONNXTransposeOp>(definingOp)) {
-      value = transposeOp.getData();
-      continue;
-    }
-
-    return false;
-  }
-
-  return false;
+  auto batchMatrixCompute =
+    createSpatCompute<1>(rewriter, loc, TypeRange {matrixType}, {}, ValueRange {value}, [&](Value input) {
+      spatial::SpatYieldOp::create(rewriter, loc, buildMatrix(input));
+    });
+  return batchMatrixCompute.getResult(0);
 }
 
 static Value transposeLastTwoDims(Value value, PatternRewriter& rewriter, Location loc) {
@@ -107,15 +85,31 @@ static Value transposeLastTwoDimsInCompute(Value value, PatternRewriter& rewrite
     perm = {0, 2, 1};
   }
 
-  auto transposeCompute =
-    createSpatCompute<1>(rewriter, loc, transposedType, {}, ValueRange {value}, [&](Value input) {
-      Value transposed =
-        ONNXTransposeOp::create(rewriter, loc, transposedType, input, rewriter.getI64ArrayAttr(perm));
-      spatial::SpatYieldOp::create(rewriter, loc, transposed);
-    });
+  auto transposeCompute = createSpatCompute<1>(rewriter, loc, transposedType, {}, ValueRange {value}, [&](Value input) {
+    Value transposed = ONNXTransposeOp::create(rewriter, loc, transposedType, input, rewriter.getI64ArrayAttr(perm));
+    spatial::SpatYieldOp::create(rewriter, loc, transposed);
+  });
   return transposeCompute.getResult(0);
 }
 
+static Value concatValues(ValueRange inputs, int64_t axis, PatternRewriter& rewriter, Location loc) {
+  auto firstType = cast<RankedTensorType>(inputs.front().getType());
+  SmallVector<int64_t> outputShape(firstType.getShape().begin(), firstType.getShape().end());
+  int64_t concatDimSize = 0;
+  for (Value input : inputs)
+    concatDimSize += cast<RankedTensorType>(input.getType()).getDimSize(axis);
+  outputShape[axis] = concatDimSize;
+  auto resultType = RankedTensorType::get(outputShape, firstType.getElementType(), firstType.getEncoding());
+
+  if (llvm::all_of(inputs, isHostFoldableValue))
+    return createSpatConcat(rewriter, loc, axis, inputs);
+
+  auto concatCompute = createSpatCompute(rewriter, loc, TypeRange {resultType}, {}, inputs, [&](ValueRange args) {
+    spatial::SpatYieldOp::create(rewriter, loc, createSpatConcat(rewriter, loc, axis, args));
+  });
+  return concatCompute.getResult(0);
+}
+
 struct MatMulToGemm : OpRewritePattern<ONNXMatMulOp> {
   using OpRewritePattern::OpRewritePattern;
 
@@ -157,7 +151,7 @@ struct MatMulToGemm : OpRewritePattern<ONNXMatMulOp> {
     }
 
     Location loc = matmulOp.getLoc();
-    bool useTransposedForm = isConstantLikeOperand(matmulOp.getA()) && !isConstantLikeOperand(matmulOp.getB());
+    bool useTransposedForm = isHostFoldableValue(matmulOp.getA()) && !isHostFoldableValue(matmulOp.getB());
 
     Value lhs = matmulOp.getA();
     Value rhs = matmulOp.getB();
@@ -193,8 +187,14 @@ struct MatMulToGemm : OpRewritePattern<ONNXMatMulOp> {
                                             rewriter.getBoolAttr(false),
                                             rewriter.getBoolAttr(false))
                            .getY();
-      if (useTransposedForm)
-        gemmResult = ONNXTransposeOp::create(rewriter, loc, outType, gemmResult, rewriter.getI64ArrayAttr({1, 0}));
+      if (useTransposedForm) {
+        auto transposeCompute =
+          createSpatCompute<1>(rewriter, loc, TypeRange {outType}, {}, gemmResult, [&](Value input) {
+            Value transposed = ONNXTransposeOp::create(rewriter, loc, outType, input, rewriter.getI64ArrayAttr({1, 0}));
+            spatial::SpatYieldOp::create(rewriter, loc, transposed);
+          });
+        gemmResult = transposeCompute.getResult(0);
+      }
       rewriter.replaceOp(matmulOp, gemmResult);
       return success();
     }
@@ -215,24 +215,30 @@ struct MatMulToGemm : OpRewritePattern<ONNXMatMulOp> {
                                             rewriter.getBoolAttr(false),
                                             rewriter.getBoolAttr(false))
                            .getY();
-      if (useTransposedForm)
-        gemmResult = ONNXTransposeOp::create(
-          rewriter,
-          loc,
-          RankedTensorType::get({m, n}, outType.getElementType()),
-          gemmResult,
-          rewriter.getI64ArrayAttr({1, 0}));
-      batchResults.push_back(tensor::ExpandShapeOp::create(rewriter,
-                                                           loc,
-                                                           batchedOutType,
-                                                           gemmResult,
-                                                           SmallVector<ReassociationIndices> {
-                                                             {0, 1},
-                                                             {2}
-      }));
+      auto batchResultCompute =
+        createSpatCompute<1>(rewriter, loc, TypeRange {batchedOutType}, {}, gemmResult, [&](Value input) {
+          Value resultMatrix = input;
+          if (useTransposedForm) {
+            resultMatrix = ONNXTransposeOp::create(rewriter,
+                                                   loc,
+                                                   RankedTensorType::get({m, n}, outType.getElementType()),
+                                                   input,
+                                                   rewriter.getI64ArrayAttr({1, 0}));
+          }
+          Value expanded = tensor::ExpandShapeOp::create(rewriter,
+                                                         loc,
+                                                         batchedOutType,
+                                                         resultMatrix,
+                                                         SmallVector<ReassociationIndices> {
+                                                           {0, 1},
+                                                           {2}
+          });
+          spatial::SpatYieldOp::create(rewriter, loc, expanded);
+        });
+      batchResults.push_back(batchResultCompute.getResult(0));
     }
 
-    Value result = createSpatConcat(rewriter, loc, /*axis=*/0, batchResults);
+    Value result = concatValues(batchResults, /*axis=*/0, rewriter, loc);
     rewriter.replaceOp(matmulOp, result);
     return success();
   }

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ReduceMean.cpp

@@ -6,7 +6,8 @@
 #include <algorithm>
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostFoldability.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -81,6 +82,24 @@ createAverageCompute(Value input, RankedTensorType resultType, ConversionPattern
   return computeOp.getResult(0);
 }
 
+static Value concatValues(ValueRange inputs, int64_t axis, ConversionPatternRewriter& rewriter, Location loc) {
+  auto firstType = cast<RankedTensorType>(inputs.front().getType());
+  SmallVector<int64_t> outputShape(firstType.getShape().begin(), firstType.getShape().end());
+  int64_t concatDimSize = 0;
+  for (Value input : inputs)
+    concatDimSize += cast<RankedTensorType>(input.getType()).getDimSize(axis);
+  outputShape[axis] = concatDimSize;
+  auto resultType = RankedTensorType::get(outputShape, firstType.getElementType(), firstType.getEncoding());
+
+  if (llvm::all_of(inputs, isHostFoldableValue))
+    return createSpatConcat(rewriter, loc, axis, inputs);
+
+  auto concatCompute = createSpatCompute(rewriter, loc, TypeRange {resultType}, {}, inputs, [&](ValueRange args) {
+    spatial::SpatYieldOp::create(rewriter, loc, createSpatConcat(rewriter, loc, axis, args));
+  });
+  return concatCompute.getResult(0);
+}
+
 static Value buildReduceMeanKeepdims(Value input,
                                      ArrayRef<bool> reducedAxes,
                                      int64_t axis,
@@ -100,7 +119,7 @@ static Value buildReduceMeanKeepdims(Value input,
   for (Value slice : slices)
     reducedSlices.push_back(buildReduceMeanKeepdims(slice, reducedAxes, axis + 1, leafType, rewriter, loc));
 
-  return createSpatConcat(rewriter, loc, axis, reducedSlices);
+  return concatValues(reducedSlices, axis, rewriter, loc);
 }
 
 static Value squeezeReducedAxes(Value keepdimsValue,
@@ -115,9 +134,16 @@ static Value squeezeReducedAxes(Value keepdimsValue,
     return tensor::FromElementsOp::create(rewriter, loc, resultType, ValueRange {element});
   }
 
-  return tensor::CollapseShapeOp::create(
-           rewriter, loc, resultType, keepdimsValue, buildCollapseReassociation(reducedAxes))
-    .getResult();
+  auto reassociation = buildCollapseReassociation(reducedAxes);
+  if (isHostFoldableValue(keepdimsValue))
+    return tensor::CollapseShapeOp::create(rewriter, loc, resultType, keepdimsValue, reassociation).getResult();
+
+  auto squeezeCompute =
+    createSpatCompute<1>(rewriter, loc, TypeRange {resultType}, {}, ValueRange {keepdimsValue}, [&](Value input) {
+      Value collapsed = tensor::CollapseShapeOp::create(rewriter, loc, resultType, input, reassociation);
+      spatial::SpatYieldOp::create(rewriter, loc, collapsed);
+    });
+  return squeezeCompute.getResult(0);
 }
 
 struct ReduceMeanToSpatialCompute : OpConversionPattern<ONNXReduceMeanV13Op> {

src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Pool.cpp

@@ -31,8 +31,8 @@ static int64_t getOptionalI64(std::optional<ArrayAttrT> arrayAttr, size_t index,
 }
 
 template <typename PoolOp>
-static FailureOr<Value>
-concatAlongAxis(ConversionPatternRewriter& rewriter, Location loc, PoolOp poolOp, int64_t axis, ArrayRef<Value> values) {
+static FailureOr<Value> concatAlongAxis(
+  ConversionPatternRewriter& rewriter, Location loc, PoolOp poolOp, int64_t axis, ArrayRef<Value> values) {
   if (values.empty()) {
     poolOp.emitOpError("failed to build pooled output because an intermediate concatenation input list was empty");
     return failure();
@@ -68,8 +68,8 @@ reduceWindowValues(ConversionPatternRewriter& rewriter, Location loc, Operation*
   return reduced;
 }
 
-static FailureOr<Value>
-scaleAverageWindow(ConversionPatternRewriter& rewriter, Location loc, Operation* op, Value reducedWindow, int64_t divisor) {
+static FailureOr<Value> scaleAverageWindow(
+  ConversionPatternRewriter& rewriter, Location loc, Operation* op, Value reducedWindow, int64_t divisor) {
   if (divisor <= 0) {
     op->emitOpError("AveragePool divisor must be positive");
     return failure();

src/PIM/Conversion/ONNXToSpatial/Patterns/NN/Softmax.cpp

@@ -2,7 +2,8 @@
 #include "mlir/Transforms/DialectConversion.h"
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostFoldability.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -32,6 +33,24 @@ static Value createSoftmaxCompute(Value input, ConversionPatternRewriter& rewrit
   return computeOp.getResult(0);
 }
 
+static Value concatValues(ValueRange inputs, int64_t axis, ConversionPatternRewriter& rewriter, Location loc) {
+  auto firstType = cast<RankedTensorType>(inputs.front().getType());
+  SmallVector<int64_t> outputShape(firstType.getShape().begin(), firstType.getShape().end());
+  int64_t concatDimSize = 0;
+  for (Value input : inputs)
+    concatDimSize += cast<RankedTensorType>(input.getType()).getDimSize(axis);
+  outputShape[axis] = concatDimSize;
+  auto resultType = RankedTensorType::get(outputShape, firstType.getElementType(), firstType.getEncoding());
+
+  if (llvm::all_of(inputs, isHostFoldableValue))
+    return createSpatConcat(rewriter, loc, axis, inputs);
+
+  auto concatCompute = createSpatCompute(rewriter, loc, TypeRange {resultType}, {}, inputs, [&](ValueRange args) {
+    spatial::SpatYieldOp::create(rewriter, loc, createSpatConcat(rewriter, loc, axis, args));
+  });
+  return concatCompute.getResult(0);
+}
+
 static Value
 buildSoftmax(Value input, int64_t softmaxAxis, int64_t axis, ConversionPatternRewriter& rewriter, Location loc) {
   auto inputType = cast<RankedTensorType>(input.getType());
@@ -47,7 +66,7 @@ buildSoftmax(Value input, int64_t softmaxAxis, int64_t axis, ConversionPatternRe
   for (Value slice : slices)
     rebuiltSlices.push_back(buildSoftmax(slice, softmaxAxis, axis + 1, rewriter, loc));
 
-  return createSpatConcat(rewriter, loc, axis, rebuiltSlices);
+  return concatValues(rebuiltSlices, axis, rewriter, loc);
 }
 
 struct SoftmaxToSpatialCompute : OpConversionPattern<ONNXSoftmaxOp> {
@@ -92,8 +111,13 @@ struct SoftmaxToSpatialCompute : OpConversionPattern<ONNXSoftmaxOp> {
       Value transposedInput = preTransposeCompute.getResult(0);
       Value transposedResult = buildSoftmax(
         transposedInput, /*softmaxAxis=*/inputType.getRank() - 1, /*axis=*/0, rewriter, softmaxOp.getLoc());
-      result = ONNXTransposeOp::create(
-        rewriter, softmaxOp.getLoc(), inputType, transposedResult, rewriter.getI64ArrayAttr(inversePermutation));
+      auto postTransposeCompute =
+        createSpatCompute<1>(rewriter, softmaxOp.getLoc(), TypeRange {inputType}, {}, transposedResult, [&](Value x) {
+          Value transposed = ONNXTransposeOp::create(
+            rewriter, softmaxOp.getLoc(), inputType, x, rewriter.getI64ArrayAttr(inversePermutation));
+          spatial::SpatYieldOp::create(rewriter, softmaxOp.getLoc(), transposed);
+        });
+      result = postTransposeCompute.getResult(0);
     }
 
     rewriter.replaceOp(softmaxOp, result);

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Concat.cpp

@@ -2,6 +2,8 @@
 #include "mlir/IR/PatternMatch.h"
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostFoldability.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -18,7 +20,17 @@ struct Concat : public OpConversionPattern<ONNXConcatOp> {
     auto inputs = adaptor.getInputs();
     int64_t axis = adaptor.getAxis();
 
-    rewriter.replaceOp(maxpoolOp, createSpatConcat(rewriter, maxpoolOp.getLoc(), axis, inputs));
+    if (llvm::all_of(inputs, isHostFoldableValue)) {
+      rewriter.replaceOp(maxpoolOp, createSpatConcat(rewriter, maxpoolOp.getLoc(), axis, inputs));
+      return success();
+    }
+
+    auto computeOp = createSpatCompute(
+      rewriter, maxpoolOp.getLoc(), TypeRange {maxpoolOp.getResult().getType()}, {}, inputs, [&](ValueRange args) {
+        spatial::SpatYieldOp::create(
+          rewriter, maxpoolOp.getLoc(), createSpatConcat(rewriter, maxpoolOp.getLoc(), axis, args));
+      });
+    rewriter.replaceOp(maxpoolOp, computeOp.getResults());
 
     return success();
   }

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Gather.cpp

@@ -6,7 +6,7 @@
 #include "llvm/ADT/SmallVector.h"
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Reshape.cpp

@@ -3,7 +3,10 @@
 
 #include "llvm/ADT/SmallVector.h"
 
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostFoldability.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
 using namespace mlir;
@@ -95,18 +98,33 @@ struct Reshape : OpConversionPattern<ONNXReshapeOp> {
       return success();
     }
 
+    auto replaceWithReshape = [&](auto buildReshape) -> LogicalResult {
+      if (isHostFoldableValue(adaptor.getData())) {
+        rewriter.replaceOp(reshapeOp, buildReshape(adaptor.getData()));
+        return success();
+      }
+
+      auto computeOp = createSpatCompute<1>(
+        rewriter, reshapeOp.getLoc(), TypeRange {resultType}, {}, adaptor.getData(), [&](Value data) {
+          Value reshaped = buildReshape(data);
+          spatial::SpatYieldOp::create(rewriter, reshapeOp.getLoc(), reshaped);
+        });
+      rewriter.replaceOp(reshapeOp, computeOp.getResults());
+      return success();
+    };
+
     SmallVector<ReassociationIndices> reassociation;
     if (sourceType.getRank() > resultType.getRank()
-        && inferCollapseReassociation(sourceType.getShape(), resultType.getShape(), reassociation)) {
-      rewriter.replaceOpWithNewOp<tensor::CollapseShapeOp>(reshapeOp, resultType, adaptor.getData(), reassociation);
-      return success();
-    }
+        && inferCollapseReassociation(sourceType.getShape(), resultType.getShape(), reassociation))
+      return replaceWithReshape([&](Value data) {
+        return tensor::CollapseShapeOp::create(rewriter, reshapeOp.getLoc(), resultType, data, reassociation);
+      });
 
     if (sourceType.getRank() < resultType.getRank()
-        && inferExpandReassociation(sourceType.getShape(), resultType.getShape(), reassociation)) {
-      rewriter.replaceOpWithNewOp<tensor::ExpandShapeOp>(reshapeOp, resultType, adaptor.getData(), reassociation);
-      return success();
-    }
+        && inferExpandReassociation(sourceType.getShape(), resultType.getShape(), reassociation))
+      return replaceWithReshape([&](Value data) {
+        return tensor::ExpandShapeOp::create(rewriter, reshapeOp.getLoc(), resultType, data, reassociation);
+      });
 
     return failure();
   }

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Resize.cpp

@@ -6,7 +6,7 @@
 #include <algorithm>
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Split.cpp

@@ -2,7 +2,9 @@
 #include "mlir/Transforms/DialectConversion.h"
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/HostFoldability.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
 using namespace mlir;
@@ -47,16 +49,40 @@ struct Split : OpConversionPattern<ONNXSplitOp> {
     outputs.reserve(splitOp.getNumResults());
 
     int64_t offset = 0;
+    SmallVector<RankedTensorType> resultTypes;
+    resultTypes.reserve(splitOp.getNumResults());
+    SmallVector<int64_t> sliceSizes;
+    sliceSizes.reserve(splitOp.getNumResults());
     for (Value result : splitOp.getResults()) {
       auto resultType = dyn_cast<RankedTensorType>(result.getType());
       if (!resultType || !resultType.hasStaticShape())
         return failure();
-      int64_t sliceSize = resultType.getShape()[axis];
-      outputs.push_back(extractSliceAt(adaptor.getInput(), axis, offset, sliceSize, rewriter, splitOp.getLoc()));
-      offset += sliceSize;
+      resultTypes.push_back(resultType);
+      sliceSizes.push_back(resultType.getShape()[axis]);
     }
 
-    rewriter.replaceOp(splitOp, outputs);
+    if (isHostFoldableValue(adaptor.getInput())) {
+      for (int64_t sliceSize : sliceSizes) {
+        outputs.push_back(extractSliceAt(adaptor.getInput(), axis, offset, sliceSize, rewriter, splitOp.getLoc()));
+        offset += sliceSize;
+      }
+      rewriter.replaceOp(splitOp, outputs);
+      return success();
+    }
+
+    auto computeOp = createSpatCompute<1>(
+      rewriter, splitOp.getLoc(), TypeRange(splitOp.getResultTypes()), {}, adaptor.getInput(), [&](Value input) {
+        SmallVector<Value> runtimeOutputs;
+        runtimeOutputs.reserve(resultTypes.size());
+        int64_t runtimeOffset = 0;
+        for (int64_t sliceSize : sliceSizes) {
+          runtimeOutputs.push_back(extractSliceAt(input, axis, runtimeOffset, sliceSize, rewriter, splitOp.getLoc()));
+          runtimeOffset += sliceSize;
+        }
+        spatial::SpatYieldOp::create(rewriter, splitOp.getLoc(), runtimeOutputs);
+      });
+
+    rewriter.replaceOp(splitOp, computeOp.getResults());
     return success();
   }
 };

src/PIM/Conversion/ONNXToSpatial/PostPatterns.cpp

@@ -0,0 +1,265 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/PatternMatch.h"
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+
+#include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/WeightMaterialization.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PostPatterns.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+namespace {
+
+static bool isWeightMaterializationHelperUser(Operation* op) {
+  return isa<tensor::ExtractSliceOp, tensor::ExpandShapeOp, tensor::CollapseShapeOp, ONNXTransposeOp>(op);
+}
+
+static bool canPromoteInputBlockArgument(BlockArgument arg) {
+  return !arg.use_empty() && llvm::all_of(arg.getUsers(), isWeightMaterializationHelperUser);
+}
+
+static bool isDirectConstantValue(Value value) {
+  return isa_and_nonnull<arith::ConstantOp, ONNXConstantOp>(value.getDefiningOp());
+}
+
+// Collapses one-lane batches so later phases do not carry batch-only structure unnecessarily.
+struct FoldSingleLaneComputeBatchPattern : OpRewritePattern<spatial::SpatComputeBatch> {
+  using OpRewritePattern<spatial::SpatComputeBatch>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatComputeBatch batchOp, PatternRewriter& rewriter) const override {
+    if (batchOp.getLaneCount() != 1)
+      return rewriter.notifyMatchFailure(batchOp, "requires a single lane");
+
+    auto loc = batchOp.getLoc();
+    rewriter.setInsertionPoint(batchOp);
+    auto computeOp =
+      spatial::SpatCompute::create(rewriter, loc, batchOp.getResultTypes(), batchOp.getWeights(), batchOp.getInputs());
+    computeOp.getProperties().setOperandSegmentSizes(
+      {static_cast<int>(batchOp.getWeights().size()), static_cast<int>(batchOp.getInputs().size())});
+
+    Block& templateBlock = batchOp.getBody().front();
+    SmallVector<Type> blockArgTypes;
+    SmallVector<Location> blockArgLocs;
+    blockArgTypes.reserve(templateBlock.getNumArguments());
+    blockArgLocs.reserve(templateBlock.getNumArguments());
+    for (BlockArgument arg : templateBlock.getArguments()) {
+      blockArgTypes.push_back(arg.getType());
+      blockArgLocs.push_back(loc);
+    }
+
+    auto* newBlock =
+      rewriter.createBlock(&computeOp.getBody(), computeOp.getBody().end(), TypeRange(blockArgTypes), blockArgLocs);
+    IRMapping mapper;
+    for (auto [oldArg, newArg] : llvm::zip(templateBlock.getArguments(), newBlock->getArguments()))
+      mapper.map(oldArg, newArg);
+
+    rewriter.setInsertionPointToEnd(newBlock);
+    for (Operation& op : templateBlock)
+      rewriter.clone(op, mapper);
+
+    batchOp->replaceAllUsesWith(computeOp->getResults());
+    rewriter.eraseOp(batchOp);
+    return success();
+  }
+};
+
+// Promotes foldable helper chains from runtime inputs to weights to avoid artificial compute inputs.
+struct PromoteWeightLikeComputeInputsPattern : OpRewritePattern<spatial::SpatCompute> {
+  using OpRewritePattern<spatial::SpatCompute>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatCompute compute, PatternRewriter& rewriter) const override {
+    SmallVector<bool> promoteInput(compute.getInputs().size(), false);
+    bool needsRewrite = false;
+    Block& oldBlock = compute.getBody().front();
+    for (auto [inputIdx, input] : llvm::enumerate(compute.getInputs())) {
+      if (inputIdx >= oldBlock.getNumArguments())
+        continue;
+      if (!isWeightLikeComputeOperand(input))
+        continue;
+      if (isDirectConstantValue(input) && !canPromoteInputBlockArgument(oldBlock.getArgument(inputIdx)))
+        continue;
+      promoteInput[inputIdx] = true;
+      needsRewrite = true;
+    }
+    if (!needsRewrite)
+      return rewriter.notifyMatchFailure(compute, "no weight-like inputs to promote");
+
+    rewriter.setInsertionPointAfter(compute);
+
+    SmallVector<Value> newWeights(compute.getWeights().begin(), compute.getWeights().end());
+    SmallVector<Value> newInputs;
+    SmallVector<Type> newInputTypes;
+    SmallVector<Location> newInputLocs;
+    newWeights.reserve(compute.getWeights().size() + compute.getInputs().size());
+    newInputs.reserve(compute.getInputs().size());
+    newInputTypes.reserve(compute.getInputs().size());
+    newInputLocs.reserve(compute.getInputs().size());
+
+    for (auto [inputIdx, input] : llvm::enumerate(compute.getInputs())) {
+      if (promoteInput[inputIdx]) {
+        newWeights.push_back(input);
+        continue;
+      }
+      newInputs.push_back(input);
+      newInputTypes.push_back(input.getType());
+      newInputLocs.push_back(input.getLoc());
+    }
+
+    auto newCompute =
+      spatial::SpatCompute::create(rewriter, compute.getLoc(), compute.getResultTypes(), newWeights, newInputs);
+    auto* newBlock =
+      rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), newInputTypes, newInputLocs);
+    newCompute.getProperties().setOperandSegmentSizes(
+      {static_cast<int>(newWeights.size()), static_cast<int>(newInputs.size())});
+    rewriter.setInsertionPointToStart(newBlock);
+
+    IRRewriter bodyRewriter(rewriter.getContext());
+    bodyRewriter.setInsertionPointToStart(newBlock);
+
+    IRMapping mapper;
+    size_t newInputIdx = 0;
+    for (auto [oldInputIdx, oldArg] : llvm::enumerate(oldBlock.getArguments())) {
+      if (!promoteInput[oldInputIdx]) {
+        mapper.map(oldArg, newBlock->getArgument(newInputIdx++));
+        continue;
+      }
+
+      auto clonedValue = materializeWeightLikeValueInBlock(compute.getInputs()[oldInputIdx], bodyRewriter, mapper);
+      if (failed(clonedValue))
+        return rewriter.notifyMatchFailure(compute, "failed to materialize promoted weight-like operand");
+      mapper.map(oldArg, *clonedValue);
+    }
+
+    for (Operation& op : oldBlock.without_terminator())
+      rewriter.clone(op, mapper);
+
+    auto oldYield = cast<spatial::SpatYieldOp>(oldBlock.getTerminator());
+    SmallVector<Value> newYieldOperands;
+    newYieldOperands.reserve(oldYield.getOutputs().size());
+    for (Value operand : oldYield.getOutputs()) {
+      auto mapped = mapper.lookupOrNull(operand);
+      newYieldOperands.push_back(mapped ? cast<Value>(mapped) : operand);
+    }
+    spatial::SpatYieldOp::create(rewriter, oldYield.getLoc(), newYieldOperands);
+
+    rewriter.replaceOp(compute, newCompute.getResults());
+    return success();
+  }
+};
+
+// Promotes foldable batch helper chains to weights while preserving compact compute_batch IR.
+struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::SpatComputeBatch> {
+  using OpRewritePattern<spatial::SpatComputeBatch>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatComputeBatch compute, PatternRewriter& rewriter) const override {
+    SmallVector<bool> promoteInput(compute.getInputs().size(), false);
+    bool needsRewrite = false;
+    Block& oldBlock = compute.getBody().front();
+    for (auto [inputIdx, input] : llvm::enumerate(compute.getInputs())) {
+      if (inputIdx >= oldBlock.getNumArguments())
+        continue;
+      if (!isWeightLikeComputeOperand(input))
+        continue;
+      if (isDirectConstantValue(input) && !canPromoteInputBlockArgument(oldBlock.getArgument(inputIdx)))
+        continue;
+      promoteInput[inputIdx] = true;
+      needsRewrite = true;
+    }
+    if (!needsRewrite)
+      return rewriter.notifyMatchFailure(compute, "no weight-like batch inputs to promote");
+
+    rewriter.setInsertionPointAfter(compute);
+
+    SmallVector<Value> newWeights(compute.getWeights().begin(), compute.getWeights().end());
+    SmallVector<Value> newInputs;
+    SmallVector<Type> newInputTypes;
+    SmallVector<Location> newInputLocs;
+    newWeights.reserve(compute.getWeights().size() + compute.getInputs().size());
+    newInputs.reserve(compute.getInputs().size());
+    newInputTypes.reserve(compute.getInputs().size());
+    newInputLocs.reserve(compute.getInputs().size());
+
+    for (auto [inputIdx, input] : llvm::enumerate(compute.getInputs())) {
+      if (promoteInput[inputIdx]) {
+        newWeights.push_back(input);
+        continue;
+      }
+      newInputs.push_back(input);
+      newInputTypes.push_back(input.getType());
+      newInputLocs.push_back(input.getLoc());
+    }
+
+    auto newCompute =
+      spatial::SpatComputeBatch::create(rewriter,
+                                        compute.getLoc(),
+                                        compute.getResultTypes(),
+                                        rewriter.getI32IntegerAttr(static_cast<int32_t>(compute.getLaneCount())),
+                                        newWeights,
+                                        newInputs);
+    auto* newBlock =
+      rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), newInputTypes, newInputLocs);
+    newCompute.getProperties().setOperandSegmentSizes(
+      {static_cast<int>(newWeights.size()), static_cast<int>(newInputs.size())});
+    rewriter.setInsertionPointToStart(newBlock);
+
+    IRRewriter bodyRewriter(rewriter.getContext());
+    bodyRewriter.setInsertionPointToStart(newBlock);
+
+    IRMapping mapper;
+    size_t newInputIdx = 0;
+    for (auto [oldInputIdx, oldArg] : llvm::enumerate(oldBlock.getArguments())) {
+      if (!promoteInput[oldInputIdx]) {
+        mapper.map(oldArg, newBlock->getArgument(newInputIdx++));
+        continue;
+      }
+
+      auto clonedValue = materializeWeightLikeValueInBlock(compute.getInputs()[oldInputIdx], bodyRewriter, mapper);
+      if (failed(clonedValue))
+        return rewriter.notifyMatchFailure(compute, "failed to materialize promoted batch weight-like operand");
+      mapper.map(oldArg, *clonedValue);
+    }
+
+    for (Operation& op : oldBlock.without_terminator())
+      rewriter.clone(op, mapper);
+
+    auto oldYield = cast<spatial::SpatYieldOp>(oldBlock.getTerminator());
+    SmallVector<Value> newYieldOperands;
+    newYieldOperands.reserve(oldYield.getOutputs().size());
+    for (Value operand : oldYield.getOutputs()) {
+      auto mapped = mapper.lookupOrNull(operand);
+      newYieldOperands.push_back(mapped ? cast<Value>(mapped) : operand);
+    }
+    spatial::SpatYieldOp::create(rewriter, oldYield.getLoc(), newYieldOperands);
+
+    rewriter.replaceOp(compute, newCompute.getResults());
+    return success();
+  }
+};
+
+} // namespace
+
+void populateEarlyPostPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<FoldSingleLaneComputeBatchPattern>(ctx);
+}
+
+void populatePostPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
+  patterns.add<PromoteWeightLikeComputeInputsPattern, PromoteWeightLikeComputeBatchInputsPattern>(ctx);
+}
+
+void annotateWeightsConstants(func::FuncOp funcOp) {
+  funcOp.walk([&](arith::ConstantOp constantOp) {
+    if (hasOnlySpatialMvmVmmWeightUses(constantOp.getResult()))
+      markWeightAlways(constantOp);
+  });
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/PostPatterns.hpp

@@ -0,0 +1,14 @@
+#pragma once
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/MLIRContext.h"
+
+namespace onnx_mlir {
+
+void populateEarlyPostPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+
+void populatePostPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+
+void annotateWeightsConstants(mlir::func::FuncOp funcOp);
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/PrePatterns.cpp

@@ -0,0 +1,25 @@
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ConversionPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PrePatterns.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+namespace {
+
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatial.hpp.inc"
+
+} // namespace
+
+void populatePrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx) {
+  patterns.add<onnxToArithConstant>(ctx);
+  patterns.add<convAddToConvWithBiasLeft>(ctx);
+  patterns.add<convAddToConvWithBiasRight>(ctx);
+  patterns.add<matMulAddToGemm>(ctx);
+  patterns.add<matMulToGemm>(ctx);
+  patterns.add<removeFlattenSameShape>(ctx);
+  populateMatMulRewritePatterns(patterns, ctx);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/PrePatterns.hpp

@@ -0,0 +1,10 @@
+#pragma once
+
+#include "mlir/IR/MLIRContext.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace onnx_mlir {
+
+void populatePrePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp

@@ -0,0 +1,224 @@
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/IRMapping.h"
+
+#include "Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/TensorPackingPatterns.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+
+using namespace mlir;
+using namespace onnx_mlir::pim;
+
+namespace onnx_mlir {
+namespace {
+
+static int32_t translateSpatialCoreIdToPimCoreId(size_t spatialCoreId) { return static_cast<int32_t>(spatialCoreId); }
+
+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());
+
+  SmallVector<int32_t> coreIds;
+  coreIds.reserve(static_cast<size_t>(computeBatchOp.getLaneCount()));
+  for (uint32_t lane = 0; lane < computeBatchOp.getLaneCount(); ++lane)
+    coreIds.push_back(static_cast<int32_t>(fallbackCoreId++));
+  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));
+
+  Value input = mapper.lookup(sendTensorBatchOp.getInput());
+  if (auto concatOp = input.getDefiningOp<tensor::ConcatOp>())
+    if (concatOp.getDim() == 0)
+      if (Value packedInput =
+            createPackedExtractSliceTensor(concatOp.getInputs(), rewriter, sendTensorBatchOp.getLoc()))
+        input = packedInput;
+
+  pim::PimSendTensorBatchOp::create(
+    rewriter, sendTensorBatchOp.getLoc(), input, rewriter.getDenseI32ArrayAttr(targetCoreIds));
+}
+
+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));
+
+  auto outputType = cast<ShapedType>(receiveTensorBatchOp.getOutput().getType());
+  auto outputBuffer = createEmptyTensorFromShaped(rewriter, receiveTensorBatchOp.getLoc(), outputType);
+  Value received = pim::PimReceiveTensorBatchOp::create(rewriter,
+                                                        receiveTensorBatchOp.getLoc(),
+                                                        outputBuffer.getType(),
+                                                        outputBuffer,
+                                                        rewriter.getDenseI32ArrayAttr(sourceCoreIds))
+                     .getOutput();
+  mapper.map(receiveTensorBatchOp.getOutput(), received);
+}
+
+} // 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");
+
+  SmallVector<int32_t> coreIds = getPimCoreIdsForBatchOp(computeBatchOp, state.nextCoreId);
+  SmallVector<Value> batchWeights(computeBatchOp.getWeights().begin(), computeBatchOp.getWeights().end());
+  SmallVector<Value> batchInputs;
+  if (!computeBatchOp.getInputs().empty())
+    batchInputs.append(computeBatchOp.getInputs().begin(), computeBatchOp.getInputs().end());
+
+  rewriter.setInsertionPointAfter(computeBatchOp);
+  auto coreBatchOp = pim::PimCoreBatchOp::create(rewriter,
+                                                 loc,
+                                                 rewriter.getI32IntegerAttr(computeBatchOp.getLaneCount()),
+                                                 ValueRange(batchWeights),
+                                                 ValueRange(batchInputs));
+  coreBatchOp.getProperties().setOperandSegmentSizes(
+    {static_cast<int>(batchWeights.size()), static_cast<int>(batchInputs.size())});
+  coreBatchOp->setAttr(onnx_mlir::kCoreIdsAttrName, rewriter.getDenseI32ArrayAttr(coreIds));
+
+  SmallVector<Type> blockArgTypes;
+  SmallVector<Location> blockArgLocs;
+  for (BlockArgument arg : oldBlock.getArguments()) {
+    blockArgTypes.push_back(arg.getType());
+    blockArgLocs.push_back(arg.getLoc());
+  }
+  Block* newBlock =
+    rewriter.createBlock(&coreBatchOp.getBody(), coreBatchOp.getBody().end(), TypeRange(blockArgTypes), blockArgLocs);
+
+  IRMapping mapper;
+  rewriter.setInsertionPointToStart(newBlock);
+  for (auto [oldArg, newArg] : llvm::zip(oldBlock.getArguments(), newBlock->getArguments())) {
+    auto newArgType = cast<ShapedType>(newArg.getType());
+    auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, newArgType);
+    auto copied = pim::PimMemCopyHostToDevBatchOp::create(rewriter,
+                                                          loc,
+                                                          outputBuffer.getType(),
+                                                          outputBuffer,
+                                                          newArg,
+                                                          rewriter.getI32IntegerAttr(0),
+                                                          rewriter.getI32IntegerAttr(0),
+                                                          getTensorSizeInBytesAttr(rewriter, newArg))
+                    .getOutput();
+    mapper.map(oldArg, copied);
+  }
+
+  auto materializeCapturedTensor = [&](Value capturedTensor) -> Value {
+    if (auto mapped = mapper.lookupOrNull(capturedTensor))
+      return mapped;
+
+    auto capturedType = cast<ShapedType>(capturedTensor.getType());
+    auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, capturedType);
+    auto copied = pim::PimMemCopyHostToDevBatchOp::create(rewriter,
+                                                          loc,
+                                                          outputBuffer.getType(),
+                                                          outputBuffer,
+                                                          capturedTensor,
+                                                          rewriter.getI32IntegerAttr(0),
+                                                          rewriter.getI32IntegerAttr(0),
+                                                          getTensorSizeInBytesAttr(rewriter, capturedTensor))
+                    .getOutput();
+    mapper.map(capturedTensor, copied);
+    return copied;
+  };
+
+  rewriter.setInsertionPointToEnd(newBlock);
+  for (Operation& op : oldBlock) {
+    if (isa<spatial::SpatYieldOp>(op))
+      continue;
+
+    if (auto sendBatchOp = dyn_cast<spatial::SpatChannelSendBatchOp>(op)) {
+      pim::PimSendBatchOp::create(rewriter,
+                                  loc,
+                                  mapper.lookup(sendBatchOp.getInput()),
+                                  getTensorSizeInBytesAttr(rewriter, mapper.lookup(sendBatchOp.getInput())),
+                                  sendBatchOp.getTargetCoreIdsAttr());
+      continue;
+    }
+
+    if (auto sendTensorBatchOp = dyn_cast<spatial::SpatChannelSendTensorBatchOp>(op)) {
+      lowerChannelSendTensorBatch(sendTensorBatchOp, mapper, rewriter);
+      continue;
+    }
+
+    if (auto receiveBatchOp = dyn_cast<spatial::SpatChannelReceiveBatchOp>(op)) {
+      auto outputType = cast<ShapedType>(receiveBatchOp.getOutput().getType());
+      auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, outputType);
+      auto received = pim::PimReceiveBatchOp::create(rewriter,
+                                                     loc,
+                                                     outputBuffer.getType(),
+                                                     outputBuffer,
+                                                     getTensorSizeInBytesAttr(rewriter, receiveBatchOp.getOutput()),
+                                                     receiveBatchOp.getSourceCoreIdsAttr())
+                        .getOutput();
+      mapper.map(receiveBatchOp.getOutput(), received);
+      continue;
+    }
+
+    if (auto receiveTensorBatchOp = dyn_cast<spatial::SpatChannelReceiveTensorBatchOp>(op)) {
+      lowerChannelReceiveTensorBatch(receiveTensorBatchOp, mapper, rewriter);
+      continue;
+    }
+
+    if (auto toTensorOp = dyn_cast<bufferization::ToTensorOp>(op)) {
+      if (isa_and_present<memref::GetGlobalOp>(toTensorOp.getBuffer().getDefiningOp())) {
+        Operation* cloned = rewriter.clone(op, mapper);
+        auto clonedTensor = cloned->getResult(0);
+        auto clonedType = cast<ShapedType>(clonedTensor.getType());
+        auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, clonedType);
+        auto copied = pim::PimMemCopyHostToDevBatchOp::create(rewriter,
+                                                              loc,
+                                                              outputBuffer.getType(),
+                                                              outputBuffer,
+                                                              clonedTensor,
+                                                              rewriter.getI32IntegerAttr(0),
+                                                              rewriter.getI32IntegerAttr(0),
+                                                              getTensorSizeInBytesAttr(rewriter, clonedTensor))
+                        .getOutput();
+        mapper.map(toTensorOp.getResult(), copied);
+        continue;
+      }
+    }
+
+    for (Value operand : op.getOperands()) {
+      if (!isa<TensorType>(operand.getType()) || mapper.contains(operand))
+        continue;
+
+      Operation* definingOp = operand.getDefiningOp();
+      if (definingOp && definingOp->getBlock() == &oldBlock)
+        continue;
+
+      materializeCapturedTensor(operand);
+    }
+
+    Operation* cloned = rewriter.clone(op, mapper);
+    for (auto [originalResult, clonedResult] : llvm::zip(op.getResults(), cloned->getResults()))
+      mapper.map(originalResult, clonedResult);
+  }
+
+  rewriter.setInsertionPointToEnd(newBlock);
+  PimHaltOp::create(rewriter, loc);
+  return success();
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.hpp

@@ -0,0 +1,10 @@
+#pragma once
+
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.hpp"
+
+namespace onnx_mlir {
+
+mlir::LogicalResult
+lowerComputeBatchOp(spatial::SpatComputeBatch computeBatchOp, CoreLoweringState& state, mlir::IRRewriter& rewriter);
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/CMakeLists.txt

@@ -4,8 +4,16 @@ add_public_tablegen_target(SpatialToPimIncGen)
 
 add_pim_library(OMSpatialToPim
   SpatialToPimPass.cpp
+  BatchCoreLoweringPatterns.cpp
+  ChannelLoweringPatterns.cpp
+  Cleanup.cpp
   Common.cpp
-  Patterns.cpp
+  ComputeLikeRegionUtils.cpp
+  CoreLoweringPatterns.cpp
+  GlobalTensorMaterialization.cpp
+  PhaseVerification.cpp
+  ReturnPathNormalization.cpp
+  TensorPackingPatterns.cpp
 
   EXCLUDE_FROM_OM_LIBS
 

src/PIM/Conversion/SpatialToPim/ChannelLoweringPatterns.cpp

@@ -0,0 +1,136 @@
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/ChannelLoweringPatterns.hpp"
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+static int32_t toPimCoreId(int32_t spatialCoreId) { return spatialCoreId; }
+
+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())));
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
+struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatChannelReceiveOp op, PatternRewriter& rewriter) const override {
+    if (op->use_empty()) {
+      rewriter.eraseOp(op);
+      return success();
+    }
+    auto outputType = cast<ShapedType>(op.getResult().getType());
+    Value outputBuffer =
+      tensor::EmptyOp::create(rewriter, op.getLoc(), outputType.getShape(), outputType.getElementType()).getResult();
+    Value received = pim::PimReceiveOp::create(rewriter,
+                                               op.getLoc(),
+                                               op.getResult().getType(),
+                                               outputBuffer,
+                                               getTensorSizeInBytesAttr(rewriter, op.getResult()),
+                                               rewriter.getI32IntegerAttr(toPimCoreId(op.getSourceCoreId())))
+                       .getOutput();
+    rewriter.replaceOp(op, received);
+    return success();
+  }
+};
+
+struct ChannelSendTensorLowering : OpRewritePattern<spatial::SpatChannelSendTensorOp> {
+  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));
+    rewriter.eraseOp(op);
+    return success();
+  }
+};
+
+struct ChannelReceiveTensorLowering : OpRewritePattern<spatial::SpatChannelReceiveTensorOp> {
+  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));
+    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))
+        .getOutput();
+    rewriter.replaceOp(op, received);
+    return success();
+  }
+};
+
+struct ExtractRowsLowering : OpRewritePattern<spatial::SpatExtractRowsOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatExtractRowsOp op, PatternRewriter& rewriter) const override {
+    auto inputType = cast<RankedTensorType>(op.getInput().getType());
+    SmallVector<Value> replacements;
+    replacements.reserve(op.getNumResults());
+    for (auto [rowIndex, output] : llvm::enumerate(op.getOutputs())) {
+      auto outputType = cast<RankedTensorType>(output.getType());
+      SmallVector<OpFoldResult> offsets = {
+        rewriter.getIndexAttr(static_cast<int64_t>(rowIndex) * outputType.getDimSize(0)), rewriter.getIndexAttr(0)};
+      SmallVector<OpFoldResult> sizes = {rewriter.getIndexAttr(outputType.getDimSize(0)),
+                                         rewriter.getIndexAttr(inputType.getDimSize(1))};
+      SmallVector<OpFoldResult> strides = {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+      replacements.push_back(
+        tensor::ExtractSliceOp::create(rewriter, op.getLoc(), outputType, op.getInput(), offsets, sizes, strides)
+          .getResult());
+    }
+    rewriter.replaceOp(op, replacements);
+    return success();
+  }
+};
+
+struct ConcatLowering : OpRewritePattern<spatial::SpatConcatOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatConcatOp op, PatternRewriter& rewriter) const override {
+    auto outputType = cast<ShapedType>(op.getOutput().getType());
+    Value outputBuffer =
+      tensor::EmptyOp::create(rewriter, op.getLoc(), outputType.getShape(), outputType.getElementType()).getResult();
+    Value concatenated =
+      pim::PimConcatOp::create(
+        rewriter, op.getLoc(), op.getOutput().getType(), op.getAxisAttr(), op.getInputs(), outputBuffer)
+        .getOutput();
+    rewriter.replaceOp(op, concatenated);
+    return success();
+  }
+};
+
+} // namespace
+
+void populateChannelLoweringPatterns(RewritePatternSet& patterns) {
+  patterns.add<ChannelSendLowering,
+               ChannelReceiveLowering,
+               ChannelSendTensorLowering,
+               ChannelReceiveTensorLowering,
+               ExtractRowsLowering,
+               ConcatLowering>(patterns.getContext());
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/ChannelLoweringPatterns.hpp

@@ -0,0 +1,9 @@
+#pragma once
+
+#include "mlir/IR/PatternMatch.h"
+
+namespace onnx_mlir {
+
+void populateChannelLoweringPatterns(mlir::RewritePatternSet& patterns);
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/Cleanup.cpp

@@ -0,0 +1,42 @@
+#include "llvm/ADT/STLExtras.h"
+
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/Cleanup.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+LogicalResult erasePendingOps(SmallVectorImpl<Operation*>& pendingOps, IRRewriter& rewriter) {
+  while (!pendingOps.empty()) {
+    bool erasedAnyOp = false;
+    for (auto it = pendingOps.begin(); it != pendingOps.end();) {
+      Operation* opToRemove = *it;
+      if (!opToRemove->use_empty()) {
+        ++it;
+        continue;
+      }
+
+      rewriter.eraseOp(opToRemove);
+      it = pendingOps.erase(it);
+      erasedAnyOp = true;
+    }
+
+    if (erasedAnyOp)
+      continue;
+
+    for (Operation* opToRemove : pendingOps) {
+      InFlightDiagnostic diag = opToRemove->emitError("pending Spatial-to-PIM cleanup could not erase operation");
+      diag << "; op has " << llvm::range_size(opToRemove->getUsers()) << " remaining user(s)";
+      for (Operation* user : opToRemove->getUsers()) {
+        bool userPendingRemoval = llvm::is_contained(pendingOps, user);
+        opToRemove->emitRemark() << "remaining user `" << user->getName() << "`"
+                                 << (userPendingRemoval ? " is also pending removal" : " is not pending removal");
+      }
+    }
+    return failure();
+  }
+
+  return success();
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/Cleanup.hpp

@@ -0,0 +1,11 @@
+#pragma once
+
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Support/LLVM.h"
+
+namespace onnx_mlir {
+
+mlir::LogicalResult erasePendingOps(llvm::SmallVectorImpl<mlir::Operation*>& pendingOps, mlir::IRRewriter& rewriter);
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/ComputeLikeRegionUtils.cpp

@@ -0,0 +1,44 @@
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/ComputeLikeRegionUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+std::optional<unsigned> getDirectComputeLikeInputIndex(Operation* owner, unsigned operandNumber) {
+  auto getInputIndex = [operandNumber](Operation* op, unsigned inputCount) -> std::optional<unsigned> {
+    if (inputCount == 0)
+      return std::nullopt;
+    unsigned inputBegin = op->getNumOperands() - inputCount;
+    if (operandNumber < inputBegin)
+      return std::nullopt;
+    return operandNumber - inputBegin;
+  };
+
+  if (auto compute = dyn_cast<spatial::SpatCompute>(owner))
+    return getInputIndex(owner, compute.getInputs().size());
+
+  if (auto computeBatch = dyn_cast<spatial::SpatComputeBatch>(owner))
+    return getInputIndex(owner, computeBatch.getInputs().size());
+
+  return std::nullopt;
+}
+
+void replaceAndEraseDirectComputeLikeInput(PatternRewriter& rewriter,
+                                           Operation* owner,
+                                           unsigned inputIndex,
+                                           Value replacement) {
+  Block& body = owner->getRegion(0).front();
+  BlockArgument bodyArgument = body.getArgument(inputIndex);
+
+  rewriter.startOpModification(owner);
+  bodyArgument.replaceAllUsesWith(replacement);
+  if (auto compute = dyn_cast<spatial::SpatCompute>(owner))
+    compute.getInputsMutable().erase(inputIndex);
+  else
+    cast<spatial::SpatComputeBatch>(owner).getInputsMutable().erase(inputIndex);
+  body.eraseArgument(inputIndex);
+  rewriter.finalizeOpModification(owner);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/ComputeLikeRegionUtils.hpp

@@ -0,0 +1,17 @@
+#pragma once
+
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/PatternMatch.h"
+
+#include <optional>
+
+namespace onnx_mlir {
+
+std::optional<unsigned> getDirectComputeLikeInputIndex(mlir::Operation* owner, unsigned operandNumber);
+
+void replaceAndEraseDirectComputeLikeInput(mlir::PatternRewriter& rewriter,
+                                           mlir::Operation* owner,
+                                           unsigned inputIndex,
+                                           mlir::Value replacement);
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp

@@ -0,0 +1,213 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/IR/IRMapping.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/CoreLoweringPatterns.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+using namespace onnx_mlir::pim;
+
+namespace onnx_mlir {
+namespace {
+
+static bool isChannelUseChainOp(Operation* op) {
+  return isa<tensor::ExtractSliceOp,
+             tensor::CollapseShapeOp,
+             tensor::ExpandShapeOp,
+             tensor::CastOp,
+             tosa::ReshapeOp,
+             ONNXTransposeOp,
+             pim::PimTransposeOp>(op);
+}
+
+static void cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewriter& rewriter) {
+  for (Value operand : op->getOperands()) {
+    if (mapping.lookupOrNull(operand))
+      continue;
+
+    Operation* definingOp = operand.getDefiningOp();
+    if (!definingOp)
+      continue;
+
+    if (!isa<tensor::EmptyOp, arith::ConstantOp>(definingOp))
+      continue;
+
+    Operation* clonedOp = rewriter.clone(*definingOp, mapping);
+    for (auto [originalResult, newResult] : llvm::zip(definingOp->getResults(), clonedOp->getResults()))
+      mapping.map(originalResult, newResult);
+    rewriter.setInsertionPointAfter(clonedOp);
+  }
+}
+
+static int32_t translateSpatialCoreIdToPimCoreId(size_t spatialCoreId) { return static_cast<int32_t>(spatialCoreId); }
+
+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());
+  return static_cast<int32_t>(fallbackCoreId++);
+}
+
+static LogicalResult collectHelperComputeChain(spatial::SpatCompute computeOp,
+                                               SmallVectorImpl<Operation*>& helperChain,
+                                               bool requireReturnUse = true) {
+  if (computeOp.getInputs().size() != 1 || computeOp.getNumResults() != 1)
+    return failure();
+  if (requireReturnUse
+      && (!computeOp.getResult(0).hasOneUse() || !isa<func::ReturnOp>(*computeOp.getResult(0).getUsers().begin())))
+    return failure();
+
+  Block& block = computeOp.getBody().front();
+  if (block.getNumArguments() != 1)
+    return failure();
+
+  auto yieldOp = dyn_cast<spatial::SpatYieldOp>(block.getTerminator());
+  if (!yieldOp || yieldOp.getNumOperands() != 1)
+    return failure();
+
+  SmallVector<Operation*> reverseChain;
+  Value currentValue = yieldOp.getOperands().front();
+  Value blockArg = block.getArgument(0);
+
+  while (currentValue != blockArg) {
+    Operation* definingOp = currentValue.getDefiningOp();
+    if (!definingOp || definingOp->getBlock() != &block || !isChannelUseChainOp(definingOp))
+      return failure();
+    reverseChain.push_back(definingOp);
+    currentValue = definingOp->getOperand(0);
+  }
+
+  SmallPtrSet<Operation*, 8> chainSet(reverseChain.begin(), reverseChain.end());
+  for (Operation& op : llvm::make_early_inc_range(block.without_terminator()))
+    if (!chainSet.contains(&op) && !isa<tensor::EmptyOp, arith::ConstantOp>(op))
+      return failure();
+
+  helperChain.assign(reverseChain.rbegin(), reverseChain.rend());
+  return success();
+}
+
+static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatCompute computeOp, IRRewriter& rewriter) {
+  if (!computeOp.getInputs().empty() || computeOp.getNumResults() != 1)
+    return false;
+  if (!llvm::all_of(computeOp.getResult(0).getUsers(), [](Operation* user) {
+        return isa<spatial::SpatCompute, spatial::SpatComputeBatch, pim::PimCoreOp, pim::PimCoreBatchOp>(user);
+      }))
+    return false;
+
+  Block& block = computeOp.getBody().front();
+  if (block.getNumArguments() != 0)
+    return false;
+
+  auto yieldOp = dyn_cast<spatial::SpatYieldOp>(block.getTerminator());
+  if (!yieldOp || yieldOp.getNumOperands() != 1)
+    return false;
+
+  rewriter.setInsertionPoint(computeOp);
+  IRMapping mapping;
+  for (Operation& op : block.without_terminator()) {
+    cloneMappedHelperOperands(&op, mapping, rewriter);
+    Operation* clonedOp = rewriter.clone(op, mapping);
+    for (auto [originalResult, newResult] : llvm::zip(op.getResults(), clonedOp->getResults()))
+      mapping.map(originalResult, newResult);
+    rewriter.setInsertionPointAfter(clonedOp);
+  }
+
+  Value replacement = mapping.lookupOrDefault(yieldOp.getOperand(0));
+  computeOp.getResult(0).replaceAllUsesWith(replacement);
+  return true;
+}
+
+} // namespace
+
+void markOpToRemove(CoreLoweringState& state, Operation* op) {
+  if (!llvm::is_contained(state.operationsToRemove, op))
+    state.operationsToRemove.push_back(op);
+}
+
+LogicalResult lowerComputeOp(spatial::SpatCompute computeOp, CoreLoweringState& state, IRRewriter& rewriter) {
+  Location loc = computeOp->getLoc();
+
+  if (inlineInputlessHelperComputeForWeightLikeUsers(computeOp, rewriter))
+    return success();
+
+  SmallVector<Operation*> helperChain;
+  if (succeeded(collectHelperComputeChain(computeOp, helperChain)))
+    return success();
+
+  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())
+      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);
+  }
+
+  if (computeOp.getNumResults() != yieldOp.getNumOperands())
+    llvm_unreachable("ComputeOp must have same number of results as yieldOp operands");
+
+  for (auto [result, yieldValue] : llvm::zip(computeOp.getResults(), yieldOp.getOperands())) {
+    if (result.use_empty())
+      continue;
+
+    ReturnPathState returnPathState {state.outputTensors, state.operationsToRemove};
+    ReturnPathLoweringResult returnPathResult =
+      lowerComputeResultReturnPath(computeOp, cast<OpResult>(result), yieldValue, returnPathState, rewriter);
+    if (returnPathResult == ReturnPathLoweringResult::Failure)
+      return failure();
+    if (returnPathResult == ReturnPathLoweringResult::Handled)
+      continue;
+
+    auto resultUses = result.getUses();
+    if (rangeLength(resultUses) == 1) {
+      OpOperand& resultUse = *resultUses.begin();
+      Operation* resultUser = resultUse.getOwner();
+      if (isa<spatial::SpatChannelSendOp>(resultUser))
+        continue;
+    }
+
+    return computeOp.emitOpError("has an unsupported remaining result use during Spatial-to-PIM lowering");
+  }
+
+  rewriter.setInsertionPoint(yieldOp);
+  rewriter.replaceOpWithNewOp<PimHaltOp>(yieldOp);
+
+  SmallVector<Value> computeWeights;
+  if (!computeOp.getWeights().empty())
+    computeWeights.append(computeOp.getWeights().begin(), computeOp.getWeights().end());
+  rewriter.setInsertionPointAfter(computeOp);
+  auto coreOp = PimCoreOp::create(rewriter,
+                                  loc,
+                                  ValueRange(computeWeights),
+                                  rewriter.getI32IntegerAttr(getPimCoreIdForComputeOp(computeOp, state.nextCoreId)));
+  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());
+  coreOpBlocks.splice(coreOpBlocks.begin(), computeOp.getBody().getBlocks());
+  Block* tempComputeBlock = new Block();
+  computeOp.getBody().push_back(tempComputeBlock);
+  rewriter.setInsertionPointToEnd(tempComputeBlock);
+  PimHaltOp::create(rewriter, computeOp.getLoc());
+  return success();
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.hpp

@@ -0,0 +1,21 @@
+#pragma once
+
+#include "mlir/IR/PatternMatch.h"
+
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/ReturnPathNormalization.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+namespace onnx_mlir {
+
+struct CoreLoweringState {
+  size_t& nextCoreId;
+  llvm::SmallVectorImpl<OutputTensorFactory>& outputTensors;
+  llvm::SmallVectorImpl<mlir::Operation*>& operationsToRemove;
+};
+
+void markOpToRemove(CoreLoweringState& state, mlir::Operation* op);
+
+mlir::LogicalResult
+lowerComputeOp(spatial::SpatCompute computeOp, CoreLoweringState& state, mlir::IRRewriter& rewriter);
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/GlobalTensorMaterialization.cpp

@@ -6,16 +6,17 @@
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/SymbolTable.h"
 #include "mlir/IR/Value.h"
-#include "mlir/Support/LLVM.h"
 
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/Casting.h"
-#include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/LogicalResult.h"
 
 #include "Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/ComputeLikeRegionUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/GlobalTensorMaterialization.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
@@ -23,33 +24,33 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {
 
-
-static std::optional<unsigned> getDirectComputeInputIndex(Operation* owner, unsigned operandNumber) {
-  if (auto compute = dyn_cast<spatial::SpatCompute>(owner)) {
-    unsigned inputCount = compute.getInputs().size();
-    if (inputCount == 0)
-      return std::nullopt;
-
-    unsigned inputBegin = compute->getNumOperands() - inputCount;
-    if (operandNumber < inputBegin)
-      return std::nullopt;
-    return operandNumber - inputBegin;
-  }
-
-  if (auto computeBatch = dyn_cast<spatial::SpatComputeBatch>(owner)) {
-    unsigned inputCount = computeBatch.getInputs().size();
-    if (inputCount == 0)
-      return std::nullopt;
-
-    unsigned inputBegin = computeBatch->getNumOperands() - inputCount;
-    if (operandNumber < inputBegin)
-      return std::nullopt;
-    return operandNumber - inputBegin;
-  }
-
-  return std::nullopt;
+static std::string makeUniqueSymbolName(Operation* symbolTableOp, StringRef baseName) {
+  std::string name = baseName.str();
+  unsigned suffix = 0;
+  while (SymbolTable::lookupSymbolIn(symbolTableOp, name))
+    name = (baseName + "_" + Twine(suffix++)).str();
+  return name;
 }
 
+static memref::GlobalOp createPrivateMemrefGlobalWithUniqueName(PatternRewriter& rewriter,
+                                                                Location loc,
+                                                                ModuleOp moduleOp,
+                                                                StringRef baseName,
+                                                                MemRefType type,
+                                                                Attribute initialValue = {},
+                                                                UnitAttr constant = {}) {
+  std::string symbolName = makeUniqueSymbolName(moduleOp, baseName);
+  return memref::GlobalOp::create(rewriter,
+                                  loc,
+                                  rewriter.getStringAttr(symbolName),
+                                  rewriter.getStringAttr("private"),
+                                  TypeAttr::get(type),
+                                  initialValue,
+                                  constant,
+                                  IntegerAttr {});
+}
+
+// Sinks top-level tensor slices into compute regions so later lowering sees local runtime work.
 struct MoveExtractSliceIntoCompute final : OpRewritePattern<mlir::tensor::ExtractSliceOp> {
   using OpRewritePattern::OpRewritePattern;
 
@@ -59,7 +60,7 @@ struct MoveExtractSliceIntoCompute final : OpRewritePattern<mlir::tensor::Extrac
 
     for (auto& uses : extractSliceOp->getUses()) {
       if (isa<spatial::SpatCompute>(uses.getOwner())) {
-        if (!getDirectComputeInputIndex(uses.getOwner(), uses.getOperandNumber()))
+        if (!getDirectComputeLikeInputIndex(uses.getOwner(), uses.getOperandNumber()))
           return failure();
       }
       else if (isa_and_present<func::FuncOp>(uses.getOwner()->getParentOp())) {
@@ -72,7 +73,7 @@ struct MoveExtractSliceIntoCompute final : OpRewritePattern<mlir::tensor::Extrac
     for (auto& uses : llvm::make_early_inc_range(extractSliceOp->getUses())) {
 
       if (auto spatCompute = dyn_cast<spatial::SpatCompute>(uses.getOwner())) {
-        auto inputIndex = getDirectComputeInputIndex(spatCompute, uses.getOperandNumber());
+        auto inputIndex = getDirectComputeLikeInputIndex(spatCompute, uses.getOperandNumber());
         if (!inputIndex)
           return failure();
         auto BBArgIndex = *inputIndex;
@@ -87,14 +88,11 @@ struct MoveExtractSliceIntoCompute final : OpRewritePattern<mlir::tensor::Extrac
           mapSpatToExtract.insert({spatCompute.getOperation(), newExtractSlice->getResult(0)});
         }
 
-        rewriter.startOpModification(spatCompute.getOperation());
-        BBArgValue.replaceAllUsesWith(mapSpatToExtract[spatCompute.getOperation()]);
-        spatCompute.getInputsMutable().erase(BBArgIndex);
-        spatCompute.getBody().front().eraseArgument(BBArgIndex);
-        rewriter.finalizeOpModification(spatCompute.getOperation());
+        replaceAndEraseDirectComputeLikeInput(
+          rewriter, spatCompute.getOperation(), BBArgIndex, mapSpatToExtract[spatCompute.getOperation()]);
       }
       else if (auto spatComputeBatch = dyn_cast<spatial::SpatComputeBatch>(uses.getOwner())) {
-        auto inputIndex = getDirectComputeInputIndex(spatComputeBatch, uses.getOperandNumber());
+        auto inputIndex = getDirectComputeLikeInputIndex(spatComputeBatch, uses.getOperandNumber());
         if (!inputIndex)
           return failure();
         auto BBArgIndex = *inputIndex;
@@ -109,11 +107,8 @@ struct MoveExtractSliceIntoCompute final : OpRewritePattern<mlir::tensor::Extrac
           mapSpatToExtract.insert({spatComputeBatch.getOperation(), newExtractSlice->getResult(0)});
         }
 
-        rewriter.startOpModification(spatComputeBatch.getOperation());
-        BBArgValue.replaceAllUsesWith(mapSpatToExtract[spatComputeBatch.getOperation()]);
-        spatComputeBatch.getInputsMutable().erase(BBArgIndex);
-        spatComputeBatch.getBody().front().eraseArgument(BBArgIndex);
-        rewriter.finalizeOpModification(spatComputeBatch.getOperation());
+        replaceAndEraseDirectComputeLikeInput(
+          rewriter, spatComputeBatch.getOperation(), BBArgIndex, mapSpatToExtract[spatComputeBatch.getOperation()]);
       }
       else {
         {
@@ -148,11 +143,11 @@ struct MoveExtractSliceIntoCompute final : OpRewritePattern<mlir::tensor::Extrac
   }
 };
 
+// Turns runtime constants consumed by compute regions into private globals and local loads.
 struct ArithConstToGlobalMemoryPattern final : OpRewritePattern<mlir::arith::ConstantOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(mlir::arith::ConstantOp constantOp, PatternRewriter& rewriter) const override {
-    static int i = 0;
     Location loc = constantOp.getLoc();
 
     if (hasWeightAlways(constantOp))
@@ -177,15 +172,14 @@ struct ArithConstToGlobalMemoryPattern final : OpRewritePattern<mlir::arith::Con
     if (constRankedTensorType) {
       mlir::MemRefType memRefType =
         mlir::MemRefType::get(constRankedTensorType.getShape(), constRankedTensorType.getElementType());
-      std::string argName = "const_" + std::to_string(i++);
-      memref::GlobalOp::create(rewriter,
-                               loc,
-                               rewriter.getStringAttr(argName),
-                               rewriter.getStringAttr("private"),
-                               TypeAttr::get(memRefType),
-                               constantOp.getValueAttr(),
-                               rewriter.getUnitAttr(),
-                               {});
+      auto globalOp = createPrivateMemrefGlobalWithUniqueName(rewriter,
+                                                              loc,
+                                                              constantOp->getParentOfType<ModuleOp>(),
+                                                              "const",
+                                                              memRefType,
+                                                              constantOp.getValueAttr(),
+                                                              rewriter.getUnitAttr());
+      std::string argName = globalOp.getSymName().str();
 
       llvm::DenseMap<Operation*, Value> mapSpatComputeToConst;
 
@@ -193,11 +187,10 @@ struct ArithConstToGlobalMemoryPattern final : OpRewritePattern<mlir::arith::Con
         auto constUsers = constUses.getOwner();
 
         if (auto spatCompute = llvm::dyn_cast<spatial::SpatCompute>(constUsers)) {
-          auto inputIndex = getDirectComputeInputIndex(spatCompute, constUses.getOperandNumber());
+          auto inputIndex = getDirectComputeLikeInputIndex(spatCompute, constUses.getOperandNumber());
           if (!inputIndex)
             return failure();
           auto BBArgIndex = *inputIndex;
-          auto BBArgValue = spatCompute.getBody().front().getArgument(BBArgIndex);
           rewriter.setInsertionPoint(&spatCompute.getBody().front().front());
           if (!mapSpatComputeToConst.contains(spatCompute.getOperation())) {
             auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, argName);
@@ -206,18 +199,14 @@ struct ArithConstToGlobalMemoryPattern final : OpRewritePattern<mlir::arith::Con
             mapSpatComputeToConst.insert({spatCompute.getOperation(), toTensor.getResult()});
           }
 
-          rewriter.startOpModification(spatCompute.getOperation());
-          BBArgValue.replaceAllUsesWith(mapSpatComputeToConst[spatCompute.getOperation()]);
-          spatCompute.getInputsMutable().erase(BBArgIndex);
-          spatCompute.getBody().front().eraseArgument(BBArgIndex);
-          rewriter.finalizeOpModification(spatCompute.getOperation());
+          replaceAndEraseDirectComputeLikeInput(
+            rewriter, spatCompute.getOperation(), BBArgIndex, mapSpatComputeToConst[spatCompute.getOperation()]);
         }
         else if (auto spatComputeBatch = llvm::dyn_cast<spatial::SpatComputeBatch>(constUsers)) {
-          auto inputIndex = getDirectComputeInputIndex(spatComputeBatch, constUses.getOperandNumber());
+          auto inputIndex = getDirectComputeLikeInputIndex(spatComputeBatch, constUses.getOperandNumber());
           if (!inputIndex)
             return failure();
           auto BBArgIndex = *inputIndex;
-          auto BBArgValue = spatComputeBatch.getBody().front().getArgument(BBArgIndex);
           rewriter.setInsertionPoint(&spatComputeBatch.getBody().front().front());
           if (!mapSpatComputeToConst.contains(spatComputeBatch.getOperation())) {
             auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, argName);
@@ -226,11 +215,10 @@ struct ArithConstToGlobalMemoryPattern final : OpRewritePattern<mlir::arith::Con
             mapSpatComputeToConst.insert({spatComputeBatch.getOperation(), toTensor.getResult()});
           }
 
-          rewriter.startOpModification(spatComputeBatch.getOperation());
-          BBArgValue.replaceAllUsesWith(mapSpatComputeToConst[spatComputeBatch.getOperation()]);
-          spatComputeBatch.getInputsMutable().erase(BBArgIndex);
-          spatComputeBatch.getBody().front().eraseArgument(BBArgIndex);
-          rewriter.finalizeOpModification(spatComputeBatch.getOperation());
+          replaceAndEraseDirectComputeLikeInput(rewriter,
+                                                spatComputeBatch.getOperation(),
+                                                BBArgIndex,
+                                                mapSpatComputeToConst[spatComputeBatch.getOperation()]);
         }
         else {
           {
@@ -272,34 +260,26 @@ struct ArithConstToGlobalMemoryPattern final : OpRewritePattern<mlir::arith::Con
         auto constUsers = constUses.getOwner();
 
         if (auto spatCompute = llvm::dyn_cast<spatial::SpatCompute>(constUsers)) {
-          auto inputIndex = getDirectComputeInputIndex(spatCompute, constUses.getOperandNumber());
+          auto inputIndex = getDirectComputeLikeInputIndex(spatCompute, constUses.getOperandNumber());
           if (!inputIndex)
             return failure();
           auto BBArgIndex = *inputIndex;
-          auto BBArgValue = spatCompute.getBody().front().getArgument(BBArgIndex);
           rewriter.setInsertionPoint(&spatCompute.getBody().front().front());
           auto newConst = rewriter.clone(*constantOp);
 
-          rewriter.startOpModification(spatCompute.getOperation());
-          BBArgValue.replaceAllUsesWith(newConst->getResult(0));
-          spatCompute.getInputsMutable().erase(BBArgIndex);
-          spatCompute.getBody().front().eraseArgument(BBArgIndex);
-          rewriter.finalizeOpModification(spatCompute.getOperation());
+          replaceAndEraseDirectComputeLikeInput(
+            rewriter, spatCompute.getOperation(), BBArgIndex, newConst->getResult(0));
         }
         else if (auto spatComputeBatch = llvm::dyn_cast<spatial::SpatComputeBatch>(constUsers)) {
-          auto inputIndex = getDirectComputeInputIndex(spatComputeBatch, constUses.getOperandNumber());
+          auto inputIndex = getDirectComputeLikeInputIndex(spatComputeBatch, constUses.getOperandNumber());
           if (!inputIndex)
             return failure();
           auto BBArgIndex = *inputIndex;
-          auto BBArgValue = spatComputeBatch.getBody().front().getArgument(BBArgIndex);
           rewriter.setInsertionPoint(&spatComputeBatch.getBody().front().front());
           auto newConst = rewriter.clone(*constantOp);
 
-          rewriter.startOpModification(spatComputeBatch.getOperation());
-          BBArgValue.replaceAllUsesWith(newConst->getResult(0));
-          spatComputeBatch.getInputsMutable().erase(BBArgIndex);
-          spatComputeBatch.getBody().front().eraseArgument(BBArgIndex);
-          rewriter.finalizeOpModification(spatComputeBatch.getOperation());
+          replaceAndEraseDirectComputeLikeInput(
+            rewriter, spatComputeBatch.getOperation(), BBArgIndex, newConst->getResult(0));
         }
         else if (auto parent = constUsers->getParentOfType<spatial::SpatCompute>()) {
           if (!mapSpatComputeToConst.contains(parent)) {
@@ -321,11 +301,13 @@ struct ArithConstToGlobalMemoryPattern final : OpRewritePattern<mlir::arith::Con
         }
       }
     }
-    rewriter.eraseOp(constantOp);
+    if (constantOp->use_empty())
+      rewriter.eraseOp(constantOp);
     return success();
   }
 };
 
+// Materializes public function tensor inputs as globals so compute bodies can load them uniformly.
 struct FuncOpArgToGlobalMemoryPattern final : OpRewritePattern<mlir::func::FuncOp> {
   using OpRewritePattern::OpRewritePattern;
 
@@ -352,52 +334,36 @@ struct FuncOpArgToGlobalMemoryPattern final : OpRewritePattern<mlir::func::FuncO
       mlir::MemRefType memRefType =
         mlir::MemRefType::get(argRankedTensorType.getShape(), argRankedTensorType.getElementType());
 
-      std::string argName = "arg_" + std::to_string(index);
-
-      memref::GlobalOp::create(rewriter,
-                               loc,
-                               rewriter.getStringAttr(argName),
-                               rewriter.getStringAttr("private"),
-                               TypeAttr::get(memRefType),
-                               {},
-                               {},
-                               {});
+      std::string baseName = ("arg_" + Twine(index)).str();
+      auto globalOp = createPrivateMemrefGlobalWithUniqueName(
+        rewriter, loc, funcOp->getParentOfType<ModuleOp>(), baseName, memRefType);
+      std::string argName = globalOp.getSymName().str();
 
       for (auto& argUses : llvm::make_early_inc_range(arg.getUses())) {
         auto argUser = argUses.getOwner();
         if (auto spatCompute = dyn_cast<spatial::SpatCompute>(argUser)) {
-          auto inputIndex = getDirectComputeInputIndex(spatCompute, argUses.getOperandNumber());
+          auto inputIndex = getDirectComputeLikeInputIndex(spatCompute, argUses.getOperandNumber());
           if (!inputIndex)
             return failure();
           auto BBArgIndex = *inputIndex;
-          auto BBArgValue = spatCompute.getBody().front().getArgument(BBArgIndex);
           rewriter.setInsertionPoint(&spatCompute.getBody().front().front());
           auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, argName);
           auto toTensor = bufferization::ToTensorOp::create(
             rewriter, loc, argRankedTensorType, getGlobalOp, rewriter.getUnitAttr(), rewriter.getUnitAttr());
 
-          rewriter.startOpModification(spatCompute.getOperation());
-          BBArgValue.replaceAllUsesWith(toTensor);
-          spatCompute.getInputsMutable().erase(BBArgIndex);
-          spatCompute.getBody().front().eraseArgument(BBArgIndex);
-          rewriter.finalizeOpModification(spatCompute.getOperation());
+          replaceAndEraseDirectComputeLikeInput(rewriter, spatCompute.getOperation(), BBArgIndex, toTensor);
         }
         else if (auto spatComputeBatch = dyn_cast<spatial::SpatComputeBatch>(argUser)) {
-          auto inputIndex = getDirectComputeInputIndex(spatComputeBatch, argUses.getOperandNumber());
+          auto inputIndex = getDirectComputeLikeInputIndex(spatComputeBatch, argUses.getOperandNumber());
           if (!inputIndex)
             return failure();
           auto BBArgIndex = *inputIndex;
-          auto BBArgValue = spatComputeBatch.getBody().front().getArgument(BBArgIndex);
           rewriter.setInsertionPoint(&spatComputeBatch.getBody().front().front());
           auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, argName);
           auto toTensor = bufferization::ToTensorOp::create(
             rewriter, loc, argRankedTensorType, getGlobalOp, rewriter.getUnitAttr(), rewriter.getUnitAttr());
 
-          rewriter.startOpModification(spatComputeBatch.getOperation());
-          BBArgValue.replaceAllUsesWith(toTensor);
-          spatComputeBatch.getInputsMutable().erase(BBArgIndex);
-          spatComputeBatch.getBody().front().eraseArgument(BBArgIndex);
-          rewriter.finalizeOpModification(spatComputeBatch.getOperation());
+          replaceAndEraseDirectComputeLikeInput(rewriter, spatComputeBatch.getOperation(), BBArgIndex, toTensor);
         }
         else {
           rewriter.setInsertionPoint(argUser);
@@ -416,7 +382,7 @@ struct FuncOpArgToGlobalMemoryPattern final : OpRewritePattern<mlir::func::FuncO
 };
 
 } // namespace
-void populateGlobalTensorToMemrefPatterns(RewritePatternSet& patterns) {
+void populateGlobalTensorMaterializationPatterns(RewritePatternSet& patterns) {
   patterns.add<MoveExtractSliceIntoCompute, FuncOpArgToGlobalMemoryPattern, ArithConstToGlobalMemoryPattern>(
     patterns.getContext());
 }

src/PIM/Conversion/SpatialToPim/GlobalTensorMaterialization.hpp

@@ -0,0 +1,9 @@
+#pragma once
+
+#include "mlir/IR/PatternMatch.h"
+
+namespace onnx_mlir {
+
+void populateGlobalTensorMaterializationPatterns(mlir::RewritePatternSet& patterns);
+
+}

src/PIM/Conversion/SpatialToPim/Patterns.hpp

@@ -1,10 +0,0 @@
-#pragma once
-
-#include "mlir/IR/PatternMatch.h"
-
-
-namespace onnx_mlir {
-
-void populateGlobalTensorToMemrefPatterns(mlir::RewritePatternSet& patterns); 
-
-}

src/PIM/Conversion/SpatialToPim/PhaseVerification.cpp

@@ -0,0 +1,20 @@
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/PhaseVerification.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+LogicalResult verifySpatialToPimBoundary(ModuleOp moduleOp) {
+  bool hasFailure = false;
+  moduleOp.walk([&](Operation* op) {
+    if (op->getDialect()->getNamespace() != "spat")
+      return;
+
+    op->emitError("illegal Spatial operation remains after Spatial-to-PIM lowering");
+    hasFailure = true;
+  });
+  return success(!hasFailure);
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/PhaseVerification.hpp

@@ -0,0 +1,9 @@
+#pragma once
+
+#include "mlir/IR/BuiltinOps.h"
+
+namespace onnx_mlir {
+
+mlir::LogicalResult verifySpatialToPimBoundary(mlir::ModuleOp moduleOp);
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp

@@ -0,0 +1,587 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Bufferization/IR/Bufferization.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Tosa/IR/TosaOps.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/IR/SymbolTable.h"
+
+#include "Conversion/ONNXToSpatial/Common/Common.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"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+using namespace onnx_mlir::pim;
+
+namespace onnx_mlir {
+namespace {
+
+struct ReturnUseInfo {
+  size_t returnIndex;
+  SmallVector<Operation*> helperChain;
+};
+
+struct ConcatReturnUseInfo {
+  size_t returnIndex;
+  SmallVector<int64_t> sliceOffsets;
+  SmallVector<int64_t> concatShape;
+  SmallVector<Operation*> concatChain;
+  SmallVector<Operation*> helperChain;
+};
+
+static bool isReturnHelperChainOp(Operation* op) {
+  return isa<tensor::ExtractSliceOp,
+             tensor::CollapseShapeOp,
+             tensor::ExpandShapeOp,
+             tensor::CastOp,
+             tosa::ReshapeOp,
+             ONNXTransposeOp,
+             pim::PimTransposeOp>(op);
+}
+
+static void markOpToRemove(ReturnPathState& state, Operation* op) {
+  if (!llvm::is_contained(state.operationsToRemove, op))
+    state.operationsToRemove.push_back(op);
+}
+
+static std::string makeUniqueSymbolName(Operation* symbolTableOp, StringRef baseName) {
+  std::string name = baseName.str();
+  unsigned suffix = 0;
+  while (SymbolTable::lookupSymbolIn(symbolTableOp, name))
+    name = (baseName + "_" + Twine(suffix++)).str();
+  return name;
+}
+
+static int64_t computeFlatElementIndex(ArrayRef<int64_t> indices, ArrayRef<int64_t> shape) {
+  int64_t flatIndex = 0;
+  for (size_t i = 0; i < shape.size(); ++i) {
+    flatIndex *= shape[i];
+    flatIndex += indices[i];
+  }
+  return flatIndex;
+}
+
+static SmallVector<int64_t> expandFlatElementIndex(int64_t flatIndex, ArrayRef<int64_t> shape) {
+  SmallVector<int64_t> indices(shape.size(), 0);
+  for (int64_t dim = static_cast<int64_t>(shape.size()) - 1; dim >= 0; --dim) {
+    indices[dim] = flatIndex % shape[dim];
+    flatIndex /= shape[dim];
+  }
+  return indices;
+}
+
+static LogicalResult collectHelperComputeChain(spatial::SpatCompute computeOp,
+                                               SmallVectorImpl<Operation*>& helperChain) {
+  if (computeOp.getInputs().size() != 1 || computeOp.getNumResults() != 1)
+    return failure();
+  if (!computeOp.getResult(0).hasOneUse() || !isa<func::ReturnOp>(*computeOp.getResult(0).getUsers().begin()))
+    return failure();
+
+  Block& block = computeOp.getBody().front();
+  if (block.getNumArguments() != 1)
+    return failure();
+
+  auto yieldOp = dyn_cast<spatial::SpatYieldOp>(block.getTerminator());
+  if (!yieldOp || yieldOp.getNumOperands() != 1)
+    return failure();
+
+  SmallVector<Operation*> reverseChain;
+  Value currentValue = yieldOp.getOperands().front();
+  Value blockArg = block.getArgument(0);
+
+  while (currentValue != blockArg) {
+    Operation* definingOp = currentValue.getDefiningOp();
+    if (!definingOp || definingOp->getBlock() != &block || !isReturnHelperChainOp(definingOp))
+      return failure();
+    reverseChain.push_back(definingOp);
+    currentValue = definingOp->getOperand(0);
+  }
+
+  SmallPtrSet<Operation*, 8> chainSet(reverseChain.begin(), reverseChain.end());
+  for (Operation& op : llvm::make_early_inc_range(block.without_terminator()))
+    if (!chainSet.contains(&op) && !isa<tensor::EmptyOp, arith::ConstantOp>(op))
+      return failure();
+
+  helperChain.assign(reverseChain.rbegin(), reverseChain.rend());
+  return success();
+}
+
+static std::optional<ReturnUseInfo> analyzeReturnUse(Value value) {
+  auto uses = value.getUses();
+  if (rangeLength(uses) != 1)
+    return std::nullopt;
+
+  SmallVector<Operation*> helperChain;
+  Value currentValue = value;
+  Operation* currentUser = uses.begin()->getOwner();
+
+  while (isReturnHelperChainOp(currentUser)) {
+    helperChain.push_back(currentUser);
+    auto currentUses = currentUser->getResult(0).getUses();
+    if (rangeLength(currentUses) != 1)
+      return std::nullopt;
+    currentValue = currentUser->getResult(0);
+    currentUser = currentUses.begin()->getOwner();
+  }
+
+  if (!isa<func::ReturnOp>(currentUser))
+    return std::nullopt;
+
+  return ReturnUseInfo {
+    currentValue.getUses().begin()->getOperandNumber(),
+    std::move(helperChain),
+  };
+}
+
+static std::optional<ConcatReturnUseInfo> analyzeConcatReturnUse(Value value) {
+  auto getConcatResult = [](Operation* op) -> Value {
+    if (auto tensorConcat = dyn_cast<tensor::ConcatOp>(op))
+      return tensorConcat.getResult();
+    if (auto spatialConcat = dyn_cast<spatial::SpatConcatOp>(op))
+      return spatialConcat.getOutput();
+    if (auto pimConcat = dyn_cast<pim::PimConcatOp>(op))
+      return pimConcat.getOutput();
+    return {};
+  };
+  auto getConcatAxis = [](Operation* op) -> std::optional<int64_t> {
+    if (auto tensorConcat = dyn_cast<tensor::ConcatOp>(op))
+      return tensorConcat.getDim();
+    if (auto spatialConcat = dyn_cast<spatial::SpatConcatOp>(op))
+      return spatialConcat.getAxis();
+    if (auto pimConcat = dyn_cast<pim::PimConcatOp>(op))
+      return pimConcat.getAxis();
+    return std::nullopt;
+  };
+  auto getConcatOperands = [](Operation* op) -> OperandRange {
+    if (auto tensorConcat = dyn_cast<tensor::ConcatOp>(op))
+      return tensorConcat.getOperands();
+    if (auto spatialConcat = dyn_cast<spatial::SpatConcatOp>(op))
+      return spatialConcat.getInputs();
+    return cast<pim::PimConcatOp>(op).getInputs();
+  };
+
+  auto uses = value.getUses();
+  if (rangeLength(uses) != 1
+      || !isa<tensor::ConcatOp, spatial::SpatConcatOp, pim::PimConcatOp>(uses.begin()->getOwner()))
+    return std::nullopt;
+
+  auto valueType = dyn_cast<ShapedType>(value.getType());
+  if (!valueType || !valueType.hasStaticShape())
+    return std::nullopt;
+
+  SmallVector<int64_t> sliceOffsets(valueType.getRank(), 0);
+  SmallVector<int64_t> concatShape(valueType.getShape().begin(), valueType.getShape().end());
+  SmallVector<Operation*> concatChain;
+  Value currentValue = value;
+  Operation* currentUser = uses.begin()->getOwner();
+
+  while (isa<tensor::ConcatOp, spatial::SpatConcatOp, pim::PimConcatOp>(currentUser)) {
+    concatChain.push_back(currentUser);
+    size_t operandIndex = currentValue.getUses().begin()->getOperandNumber();
+    int64_t axis = *getConcatAxis(currentUser);
+    for (Value operand : getConcatOperands(currentUser).take_front(operandIndex))
+      sliceOffsets[axis] += cast<ShapedType>(operand.getType()).getShape()[axis];
+
+    Value concatResult = getConcatResult(currentUser);
+    auto concatType = dyn_cast<ShapedType>(concatResult.getType());
+    if (!concatType || !concatType.hasStaticShape())
+      return std::nullopt;
+    concatShape.assign(concatType.getShape().begin(), concatType.getShape().end());
+
+    currentValue = concatResult;
+    auto currentUses = currentValue.getUses();
+    if (rangeLength(currentUses) != 1)
+      return std::nullopt;
+    currentUser = currentUses.begin()->getOwner();
+  }
+
+  SmallVector<Operation*> helperChain;
+  if (auto helperCompute = dyn_cast<spatial::SpatCompute>(currentUser)) {
+    if (helperCompute.getInputs().size() != 1 || helperCompute.getInputs().front() != currentValue)
+      return std::nullopt;
+
+    if (failed(collectHelperComputeChain(helperCompute, helperChain)))
+      return std::nullopt;
+
+    currentValue = helperCompute.getResult(0);
+    auto currentUses = currentValue.getUses();
+    if (rangeLength(currentUses) != 1)
+      return std::nullopt;
+    currentUser = currentUses.begin()->getOwner();
+  }
+
+  while (isReturnHelperChainOp(currentUser)) {
+    helperChain.push_back(currentUser);
+    auto currentUses = currentUser->getResult(0).getUses();
+    if (rangeLength(currentUses) != 1)
+      return std::nullopt;
+    currentValue = currentUser->getResult(0);
+    currentUser = currentUses.begin()->getOwner();
+  }
+
+  if (!isa<func::ReturnOp>(currentUser))
+    return std::nullopt;
+
+  return ConcatReturnUseInfo {
+    currentValue.getUses().begin()->getOperandNumber(),
+    std::move(sliceOffsets),
+    std::move(concatShape),
+    std::move(concatChain),
+    std::move(helperChain),
+  };
+}
+
+static LogicalResult mapIndicesThroughHelperChain(ArrayRef<int64_t> sourceIndices,
+                                                  ArrayRef<int64_t> sourceShape,
+                                                  ArrayRef<Operation*> helperChain,
+                                                  SmallVectorImpl<int64_t>& mappedIndices) {
+  SmallVector<int64_t> currentIndices(sourceIndices.begin(), sourceIndices.end());
+  SmallVector<int64_t> currentShape(sourceShape.begin(), sourceShape.end());
+
+  auto reshapeToResultShape = [&](Operation* op) -> LogicalResult {
+    auto resultType = dyn_cast<ShapedType>(op->getResult(0).getType());
+    if (!resultType || !resultType.hasStaticShape())
+      return failure();
+    int64_t flatIndex = computeFlatElementIndex(currentIndices, currentShape);
+    currentShape.assign(resultType.getShape().begin(), resultType.getShape().end());
+    currentIndices = expandFlatElementIndex(flatIndex, currentShape);
+    return success();
+  };
+
+  for (Operation* op : helperChain) {
+    if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(op)) {
+      auto hasStaticValues = [](ArrayRef<int64_t> values) {
+        return llvm::all_of(values, [](int64_t value) { return !ShapedType::isDynamic(value); });
+      };
+      if (!hasStaticValues(extractSliceOp.getStaticOffsets()) || !hasStaticValues(extractSliceOp.getStaticSizes())
+          || !hasStaticValues(extractSliceOp.getStaticStrides()))
+        return failure();
+
+      SmallVector<int64_t> nextIndices;
+      nextIndices.reserve(currentIndices.size());
+      for (auto [index, offset, size, stride] : llvm::zip_equal(currentIndices,
+                                                                extractSliceOp.getStaticOffsets(),
+                                                                extractSliceOp.getStaticSizes(),
+                                                                extractSliceOp.getStaticStrides())) {
+        if (stride != 1 || index < offset || index >= offset + size)
+          return failure();
+        nextIndices.push_back(index - offset);
+      }
+
+      auto resultType = dyn_cast<ShapedType>(extractSliceOp.getResult().getType());
+      if (!resultType || !resultType.hasStaticShape())
+        return failure();
+      currentIndices = std::move(nextIndices);
+      currentShape.assign(resultType.getShape().begin(), resultType.getShape().end());
+      continue;
+    }
+
+    if (auto transposeOp = dyn_cast<ONNXTransposeOp>(op)) {
+      SmallVector<int64_t> nextIndices(currentIndices.size());
+      SmallVector<int64_t> nextShape(currentShape.size());
+      for (auto [destIndex, attr] : llvm::enumerate(transposeOp.getPermAttr().getAsRange<IntegerAttr>())) {
+        int64_t sourceIndex = attr.getInt();
+        nextIndices[destIndex] = currentIndices[sourceIndex];
+        nextShape[destIndex] = currentShape[sourceIndex];
+      }
+      currentIndices = std::move(nextIndices);
+      currentShape = std::move(nextShape);
+      continue;
+    }
+
+    if (auto transposeOp = dyn_cast<pim::PimTransposeOp>(op)) {
+      SmallVector<int64_t> nextIndices(currentIndices.size());
+      SmallVector<int64_t> nextShape(currentShape.size());
+      for (auto [destIndex, attr] : llvm::enumerate(transposeOp.getPermutation().getAsRange<IntegerAttr>())) {
+        int64_t sourceIndex = attr.getInt();
+        nextIndices[destIndex] = currentIndices[sourceIndex];
+        nextShape[destIndex] = currentShape[sourceIndex];
+      }
+      currentIndices = std::move(nextIndices);
+      currentShape = std::move(nextShape);
+      continue;
+    }
+
+    if (isa<tensor::CastOp, tosa::ReshapeOp, tensor::CollapseShapeOp, tensor::ExpandShapeOp>(op)) {
+      if (failed(reshapeToResultShape(op)))
+        return failure();
+      continue;
+    }
+
+    return failure();
+  }
+
+  mappedIndices.assign(currentIndices.begin(), currentIndices.end());
+  return success();
+}
+
+static void cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewriter& rewriter) {
+  for (Value operand : op->getOperands()) {
+    if (mapping.lookupOrNull(operand))
+      continue;
+
+    Operation* definingOp = operand.getDefiningOp();
+    if (!definingOp)
+      continue;
+
+    if (!isa<tensor::EmptyOp, arith::ConstantOp>(definingOp))
+      continue;
+
+    Operation* clonedOp = rewriter.clone(*definingOp, mapping);
+    for (auto [originalResult, newResult] : llvm::zip(definingOp->getResults(), clonedOp->getResults()))
+      mapping.map(originalResult, newResult);
+    rewriter.setInsertionPointAfter(clonedOp);
+  }
+}
+
+static void
+cloneHelperChain(Value sourceValue, ArrayRef<Operation*> helperChain, IRRewriter& rewriter, Value& clonedValue) {
+  IRMapping mapping;
+  mapping.map(sourceValue, sourceValue);
+  clonedValue = sourceValue;
+
+  rewriter.setInsertionPointAfterValue(sourceValue);
+  for (Operation* op : helperChain) {
+    cloneMappedHelperOperands(op, mapping, rewriter);
+    Operation* clonedOp = rewriter.clone(*op, mapping);
+    for (auto [originalResult, newResult] : llvm::zip(op->getResults(), clonedOp->getResults()))
+      mapping.map(originalResult, newResult);
+    clonedValue = clonedOp->getResult(0);
+    rewriter.setInsertionPointAfter(clonedOp);
+  }
+}
+
+static Value emitHostCopy(IRRewriter& rewriter,
+                          Location loc,
+                          Value outputTensor,
+                          Value sourceValue,
+                          int32_t hostTargetOffset,
+                          int32_t deviceSourceOffset,
+                          int32_t sizeInBytes) {
+  return PimMemCopyDevToHostOp::create(rewriter,
+                                       loc,
+                                       outputTensor.getType(),
+                                       outputTensor,
+                                       sourceValue,
+                                       rewriter.getI32IntegerAttr(hostTargetOffset),
+                                       rewriter.getI32IntegerAttr(deviceSourceOffset),
+                                       rewriter.getI32IntegerAttr(sizeInBytes))
+    .getOutput();
+}
+
+} // namespace
+
+void addReturnOutputBuffers(func::ReturnOp returnOp,
+                            IRRewriter& rewriter,
+                            SmallVectorImpl<OutputTensorFactory>& outputTensors) {
+  outputTensors.reserve(returnOp->getNumOperands());
+  for (auto [index, returnValue] : llvm::enumerate(returnOp->getOperands())) {
+    Value currentReturnValue = returnValue;
+    Operation* returnValueDefiningOp = currentReturnValue.getDefiningOp();
+    if (returnValueDefiningOp->hasTrait<OpTrait::ConstantLike>()) {
+      assert(!hasWeightAlways(returnValueDefiningOp));
+      outputTensors.push_back(
+        [currentReturnValue](IRRewriter& rewriter, Location loc) -> Value { return currentReturnValue; });
+    }
+    else {
+      auto outRankedTensorType = llvm::dyn_cast<RankedTensorType>(currentReturnValue.getType());
+      auto memRefType = MemRefType::get(outRankedTensorType.getShape(), outRankedTensorType.getElementType());
+
+      std::string outputBaseName = ("output_" + Twine(index)).str();
+      std::string outputName = makeUniqueSymbolName(returnOp->getParentOfType<ModuleOp>(), outputBaseName);
+      rewriter.setInsertionPoint(returnOp.getParentOp());
+      memref::GlobalOp::create(rewriter,
+                               returnOp.getLoc(),
+                               rewriter.getStringAttr(outputName),
+                               rewriter.getStringAttr("private"),
+                               TypeAttr::get(memRefType),
+                               {},
+                               {},
+                               {});
+      outputTensors.push_back([memRefType, outputName, outRankedTensorType](IRRewriter& rewriter, Location loc) {
+        auto getGlobalOp = memref::GetGlobalOp::create(rewriter, loc, memRefType, outputName);
+        auto toTensor = bufferization::ToTensorOp::create(
+          rewriter, loc, outRankedTensorType, getGlobalOp, rewriter.getUnitAttr(), rewriter.getUnitAttr());
+        return toTensor.getResult();
+      });
+    }
+  }
+}
+
+ReturnPathLoweringResult lowerComputeResultReturnPath(
+  spatial::SpatCompute computeOp, OpResult result, Value yieldValue, ReturnPathState& state, IRRewriter& rewriter) {
+  Location loc = computeOp->getLoc();
+  auto yieldType = cast<TensorType>(yieldValue.getType());
+
+  if (auto returnUse = analyzeReturnUse(result)) {
+    Value storedValue = yieldValue;
+    cloneHelperChain(yieldValue, returnUse->helperChain, rewriter, storedValue);
+    for (Operation* op : returnUse->helperChain)
+      markOpToRemove(state, op);
+
+    auto storedType = cast<ShapedType>(storedValue.getType());
+    size_t elementSize = storedType.getElementTypeBitWidth() / 8;
+    if (auto storedOp = storedValue.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));
+    return ReturnPathLoweringResult::Handled;
+  }
+
+  auto resultUses = result.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);
+      Value outputTensor = state.outputTensors[resultIndexInReturn](rewriter, loc);
+      emitHostCopy(
+        rewriter, loc, outputTensor, yieldValue, 0, 0, static_cast<int32_t>(yieldType.getNumElements() * elementSize));
+      return ReturnPathLoweringResult::Handled;
+    }
+  }
+
+  if (auto concatReturnUse = analyzeConcatReturnUse(result)) {
+    size_t elementSize = yieldType.getElementTypeBitWidth() / 8;
+    for (Operation* concatOp : concatReturnUse->concatChain)
+      markOpToRemove(state, concatOp);
+
+    if (concatReturnUse->helperChain.empty()) {
+      rewriter.setInsertionPointAfterValue(yieldValue);
+      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,
+                   static_cast<int32_t>(flatOffset * elementSize),
+                   0,
+                   static_cast<int32_t>(yieldType.getNumElements() * elementSize));
+      return ReturnPathLoweringResult::Handled;
+    }
+
+    auto storedType = dyn_cast<RankedTensorType>(yieldValue.getType());
+    if (!storedType) {
+      computeOp.emitOpError("has an unsupported non-ranked concat-return helper yield during Spatial-to-PIM lowering");
+      return ReturnPathLoweringResult::Failure;
+    }
+    rewriter.setInsertionPointAfterValue(yieldValue);
+    Value outputTensor = state.outputTensors[concatReturnUse->returnIndex](rewriter, loc);
+    auto outputType = cast<ShapedType>(outputTensor.getType());
+    for (int64_t linearIndex = 0; linearIndex < storedType.getNumElements(); ++linearIndex) {
+      SmallVector<int64_t> sourceIndices = expandFlatElementIndex(linearIndex, storedType.getShape());
+      for (auto [dim, idx] : llvm::enumerate(sourceIndices))
+        sourceIndices[dim] = concatReturnUse->sliceOffsets[dim] + idx;
+
+      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");
+        return ReturnPathLoweringResult::Failure;
+      }
+
+      SmallVector<OpFoldResult> extractOffsets;
+      SmallVector<OpFoldResult> extractSizes;
+      SmallVector<OpFoldResult> extractStrides;
+      extractOffsets.reserve(storedType.getRank());
+      extractSizes.reserve(storedType.getRank());
+      extractStrides.reserve(storedType.getRank());
+      for (int64_t idx : expandFlatElementIndex(linearIndex, storedType.getShape())) {
+        extractOffsets.push_back(rewriter.getIndexAttr(idx));
+        extractSizes.push_back(rewriter.getIndexAttr(1));
+        extractStrides.push_back(rewriter.getIndexAttr(1));
+      }
+
+      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.setInsertionPointAfter(elementSlice);
+
+      int64_t destinationFlatOffset = computeFlatElementIndex(destinationIndices, outputType.getShape());
+      outputTensor = emitHostCopy(rewriter,
+                                  loc,
+                                  outputTensor,
+                                  elementSlice.getResult(),
+                                  static_cast<int32_t>(destinationFlatOffset * elementSize),
+                                  0,
+                                  static_cast<int32_t>(elementSize));
+    }
+    return ReturnPathLoweringResult::Handled;
+  }
+
+  return ReturnPathLoweringResult::NotReturnPath;
+}
+
+void replaceReturnWithOutputBuffers(func::ReturnOp returnOp, IRRewriter& rewriter, ReturnPathState& state) {
+  auto markOwnedReturnChain = [&](Operation* op, auto&& markOwnedReturnChain) -> void {
+    if (!op)
+      return;
+
+    bool isExclusivelyOwnedByReturnChain = op->use_empty();
+    if (!isExclusivelyOwnedByReturnChain && op->hasOneUse()) {
+      Operation* onlyUser = *op->getUsers().begin();
+      isExclusivelyOwnedByReturnChain =
+        isa<func::ReturnOp, tensor::ConcatOp, spatial::SpatConcatOp, pim::PimConcatOp, spatial::SpatCompute>(onlyUser)
+        || isReturnHelperChainOp(onlyUser);
+    }
+    if (!isExclusivelyOwnedByReturnChain)
+      return;
+
+    if (isReturnHelperChainOp(op)) {
+      Value source = op->getOperand(0);
+      markOpToRemove(state, op);
+      markOwnedReturnChain(source.getDefiningOp(), markOwnedReturnChain);
+      return;
+    }
+
+    if (auto computeOp = dyn_cast<spatial::SpatCompute>(op)) {
+      markOpToRemove(state, computeOp);
+      if (!computeOp.getInputs().empty())
+        for (Value input : computeOp.getInputs())
+          markOwnedReturnChain(input.getDefiningOp(), markOwnedReturnChain);
+      return;
+    }
+
+    if (auto concatOp = dyn_cast<tensor::ConcatOp>(op)) {
+      markOpToRemove(state, concatOp);
+      for (Value operand : concatOp.getOperands())
+        markOwnedReturnChain(operand.getDefiningOp(), markOwnedReturnChain);
+      return;
+    }
+
+    if (auto concatOp = dyn_cast<spatial::SpatConcatOp>(op)) {
+      markOpToRemove(state, concatOp);
+      for (Value operand : concatOp.getInputs())
+        markOwnedReturnChain(operand.getDefiningOp(), markOwnedReturnChain);
+      return;
+    }
+
+    if (auto concatOp = dyn_cast<pim::PimConcatOp>(op)) {
+      markOpToRemove(state, concatOp);
+      for (Value operand : concatOp.getInputs())
+        markOwnedReturnChain(operand.getDefiningOp(), markOwnedReturnChain);
+    }
+  };
+
+  SmallVector<Value> originalOperands(returnOp.getOperands().begin(), returnOp.getOperands().end());
+  auto loc = returnOp.getLoc();
+  for (auto it : llvm::enumerate(originalOperands)) {
+    size_t orderWithinReturn = it.index();
+    Operation* returnOperand = it.value().getDefiningOp();
+    rewriter.setInsertionPoint(returnOp);
+    Value outputTensor = state.outputTensors[orderWithinReturn](rewriter, loc);
+    rewriter.modifyOpInPlace(returnOp, [&] { returnOp.setOperand(orderWithinReturn, outputTensor); });
+    markOwnedReturnChain(returnOperand, markOwnedReturnChain);
+  }
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.hpp

@@ -0,0 +1,37 @@
+#pragma once
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/PatternMatch.h"
+
+#include <functional>
+
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+namespace onnx_mlir {
+
+using OutputTensorFactory = std::function<mlir::Value(mlir::IRRewriter& rewriter, mlir::Location loc)>;
+
+struct ReturnPathState {
+  llvm::SmallVectorImpl<OutputTensorFactory>& outputTensors;
+  llvm::SmallVectorImpl<mlir::Operation*>& operationsToRemove;
+};
+
+enum class ReturnPathLoweringResult {
+  Handled,
+  NotReturnPath,
+  Failure
+};
+
+void addReturnOutputBuffers(mlir::func::ReturnOp returnOp,
+                            mlir::IRRewriter& rewriter,
+                            llvm::SmallVectorImpl<OutputTensorFactory>& outputTensors);
+
+ReturnPathLoweringResult lowerComputeResultReturnPath(spatial::SpatCompute computeOp,
+                                                      mlir::OpResult result,
+                                                      mlir::Value yieldValue,
+                                                      ReturnPathState& state,
+                                                      mlir::IRRewriter& rewriter);
+
+void replaceReturnWithOutputBuffers(mlir::func::ReturnOp returnOp, mlir::IRRewriter& rewriter, ReturnPathState& state);
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/TensorPackingPatterns.cpp

@@ -0,0 +1,113 @@
+#include "src/Accelerators/PIM/Conversion/SpatialToPim/TensorPackingPatterns.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+// Replaces concat-of-adjacent-slices with one packed slice to keep batch sends compact.
+struct FoldConcatOfContiguousSlices : OpRewritePattern<tensor::ConcatOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(tensor::ConcatOp op, PatternRewriter& rewriter) const override {
+    if (op.getDim() != 0)
+      return failure();
+
+    Value packed = createPackedExtractSliceTensor(op.getInputs(), rewriter, op.getLoc());
+    if (!packed)
+      return failure();
+
+    rewriter.replaceOp(op, packed);
+    return success();
+  }
+};
+
+} // namespace
+
+RankedTensorType getPackedTensorType(RankedTensorType elementType, int64_t count) {
+  SmallVector<int64_t> packedShape(elementType.getShape().begin(), elementType.getShape().end());
+  packedShape[0] *= count;
+  return RankedTensorType::get(packedShape, elementType.getElementType());
+}
+
+Value createPackedExtractSliceTensor(ValueRange values, OpBuilder& builder, Location loc) {
+  if (values.empty())
+    return {};
+  if (values.size() == 1)
+    return values.front();
+
+  auto firstSliceOp = values.front().getDefiningOp<tensor::ExtractSliceOp>();
+  if (!firstSliceOp)
+    return {};
+
+  auto firstType = dyn_cast<RankedTensorType>(firstSliceOp.getResult().getType());
+  auto sourceType = dyn_cast<RankedTensorType>(firstSliceOp.getSource().getType());
+  if (!firstType || !sourceType || !firstType.hasStaticShape() || !sourceType.hasStaticShape()
+      || firstType.getRank() == 0)
+    return {};
+
+  auto hasStaticValues = [](ArrayRef<int64_t> values) {
+    return llvm::all_of(values, [](int64_t value) { return !ShapedType::isDynamic(value); });
+  };
+  if (!hasStaticValues(firstSliceOp.getStaticOffsets()) || !hasStaticValues(firstSliceOp.getStaticSizes())
+      || !hasStaticValues(firstSliceOp.getStaticStrides()))
+    return {};
+
+  ArrayRef<int64_t> firstOffsets = firstSliceOp.getStaticOffsets();
+  ArrayRef<int64_t> firstSizes = firstSliceOp.getStaticSizes();
+  ArrayRef<int64_t> firstStrides = firstSliceOp.getStaticStrides();
+  int64_t rowsPerValue = firstSizes[0];
+  if (ShapedType::isDynamic(rowsPerValue))
+    return {};
+
+  for (size_t index = 1; index < values.size(); ++index) {
+    auto sliceOp = values[index].getDefiningOp<tensor::ExtractSliceOp>();
+    if (!sliceOp || sliceOp.getSource() != firstSliceOp.getSource()
+        || sliceOp.getResult().getType() != firstSliceOp.getResult().getType()
+        || !hasStaticValues(sliceOp.getStaticOffsets()) || !hasStaticValues(sliceOp.getStaticSizes())
+        || !hasStaticValues(sliceOp.getStaticStrides()))
+      return {};
+
+    if (sliceOp.getStaticSizes() != firstSizes || sliceOp.getStaticStrides() != firstStrides)
+      return {};
+
+    if (sliceOp.getStaticOffsets()[0] != firstOffsets[0] + static_cast<int64_t>(index) * rowsPerValue)
+      return {};
+
+    for (int64_t dim = 1; dim < firstType.getRank(); ++dim)
+      if (sliceOp.getStaticOffsets()[dim] != firstOffsets[dim])
+        return {};
+  }
+
+  auto packedType = getPackedTensorType(firstType, static_cast<int64_t>(values.size()));
+  SmallVector<OpFoldResult> offsets;
+  SmallVector<OpFoldResult> sizes;
+  SmallVector<OpFoldResult> strides;
+  offsets.reserve(firstType.getRank());
+  sizes.reserve(firstType.getRank());
+  strides.reserve(firstType.getRank());
+
+  offsets.push_back(builder.getIndexAttr(firstOffsets[0]));
+  sizes.push_back(builder.getIndexAttr(rowsPerValue * static_cast<int64_t>(values.size())));
+  strides.push_back(builder.getIndexAttr(firstStrides[0]));
+  for (int64_t dim = 1; dim < firstType.getRank(); ++dim) {
+    offsets.push_back(builder.getIndexAttr(firstOffsets[dim]));
+    sizes.push_back(builder.getIndexAttr(firstSizes[dim]));
+    strides.push_back(builder.getIndexAttr(firstStrides[dim]));
+  }
+
+  bool coversWholeSource = packedType == sourceType;
+  for (int64_t dim = 0; coversWholeSource && dim < sourceType.getRank(); ++dim)
+    coversWholeSource = firstOffsets[dim] == 0 && firstStrides[dim] == 1;
+  if (coversWholeSource)
+    return firstSliceOp.getSource();
+
+  return tensor::ExtractSliceOp::create(builder, loc, packedType, firstSliceOp.getSource(), offsets, sizes, strides)
+    .getResult();
+}
+
+void populateTensorPackingPatterns(RewritePatternSet& patterns) {
+  patterns.add<FoldConcatOfContiguousSlices>(patterns.getContext());
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/TensorPackingPatterns.hpp

@@ -0,0 +1,13 @@
+#pragma once
+
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/PatternMatch.h"
+
+namespace onnx_mlir {
+
+mlir::RankedTensorType getPackedTensorType(mlir::RankedTensorType elementType, int64_t count);
+mlir::Value createPackedExtractSliceTensor(mlir::ValueRange values, mlir::OpBuilder& builder, mlir::Location loc);
+
+void populateTensorPackingPatterns(mlir::RewritePatternSet& patterns);
+
+} // namespace onnx_mlir