remove accidental copy-paste

AGENTS.md

@@ -6,6 +6,7 @@
 * Always try the release build first before building with the debug version
 * Use the debug build only when it is useful to obtain a clear stack trace with symbols, inspect names, place breakpoints, or test a small case interactively
 * The debug build is very slow, so use it only on small fast tests such as operation validations, not on network validations
+* Always prepend rtk to shell commands if missing and if rtk is available
 
 # Core engineering philosophy
 

backend-simulators/pim/pim-simulator/src/lib/instruction_set/isa.rs

@@ -258,24 +258,23 @@ where
 
     let (memory, crossbars) = core.get_memory_crossbar();
     let crossbar = crossbars.get_mut(group).unwrap();
-    let crossbar_stored_bytes = crossbar.stored_bytes();
-    let crossbar_byte_width = crossbar.width();
-
-    let crossbar_elem_width = crossbar_byte_width / size_of::<M>();
-    ensure!(
-        crossbar_byte_width % size_of::<M>() == 0,
-        "M not divisor of the crosbbar size"
-    );
-
     let crossbar_height = crossbar.height();
-    let crossbar_byte_size = crossbar_byte_width * crossbar_height;
+    let crossbar_stored_bytes = crossbar.stored_bytes();
+    let bytes_per_column = crossbar_height * size_of::<M>();
+    ensure!(bytes_per_column != 0, "crossbar height can not be zero");
+    ensure!(
+        crossbar_stored_bytes % bytes_per_column == 0,
+        "Stored crossbar bytes do not describe an integral number of columns"
+    );
+    let crossbar_elem_width = crossbar_stored_bytes / bytes_per_column;
+    ensure!(crossbar_elem_width != 0, "Crossbar contains no stored columns");
 
     let loads = memory
         .reserve_load(r1_val, crossbar_height * size_of::<F>())?
         .execute_load::<F>()?;
     let load = loads[0];
     let vec: Cow<[M]> = load.up();
-    let matrix = crossbar.load::<M>(crossbar_byte_size)?[0];
+    let matrix = crossbar.load::<M>(crossbar_stored_bytes)?[0];
 
     // --- FAER IMPLEMENTATION ---
 

src/PIM/CMakeLists.txt

@@ -117,7 +117,6 @@ add_pim_library(OMPIMAccel
   SpatialOps
   PimOps
   OMONNXToSpatial
-  OMSpatialToGraphviz
   OMSpatialToPim
   OMPimCommon
   OMPimBufferization

src/PIM/Common/CMakeLists.txt

@@ -5,9 +5,11 @@ add_pim_library(OMPimCommon
   IR/ConstantUtils.cpp
   IR/CoreBlockUtils.cpp
   IR/EntryPointUtils.cpp
+  IR/IndexingUtils.cpp
   IR/LoopUtils.cpp
   IR/ShapeUtils.cpp
   IR/SubviewUtils.cpp
+  IR/TensorSliceUtils.cpp
   IR/WeightUtils.cpp
   Support/CheckedArithmetic.cpp
   Support/DebugDump.cpp

src/PIM/Common/IR/AddressAnalysis.cpp

@@ -34,12 +34,25 @@ mlir::Value resolveAlias(mlir::Value value, const StaticValueKnowledge* knowledg
 
 llvm::FailureOr<CompiledIndexExpr> compileIndexValueImpl(mlir::Value value);
 llvm::FailureOr<CompiledAddressExpr> compileContiguousAddressExprImpl(mlir::Value value);
+mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnowledge* knowledge);
 
 template <typename... Args>
 CompiledIndexExpr makeCompiledIndexExpr(Args&&... args) {
   return CompiledIndexExpr(std::make_shared<CompiledIndexExprNode>(std::forward<Args>(args)...));
 }
 
+static mlir::Value resolveForYieldedAliasToInit(mlir::scf::ForOp forOp,
+                                                mlir::Value yieldedValue,
+                                                const StaticValueKnowledge* knowledge) {
+  yieldedValue = resolveLoopCarriedAliasImpl(yieldedValue, knowledge);
+  if (auto blockArgument = mlir::dyn_cast<mlir::BlockArgument>(yieldedValue)) {
+    if (blockArgument.getOwner() == forOp.getBody() && blockArgument.getArgNumber() > 0
+        && static_cast<unsigned>(blockArgument.getArgNumber() - 1) < forOp.getInitArgs().size())
+      return resolveLoopCarriedAliasImpl(forOp.getInitArgs()[blockArgument.getArgNumber() - 1], knowledge);
+  }
+  return yieldedValue;
+}
+
 mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnowledge* knowledge) {
   value = resolveAlias(value, knowledge);
 
@@ -56,6 +69,15 @@ mlir::Value resolveLoopCarriedAliasImpl(mlir::Value value, const StaticValueKnow
         return resolveLoopCarriedAliasImpl(tiedOperand->get(), knowledge);
   }
 
+  if (auto forOp = mlir::dyn_cast<mlir::scf::ForOp>(definingOp)) {
+    auto result = mlir::dyn_cast<mlir::OpResult>(value);
+    if (result) {
+      auto yieldOp = mlir::dyn_cast<mlir::scf::YieldOp>(forOp.getBody()->getTerminator());
+      if (yieldOp && result.getResultNumber() < yieldOp.getNumOperands())
+        return resolveForYieldedAliasToInit(forOp, yieldOp.getOperand(result.getResultNumber()), knowledge);
+    }
+  }
+
   if (auto castOp = mlir::dyn_cast<mlir::memref::CastOp>(definingOp))
     return resolveLoopCarriedAliasImpl(castOp.getSource(), knowledge);
   if (auto collapseOp = mlir::dyn_cast<mlir::memref::CollapseShapeOp>(definingOp))
@@ -499,16 +521,25 @@ llvm::FailureOr<ResolvedContiguousAddress> resolveContiguousAddressImpl(mlir::Va
         return mlir::failure();
 
       auto yieldOp = mlir::cast<mlir::scf::YieldOp>(forOp.getBody()->getTerminator());
-      mlir::Value yieldedValue = resolveLoopCarriedAliasImpl(yieldOp.getOperand(result.getResultNumber()), knowledge);
-      if (auto blockArgument = mlir::dyn_cast<mlir::BlockArgument>(yieldedValue)) {
-        if (blockArgument.getOwner() == forOp.getBody() && blockArgument.getArgNumber() > 0
-            && static_cast<unsigned>(blockArgument.getArgNumber() - 1) < forOp.getInitArgs().size()) {
-          value = resolveAlias(forOp.getInitArgs()[blockArgument.getArgNumber() - 1], knowledge);
+      value = resolveForYieldedAliasToInit(forOp, yieldOp.getOperand(result.getResultNumber()), knowledge);
       continue;
     }
-      }
 
-      value = yieldedValue;
+    if (auto ifOp = mlir::dyn_cast<mlir::scf::IfOp>(definingOp)) {
+      auto result = mlir::dyn_cast<mlir::OpResult>(value);
+      if (!result)
+        return mlir::failure();
+
+      auto condition = resolveIndexValueImpl(ifOp.getCondition(), knowledge);
+      if (failed(condition))
+        return mlir::failure();
+
+      mlir::Region& selectedRegion = *condition != 0 ? ifOp.getThenRegion() : ifOp.getElseRegion();
+      auto yieldOp = mlir::dyn_cast<mlir::scf::YieldOp>(selectedRegion.front().getTerminator());
+      if (!yieldOp || result.getResultNumber() >= yieldOp.getNumOperands())
+        return mlir::failure();
+
+      value = resolveLoopCarriedAliasImpl(yieldOp.getOperand(result.getResultNumber()), knowledge);
       continue;
     }
 
@@ -609,17 +640,35 @@ llvm::FailureOr<CompiledAddressExpr> compileContiguousAddressExprImpl(mlir::Valu
         return mlir::failure();
 
       auto yieldOp = mlir::cast<mlir::scf::YieldOp>(forOp.getBody()->getTerminator());
-      mlir::Value yieldedValue = yieldOp.getOperand(result.getResultNumber());
-      if (auto blockArgument = mlir::dyn_cast<mlir::BlockArgument>(yieldedValue)) {
-        if (blockArgument.getOwner() == forOp.getBody() && blockArgument.getArgNumber() > 0
-            && static_cast<unsigned>(blockArgument.getArgNumber() - 1) < forOp.getInitArgs().size()) {
-          value = forOp.getInitArgs()[blockArgument.getArgNumber() - 1];
+      value = resolveForYieldedAliasToInit(forOp, yieldOp.getOperand(result.getResultNumber()), nullptr);
       continue;
     }
-      }
 
-      value = yieldedValue;
-      continue;
+    if (auto ifOp = mlir::dyn_cast<mlir::scf::IfOp>(definingOp)) {
+      auto result = mlir::dyn_cast<mlir::OpResult>(value);
+      if (!result)
+        return mlir::failure();
+
+      auto thenYield = mlir::dyn_cast<mlir::scf::YieldOp>(ifOp.getThenRegion().front().getTerminator());
+      auto elseYield = mlir::dyn_cast<mlir::scf::YieldOp>(ifOp.getElseRegion().front().getTerminator());
+      if (!thenYield || !elseYield || result.getResultNumber() >= thenYield.getNumOperands()
+          || result.getResultNumber() >= elseYield.getNumOperands()) {
+        return mlir::failure();
+      }
+
+      auto thenAddress = compileContiguousAddressExprImpl(thenYield.getOperand(result.getResultNumber()));
+      auto elseAddress = compileContiguousAddressExprImpl(elseYield.getOperand(result.getResultNumber()));
+      if (failed(thenAddress) || failed(elseAddress) || thenAddress->base != elseAddress->base)
+        return mlir::failure();
+
+      auto condition = compileIndexValueImpl(ifOp.getCondition());
+      if (failed(condition))
+        return mlir::failure();
+
+      CompiledIndexExprNode selectExpr;
+      selectExpr.kind = CompiledIndexExprNode::Kind::Select;
+      selectExpr.operands = {*condition, thenAddress->byteOffset, elseAddress->byteOffset};
+      return CompiledAddressExpr {thenAddress->base, makeCompiledIndexExpr(std::move(selectExpr))};
     }
 
     if (auto subviewOp = mlir::dyn_cast<mlir::memref::SubViewOp>(definingOp)) {
@@ -801,7 +850,7 @@ llvm::FailureOr<ResolvedContiguousAddress> CompiledAddressExpr::evaluate(const S
   auto resolvedOffset = byteOffset.evaluate(knowledge);
   if (failed(resolvedOffset))
     return mlir::failure();
-  return ResolvedContiguousAddress {base, *resolvedOffset};
+  return ResolvedContiguousAddress {resolveAlias(base, &knowledge), *resolvedOffset};
 }
 
 } // namespace onnx_mlir

src/PIM/Common/IR/AffineUtils.cpp

@@ -69,6 +69,15 @@ Value affineMulConst(RewriterBase& rewriter, Location loc, Value value, int64_t
   return createOrFoldAffineApply(rewriter, loc, d0 * multiplier, ValueRange {value}, constantAnchor);
 }
 
+Value affineAddConst(RewriterBase& rewriter, Location loc, Value value, int64_t offset, Operation* constantAnchor) {
+  assert(constantAnchor && "expected a valid constant anchor");
+  if (offset == 0)
+    return value;
+
+  AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext());
+  return createOrFoldAffineApply(rewriter, loc, d0 + offset, ValueRange {value}, constantAnchor);
+}
+
 Value affineModConst(RewriterBase& rewriter, Location loc, Value value, int64_t divisor, Operation* constantAnchor) {
   assert(constantAnchor && "expected a valid constant anchor");
   assert(divisor > 0 && "expected a positive affine.mod divisor");
@@ -90,6 +99,34 @@ Value affineFloorDivConst(
   return createOrFoldAffineApply(rewriter, loc, d0.floorDiv(divisor), ValueRange {value}, constantAnchor);
 }
 
+Value affineAddModConst(
+  RewriterBase& rewriter, Location loc, Value value, int64_t offset, int64_t divisor, Operation* constantAnchor) {
+  assert(constantAnchor && "expected a valid constant anchor");
+  assert(divisor > 0 && "expected a positive affine.mod divisor");
+  if (divisor == 1)
+    return getOrCreateIndexConstant(rewriter, constantAnchor, 0);
+
+  AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext());
+  AffineExpr expr = d0;
+  if (offset != 0)
+    expr = expr + offset;
+  return createOrFoldAffineApply(rewriter, loc, expr % divisor, ValueRange {value}, constantAnchor);
+}
+
+Value affineAddFloorDivConst(
+  RewriterBase& rewriter, Location loc, Value value, int64_t offset, int64_t divisor, Operation* constantAnchor) {
+  assert(constantAnchor && "expected a valid constant anchor");
+  assert(divisor > 0 && "expected a positive affine.floor_div divisor");
+  if (divisor == 1)
+    return offset == 0 ? value : affineAddConst(rewriter, loc, value, offset, constantAnchor);
+
+  AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext());
+  AffineExpr expr = d0;
+  if (offset != 0)
+    expr = expr + offset;
+  return createOrFoldAffineApply(rewriter, loc, expr.floorDiv(divisor), ValueRange {value}, constantAnchor);
+}
+
 FailureOr<int64_t> evaluateAffineExpr(AffineExpr expr, ArrayRef<int64_t> dims, ArrayRef<int64_t> symbols) {
   if (auto constant = dyn_cast<AffineConstantExpr>(expr))
     return constant.getValue();

src/PIM/Common/IR/AffineUtils.hpp

@@ -29,6 +29,12 @@ mlir::Value affineMulConst(mlir::RewriterBase& rewriter,
                            int64_t multiplier,
                            mlir::Operation* constantAnchor);
 
+mlir::Value affineAddConst(mlir::RewriterBase& rewriter,
+                           mlir::Location loc,
+                           mlir::Value value,
+                           int64_t offset,
+                           mlir::Operation* constantAnchor);
+
 mlir::Value affineModConst(mlir::RewriterBase& rewriter,
                            mlir::Location loc,
                            mlir::Value value,
@@ -41,6 +47,20 @@ mlir::Value affineFloorDivConst(mlir::RewriterBase& rewriter,
                                 int64_t divisor,
                                 mlir::Operation* constantAnchor);
 
+mlir::Value affineAddModConst(mlir::RewriterBase& rewriter,
+                              mlir::Location loc,
+                              mlir::Value value,
+                              int64_t offset,
+                              int64_t divisor,
+                              mlir::Operation* constantAnchor);
+
+mlir::Value affineAddFloorDivConst(mlir::RewriterBase& rewriter,
+                                   mlir::Location loc,
+                                   mlir::Value value,
+                                   int64_t offset,
+                                   int64_t divisor,
+                                   mlir::Operation* constantAnchor);
+
 llvm::FailureOr<int64_t>
 evaluateAffineExpr(mlir::AffineExpr expr, llvm::ArrayRef<int64_t> dims, llvm::ArrayRef<int64_t> symbols = {});
 

src/PIM/Common/IR/BatchCoreUtils.cpp

@@ -1,5 +1,6 @@
 #include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 
 namespace onnx_mlir {
 
@@ -9,6 +10,65 @@ llvm::SmallVector<int32_t> getBatchCoreIds(pim::PimCoreBatchOp coreBatchOp) {
   return llvm::SmallVector<int32_t>(coreIdsAttr.asArrayRef().begin(), coreIdsAttr.asArrayRef().end());
 }
 
+mlir::FailureOr<std::optional<int32_t>>
+getOptionalScheduledCoreId(spatial::SpatScheduledCompute computeOp, llvm::StringRef fieldName) {
+  auto coreIdAttr = computeOp->getAttrOfType<mlir::IntegerAttr>(onnx_mlir::kCoreIdAttrName);
+  if (!coreIdAttr)
+    return std::optional<int32_t> {};
+  if (coreIdAttr.getInt() < 0) {
+    computeOp.emitOpError() << fieldName << " must be non-negative";
+    return mlir::failure();
+  }
+  auto checkedCoreId = pim::checkedI32(coreIdAttr.getInt(), computeOp, fieldName);
+  if (mlir::failed(checkedCoreId))
+    return mlir::failure();
+  return std::optional<int32_t> {*checkedCoreId};
+}
+
+mlir::FailureOr<int32_t> getRequiredScheduledCoreId(spatial::SpatScheduledCompute computeOp, llvm::StringRef fieldName) {
+  auto coreId = getOptionalScheduledCoreId(computeOp, fieldName);
+  if (mlir::failed(coreId))
+    return mlir::failure();
+  if (!*coreId) {
+    computeOp.emitOpError() << "missing required " << onnx_mlir::kCoreIdAttrName;
+    return mlir::failure();
+  }
+  return **coreId;
+}
+
+mlir::FailureOr<std::optional<llvm::SmallVector<int32_t>>>
+getOptionalScheduledBatchCoreIds(spatial::SpatScheduledComputeBatch computeBatchOp, llvm::StringRef fieldName) {
+  auto coreIdsAttr = computeBatchOp->getAttrOfType<mlir::DenseI32ArrayAttr>(onnx_mlir::kCoreIdsAttrName);
+  if (!coreIdsAttr)
+    return std::optional<llvm::SmallVector<int32_t>> {};
+
+  llvm::SmallVector<int32_t> coreIds;
+  coreIds.reserve(coreIdsAttr.size());
+  for (int32_t coreId : coreIdsAttr.asArrayRef()) {
+    if (coreId < 0) {
+      computeBatchOp.emitOpError() << fieldName << " values must be non-negative";
+      return mlir::failure();
+    }
+    auto checkedCoreId = pim::checkedI32(static_cast<int64_t>(coreId), computeBatchOp, fieldName);
+    if (mlir::failed(checkedCoreId))
+      return mlir::failure();
+    coreIds.push_back(*checkedCoreId);
+  }
+  return std::optional<llvm::SmallVector<int32_t>> {std::move(coreIds)};
+}
+
+mlir::FailureOr<llvm::SmallVector<int32_t>>
+getRequiredScheduledBatchCoreIds(spatial::SpatScheduledComputeBatch computeBatchOp, llvm::StringRef fieldName) {
+  auto coreIds = getOptionalScheduledBatchCoreIds(computeBatchOp, fieldName);
+  if (mlir::failed(coreIds))
+    return mlir::failure();
+  if (!*coreIds) {
+    computeBatchOp.emitOpError() << "missing required " << onnx_mlir::kCoreIdsAttrName;
+    return mlir::failure();
+  }
+  return std::move(**coreIds);
+}
+
 llvm::SmallVector<int32_t> getLaneChunkCoreIds(llvm::ArrayRef<int32_t> coreIds, size_t laneCount, unsigned lane) {
   llvm::SmallVector<int32_t> laneCoreIds;
   laneCoreIds.reserve(coreIds.size() / laneCount);

src/PIM/Common/IR/BatchCoreUtils.hpp

@@ -3,12 +3,26 @@
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 
+#include <optional>
+
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 namespace onnx_mlir {
 
 llvm::SmallVector<int32_t> getBatchCoreIds(pim::PimCoreBatchOp coreBatchOp);
 
+mlir::FailureOr<std::optional<int32_t>>
+getOptionalScheduledCoreId(spatial::SpatScheduledCompute computeOp, llvm::StringRef fieldName);
+
+mlir::FailureOr<int32_t> getRequiredScheduledCoreId(spatial::SpatScheduledCompute computeOp, llvm::StringRef fieldName);
+
+mlir::FailureOr<std::optional<llvm::SmallVector<int32_t>>>
+getOptionalScheduledBatchCoreIds(spatial::SpatScheduledComputeBatch computeBatchOp, llvm::StringRef fieldName);
+
+mlir::FailureOr<llvm::SmallVector<int32_t>>
+getRequiredScheduledBatchCoreIds(spatial::SpatScheduledComputeBatch computeBatchOp, llvm::StringRef fieldName);
+
 llvm::SmallVector<int32_t> getLaneChunkCoreIds(llvm::ArrayRef<int32_t> coreIds, size_t laneCount, unsigned lane);
 
 bool isExplicitHostMemCopyOperand(mlir::Operation* op, unsigned operandIndex);

src/PIM/Common/IR/CoreBlockUtils.cpp

@@ -74,6 +74,21 @@ walkPimCoreBlock(mlir::Block& block,
       continue;
     }
 
+    if (auto ifOp = mlir::dyn_cast<mlir::scf::IfOp>(op)) {
+      auto condition = resolveIndexValue(ifOp.getCondition(), knowledge);
+      if (failed(condition)) {
+        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen");
+        hasFailure = true;
+        continue;
+      }
+
+      mlir::Region& selectedRegion = *condition != 0 ? ifOp.getThenRegion() : ifOp.getElseRegion();
+      if (!selectedRegion.empty())
+        if (failed(walkPimCoreBlock(selectedRegion.front(), knowledge, callback)))
+          hasFailure = true;
+      continue;
+    }
+
     if (failed(callback(op, knowledge)))
       hasFailure = true;
   }
@@ -128,6 +143,22 @@ mlir::LogicalResult walkPimCoreBlockStructurally(
       continue;
     }
 
+    if (auto ifOp = mlir::dyn_cast<mlir::scf::IfOp>(op)) {
+      if (failed(resolveIndexValue(ifOp.getCondition(), knowledge))) {
+        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM verification");
+        hasFailure = true;
+        continue;
+      }
+
+      if (!ifOp.getThenRegion().empty())
+        if (failed(walkPimCoreBlockStructurally(ifOp.getThenRegion().front(), knowledge, callback)))
+          hasFailure = true;
+      if (!ifOp.getElseRegion().empty())
+        if (failed(walkPimCoreBlockStructurally(ifOp.getElseRegion().front(), knowledge, callback)))
+          hasFailure = true;
+      continue;
+    }
+
     if (failed(callback(op, knowledge)))
       hasFailure = true;
   }

src/PIM/Common/IR/IndexingUtils.cpp

@@ -1,6 +1,6 @@
 #include <algorithm>
 
-#include "IndexingUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/IndexingUtils.hpp"
 
 using namespace mlir;
 

src/PIM/Common/IR/ShapeUtils.cpp

@@ -1,6 +1,9 @@
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/ErrorHandling.h"
 
+#include <functional>
+
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 
 namespace onnx_mlir {
@@ -163,4 +166,80 @@ bool isContiguousSubviewWithDynamicOffsets(llvm::ArrayRef<int64_t> sourceShape,
   return true;
 }
 
+bool hasStaticPositiveShape(llvm::ArrayRef<int64_t> shape) {
+  return llvm::all_of(shape, [](int64_t dim) { return dim > 0; });
+}
+
+bool hasStaticPositiveShape(mlir::RankedTensorType type) {
+  return type.hasStaticShape() && hasStaticPositiveShape(type.getShape());
+}
+
+int64_t getStaticShapeElementCount(llvm::ArrayRef<int64_t> shape) {
+  return std::accumulate(shape.begin(), shape.end(), int64_t {1}, std::multiplies<int64_t> {});
+}
+
+llvm::SmallVector<int64_t> permuteShape(llvm::ArrayRef<int64_t> shape, llvm::ArrayRef<int64_t> permutation) {
+  llvm::SmallVector<int64_t> permutedShape;
+  permutedShape.reserve(permutation.size());
+  for (int64_t axis : permutation)
+    permutedShape.push_back(shape[axis]);
+  return permutedShape;
+}
+
+llvm::SmallVector<int64_t> invertPermutation(llvm::ArrayRef<int64_t> permutation) {
+  llvm::SmallVector<int64_t> inversePermutation(permutation.size());
+  for (auto [newIndex, oldIndex] : llvm::enumerate(permutation))
+    inversePermutation[oldIndex] = static_cast<int64_t>(newIndex);
+  return inversePermutation;
+}
+
+mlir::FailureOr<llvm::SmallVector<int64_t>>
+getTransposePermutationChecked(std::optional<mlir::ArrayAttr> permAttr, int64_t rank) {
+  llvm::SmallVector<int64_t> permutation;
+  if (!permAttr) {
+    permutation.reserve(rank);
+    for (int64_t dim = rank - 1; dim >= 0; --dim)
+      permutation.push_back(dim);
+    return permutation;
+  }
+
+  if (static_cast<int64_t>(permAttr->size()) != rank)
+    return mlir::failure();
+
+  permutation.reserve(permAttr->size());
+  llvm::SmallVector<bool> seen(rank, false);
+  for (mlir::IntegerAttr attr : permAttr->getAsRange<mlir::IntegerAttr>()) {
+    int64_t axis = attr.getInt();
+    if (axis < 0 || axis >= rank || seen[axis])
+      return mlir::failure();
+    seen[axis] = true;
+    permutation.push_back(axis);
+  }
+  return permutation;
+}
+
+llvm::SmallVector<mlir::OpFoldResult> getStaticIndexAttrs(mlir::Builder& builder, llvm::ArrayRef<int64_t> values) {
+  llvm::SmallVector<mlir::OpFoldResult> attrs;
+  attrs.reserve(values.size());
+  for (int64_t value : values)
+    attrs.push_back(builder.getIndexAttr(value));
+  return attrs;
+}
+
+llvm::SmallVector<mlir::OpFoldResult> getUnitStrides(mlir::PatternRewriter& rewriter, int64_t rank) {
+  return llvm::SmallVector<mlir::OpFoldResult>(rank, rewriter.getIndexAttr(1));
+}
+
+llvm::SmallVector<mlir::OpFoldResult> getZeroOffsets(mlir::PatternRewriter& rewriter, int64_t rank) {
+  return llvm::SmallVector<mlir::OpFoldResult>(rank, rewriter.getIndexAttr(0));
+}
+
+llvm::SmallVector<mlir::OpFoldResult> getStaticSizes(mlir::PatternRewriter& rewriter, llvm::ArrayRef<int64_t> shape) {
+  llvm::SmallVector<mlir::OpFoldResult> sizes;
+  sizes.reserve(shape.size());
+  for (int64_t dim : shape)
+    sizes.push_back(rewriter.getIndexAttr(dim));
+  return sizes;
+}
+
 } // namespace onnx_mlir

src/PIM/Common/IR/ShapeUtils.hpp

@@ -2,15 +2,23 @@
 
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/Value.h"
 
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallVector.h"
 
 #include <cstddef>
+#include <optional>
+#include <type_traits>
+#include <utility>
 
 namespace onnx_mlir {
 
+using HSliceId = size_t;
+using CoreId = size_t;
+
 llvm::SmallVector<int64_t> computeRowMajorStrides(llvm::ArrayRef<int64_t> shape);
 
 llvm::SmallVector<int64_t>
@@ -36,4 +44,69 @@ bool isContiguousSubviewWithDynamicOffsets(llvm::ArrayRef<int64_t> sourceShape,
                                            llvm::ArrayRef<int64_t> staticSizes,
                                            llvm::ArrayRef<int64_t> staticStrides);
 
+template <class A, class B, class C = std::common_type_t<A, B>>
+constexpr C ceilIntegerDivide(A a, B b) {
+  static_assert(std::is_integral_v<A>, "A must be an integer type");
+  static_assert(std::is_integral_v<B>, "B must be an integer type");
+  C ac = static_cast<C>(a);
+  C bc = static_cast<C>(b);
+  return 1 + (ac - 1) / bc;
+}
+
+template <class A, class B, class C = std::common_type_t<A, B>>
+constexpr std::pair<C, C> ceilIntegerDivideWithRemainder(A a, B b) {
+  static_assert(std::is_integral_v<A>, "A must be an integer type");
+  static_assert(std::is_integral_v<B>, "B must be an integer type");
+  C ac = static_cast<C>(a);
+  C bc = static_cast<C>(b);
+  return {ceilIntegerDivide(ac, bc), ac % bc};
+}
+
+template <class T>
+bool isVectorShape(mlir::ArrayRef<T> shape) {
+  return shape.size() == 2 && (shape[0] == 1 || shape[1] == 1);
+}
+
+template <class T>
+bool isMatrixShape(mlir::ArrayRef<T> shape) {
+  return shape.size() == 2;
+}
+
+template <class T>
+bool isHVectorShape(mlir::ArrayRef<T> shape) {
+  return shape.size() == 2 && shape[0] == 1;
+}
+
+inline auto getTensorShape(mlir::Value tensor) {
+  return mlir::cast<mlir::RankedTensorType>(tensor.getType()).getShape();
+}
+
+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();
+}
+
+bool hasStaticPositiveShape(mlir::ArrayRef<int64_t> shape);
+
+bool hasStaticPositiveShape(mlir::RankedTensorType type);
+
+int64_t getStaticShapeElementCount(mlir::ArrayRef<int64_t> shape);
+
+llvm::SmallVector<int64_t> permuteShape(mlir::ArrayRef<int64_t> shape, mlir::ArrayRef<int64_t> permutation);
+
+llvm::SmallVector<int64_t> invertPermutation(mlir::ArrayRef<int64_t> permutation);
+
+mlir::FailureOr<llvm::SmallVector<int64_t>> getTransposePermutationChecked(std::optional<mlir::ArrayAttr> permAttr,
+                                                                           int64_t rank);
+
+llvm::SmallVector<mlir::OpFoldResult> getStaticIndexAttrs(mlir::Builder& builder, llvm::ArrayRef<int64_t> values);
+
+llvm::SmallVector<mlir::OpFoldResult> getUnitStrides(mlir::PatternRewriter& rewriter, int64_t rank);
+
+llvm::SmallVector<mlir::OpFoldResult> getZeroOffsets(mlir::PatternRewriter& rewriter, int64_t rank);
+
+llvm::SmallVector<mlir::OpFoldResult> getStaticSizes(mlir::PatternRewriter& rewriter, llvm::ArrayRef<int64_t> shape);
+
 } // namespace onnx_mlir

src/PIM/Common/IR/TensorSliceUtils.cpp

@@ -0,0 +1,71 @@
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+Value extractAxisSlice(
+  PatternRewriter& rewriter, Location loc, Value source, int64_t axis, int64_t offset, int64_t size) {
+  auto sourceType = cast<RankedTensorType>(source.getType());
+  SmallVector<int64_t> resultShape(sourceType.getShape());
+  resultShape[axis] = size;
+  auto resultType = RankedTensorType::get(resultShape, sourceType.getElementType(), sourceType.getEncoding());
+
+  SmallVector<OpFoldResult> offsets = getZeroOffsets(rewriter, sourceType.getRank());
+  SmallVector<OpFoldResult> sizes = getStaticSizes(rewriter, sourceType.getShape());
+  offsets[axis] = rewriter.getIndexAttr(offset);
+  sizes[axis] = rewriter.getIndexAttr(size);
+  return tensor::ExtractSliceOp::create(
+           rewriter, loc, resultType, source, offsets, sizes, getUnitStrides(rewriter, sourceType.getRank()))
+    .getResult();
+}
+
+Value extractStaticSliceOrIdentity(RewriterBase& rewriter,
+                                   Location loc,
+                                   Value source,
+                                   RankedTensorType resultType,
+                                   ArrayRef<OpFoldResult> offsets,
+                                   ArrayRef<OpFoldResult> sizes,
+                                   ArrayRef<OpFoldResult> strides) {
+  auto sourceType = cast<RankedTensorType>(source.getType());
+  size_t rank = static_cast<size_t>(sourceType.getRank());
+
+  bool isIdentitySlice =
+    sourceType == resultType && sourceType.hasStaticShape() && offsets.size() == rank && sizes.size() == rank
+    && strides.size() == rank;
+  if (isIdentitySlice) {
+    ArrayRef<int64_t> sourceShape = sourceType.getShape();
+    for (auto [dim, offset, size, stride] : llvm::zip_equal(sourceShape, offsets, sizes, strides)) {
+      std::optional<int64_t> staticOffset = mlir::getConstantIntValue(offset);
+      std::optional<int64_t> staticSize = mlir::getConstantIntValue(size);
+      std::optional<int64_t> staticStride = mlir::getConstantIntValue(stride);
+      if (!staticOffset || !staticSize || !staticStride || *staticOffset != 0 || *staticSize != dim
+          || *staticStride != 1) {
+        isIdentitySlice = false;
+        break;
+      }
+    }
+  }
+
+  if (isIdentitySlice)
+    return source;
+
+  return tensor::ExtractSliceOp::create(rewriter, loc, resultType, source, offsets, sizes, strides).getResult();
+}
+
+Value insertStaticSlice(
+  PatternRewriter& rewriter, Location loc, Value source, Value dest, ArrayRef<OpFoldResult> offsets) {
+  auto sourceType = cast<RankedTensorType>(source.getType());
+  return tensor::InsertSliceOp::create(rewriter,
+                                       loc,
+                                       source,
+                                       dest,
+                                       offsets,
+                                       getStaticSizes(rewriter, sourceType.getShape()),
+                                       getUnitStrides(rewriter, sourceType.getRank()))
+    .getResult();
+}
+
+} // namespace onnx_mlir

src/PIM/Common/IR/TensorSliceUtils.hpp

@@ -0,0 +1,28 @@
+#pragma once
+
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/ValueRange.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+
+namespace onnx_mlir {
+
+mlir::Value extractAxisSlice(
+  mlir::PatternRewriter& rewriter, mlir::Location loc, mlir::Value source, int64_t axis, int64_t offset, int64_t size);
+
+mlir::Value extractStaticSliceOrIdentity(mlir::RewriterBase& rewriter,
+                                         mlir::Location loc,
+                                         mlir::Value source,
+                                         mlir::RankedTensorType resultType,
+                                         llvm::ArrayRef<mlir::OpFoldResult> offsets,
+                                         llvm::ArrayRef<mlir::OpFoldResult> sizes,
+                                         llvm::ArrayRef<mlir::OpFoldResult> strides);
+
+mlir::Value insertStaticSlice(mlir::PatternRewriter& rewriter,
+                              mlir::Location loc,
+                              mlir::Value source,
+                              mlir::Value dest,
+                              llvm::ArrayRef<mlir::OpFoldResult> offsets);
+
+} // namespace onnx_mlir

src/PIM/Common/LabeledList.hpp

@@ -1,315 +0,0 @@
-#pragma once
-
-#include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/ilist_node.h"
-#include "llvm/ADT/simple_ilist.h"
-
-#include <cassert>
-#include <iterator>
-#include <limits>
-#include <type_traits>
-
-namespace onnx_mlir {
-
-template <typename NodeT>
-class LabeledList;
-
-template <typename NodeT>
-class LabeledListNode : public llvm::ilist_node<NodeT> {
-  friend class LabeledList<NodeT>;
-
-public:
-  using Label = uint64_t;
-
-  LabeledListNode() = default;
-  LabeledListNode(const LabeledListNode&) = delete;
-  LabeledListNode(LabeledListNode&&) = default;
-  LabeledListNode& operator=(LabeledListNode&&) = delete;
-
-  ~LabeledListNode() { assert(owner_ == nullptr && "destroying a linked LabeledListNode"); }
-
-  bool isLinked() const { return owner_ != nullptr; }
-  Label getOrderLabel() const { return label; }
-
-  friend bool operator<(const LabeledListNode& lft, const LabeledListNode& rgt) { return lft.label < rgt.label; }
-
-private:
-  const void* owner_ = nullptr;
-  Label label = 0;
-};
-
-template <typename NodeT>
-class LabeledList {
-
-  using Label = typename NodeT::Label;
-
-  static constexpr Label kLowerSentinel = 0;
-  static constexpr Label kUpperSentinel = std::numeric_limits<Label>::max();
-  static constexpr Label kRelabelGap = 2;
-
-public:
-  using List = llvm::simple_ilist<NodeT>;
-  using Iterator = typename List::iterator;
-  using RIterator = typename List::reverse_iterator;
-  using ConstIterator = typename List::const_iterator;
-
-  LabeledList() = default;
-  LabeledList(const LabeledList&) = delete;
-  LabeledList& operator=(const LabeledList&) = delete;
-  LabeledList(LabeledList&&) = delete;
-  LabeledList& operator=(LabeledList&&) = delete;
-
-  ~LabeledList() { clear(); }
-
-  bool empty() const { return size_ == 0; }
-  size_t size() const { return size_; }
-
-  NodeT* front() { return empty() ? nullptr : &nodes_.front(); }
-  const NodeT* front() const { return empty() ? nullptr : &nodes_.front(); }
-
-  NodeT* back() { return empty() ? nullptr : &nodes_.back(); }
-  const NodeT* back() const { return empty() ? nullptr : &nodes_.back(); }
-
-  static NodeT* previous(NodeT* node) {
-    if (!node || !owner(node))
-      return nullptr;
-    auto* list = owner(node);
-    auto it = node->getIterator();
-    if (it == list->nodes_.begin())
-      return nullptr;
-    return &*std::prev(it);
-  }
-
-  static const NodeT* previous(const NodeT* node) {
-    if (!node || !owner(node))
-      return nullptr;
-    const auto* list = owner(node);
-    auto it = const_cast<NodeT*>(node)->getIterator();
-    if (it == list->nodes_.begin())
-      return nullptr;
-    return &*std::prev(it);
-  }
-
-  static NodeT* next(NodeT* node) {
-    if (!node || !owner(node))
-      return nullptr;
-    auto* list = owner(node);
-    auto it = std::next(node->getIterator());
-    if (it == list->nodes_.end())
-      return nullptr;
-    return &*it;
-  }
-
-  static const NodeT* next(const NodeT* node) {
-    if (!node || !owner(node))
-      return nullptr;
-    const auto* list = owner(node);
-    auto it = std::next(const_cast<NodeT*>(node)->getIterator());
-    if (it == list->nodes_.end())
-      return nullptr;
-    return &*it;
-  }
-
-  bool contains(const NodeT* node) const { return node && node->owner_ == this; }
-
-  Label getOrderLabel(const NodeT* node) const {
-    assert(contains(node) && "node must belong to this list");
-    return node->label;
-  }
-
-  bool comesBefore(const NodeT* lhs, const NodeT* rhs) const {
-    assert(contains(lhs) && contains(rhs) && "nodes must belong to this list");
-    return lhs->label < rhs->label;
-  }
-
-  void pushFront(NodeT* node) { insertBefore(front(), node); }
-
-  void pushBack(NodeT* node) { insertBefore(nullptr, node); }
-
-  void insertBefore(NodeT* nextNode, NodeT* node) {
-    assert(node && "cannot insert a null node");
-    assert(!node->owner_ && "node is already linked");
-    assert(nextNode == nullptr || contains(nextNode));
-
-    Iterator nextIt = nextNode ? getIteratorFor(nextNode) : nodes_.end();
-    nodes_.insert(nextIt, *node);
-    node->owner_ = this;
-    ++size_;
-    assignLabel(getIteratorFor(node));
-  }
-
-  void insertAfter(NodeT* prevNode, NodeT* node) {
-    assert(prevNode == nullptr || contains(prevNode));
-    if (prevNode == nullptr)
-      insertBefore(front(), node);
-    else
-      insertBefore(next(prevNode), node);
-  }
-
-  void remove(NodeT* node) {
-    assert(contains(node) && "node must belong to this list");
-    nodes_.remove(*node);
-    node->owner_ = nullptr;
-    node->label = 0;
-    --size_;
-  }
-
-  void moveBefore(NodeT* node, NodeT* nextNode) {
-    assert(contains(node) && "node must belong to this list");
-    assert(nextNode == nullptr || contains(nextNode));
-
-    Iterator nodeIt = getIteratorFor(node);
-    Iterator nextIt = nextNode ? getIteratorFor(nextNode) : nodes_.end();
-    if (nodeIt == nextIt || std::next(nodeIt) == nextIt)
-      return;
-
-    nodes_.splice(nextIt, nodes_, nodeIt);
-    assignLabel(getIteratorFor(node));
-  }
-
-  void moveAfter(NodeT* node, NodeT* prevNode) {
-    assert(contains(node) && "node must belong to this list");
-    assert(prevNode == nullptr || contains(prevNode));
-
-    Iterator nextIt = prevNode ? std::next(getIteratorFor(prevNode)) : nodes_.begin();
-    if (getIteratorFor(node) == nextIt)
-      return;
-    moveBefore(node, nextIt == nodes_.end() ? nullptr : &*nextIt);
-  }
-
-  void clear() {
-    while (!nodes_.empty()) {
-      NodeT* node = &nodes_.front();
-      node->owner_ = nullptr;
-      node->label = 0;
-      nodes_.remove(*node);
-    }
-    size_ = 0;
-  }
-
-  Iterator begin() { return nodes_.begin(); }
-  Iterator end() { return nodes_.end(); }
-
-  RIterator rbegin() { return nodes_.rbegin(); }
-  RIterator rend() { return nodes_.rend(); }
-
-private:
-  static const LabeledList* owner(const NodeT* node) { return static_cast<const LabeledList*>(node->owner_); }
-  static LabeledList* owner(NodeT* node) { return static_cast<LabeledList*>(const_cast<void*>(node->owner_)); }
-
-  static Label lowerLabel(const NodeT* node) { return node ? node->label : kLowerSentinel; }
-  static Label upperLabel(const NodeT* node) { return node ? node->label : kUpperSentinel; }
-
-  static Label labelGap(Label lower, Label upper) {
-    assert(lower < upper && "labels must be strictly ordered");
-    return upper - lower;
-  }
-
-  static bool hasMidpoint(Label lower, Label upper) { return labelGap(lower, upper) > 1; }
-
-  static bool hasRelabelSlack(Label lower, Label upper, size_t nodeCount) {
-    Label gap = labelGap(lower, upper);
-    return gap / static_cast<Label>(nodeCount + 1) >= kRelabelGap;
-  }
-
-  Iterator getIteratorFor(NodeT* node) { return node->getIterator(); }
-  ConstIterator getiteratorFor(const NodeT* node) const { return node->getIterator(); }
-
-  NodeT* previousNode(Iterator it) {
-    if (it == nodes_.begin())
-      return nullptr;
-    return &*std::prev(it);
-  }
-
-  const NodeT* previousNode(ConstIterator it) const {
-    if (it == nodes_.begin())
-      return nullptr;
-    return &*std::prev(it);
-  }
-
-  NodeT* nextNode(Iterator it) {
-    ++it;
-    if (it == nodes_.end())
-      return nullptr;
-    return &*it;
-  }
-
-  const NodeT* nextNode(ConstIterator it) const {
-    ++it;
-    if (it == nodes_.end())
-      return nullptr;
-    return &*it;
-  }
-
-  void assignLabel(Iterator it) {
-    Label lower = lowerLabel(previousNode(it));
-    Label upper = upperLabel(nextNode(it));
-    if (hasMidpoint(lower, upper)) {
-      (*it).label = lower + static_cast<Label>(labelGap(lower, upper) / 2);
-      return;
-    }
-
-    relabelAround(it);
-  }
-
-  void relabelAround(Iterator center) {
-    size_t targetCount = 1;
-    while (true) {
-      Iterator left = center;
-      Iterator right = center;
-      size_t actualCount = 1;
-      expandWindow(center, targetCount, left, right, actualCount);
-
-      Label lower = lowerLabel(previousNode(left));
-      Label upper = upperLabel(nextNode(right));
-      if (hasRelabelSlack(lower, upper, actualCount)) {
-        relabelWindow(left, actualCount, lower, upper);
-        return;
-      }
-
-      if (left == nodes_.begin() && nextNode(right) == nullptr) {
-        assert(hasRelabelSlack(lower, upper, actualCount) && "label space exhausted");
-        relabelWindow(left, actualCount, lower, upper);
-        return;
-      }
-
-      targetCount *= 2;
-    }
-  }
-
-  void expandWindow(Iterator center, size_t targetCount, Iterator& left, Iterator& right, size_t& actualCount) {
-    left = center;
-    right = center;
-    actualCount = 1;
-
-    while (actualCount < targetCount && (left != nodes_.begin() || nextNode(right) != nullptr)) {
-      if (left != nodes_.begin()) {
-        --left;
-        ++actualCount;
-        if (actualCount == targetCount)
-          break;
-      }
-      if (nextNode(right) != nullptr) {
-        ++right;
-        ++actualCount;
-      }
-    }
-  }
-
-  void relabelWindow(Iterator left, size_t nodeCount, Label lower, Label upper) {
-    assert(nodeCount > 0 && "relabel window must not be empty");
-    Label step = labelGap(lower, upper) / static_cast<Label>(nodeCount + 1);
-    assert(step >= 1 && "relabel step must be positive");
-
-    Iterator it = left;
-    for (size_t index = 1; index <= nodeCount; ++index) {
-      (*it).label = lower + step * index;
-      ++it;
-    }
-  }
-
-  List nodes_;
-  size_t size_ = 0;
-};
-
-} // namespace onnx_mlir

src/PIM/Common/PimCommon.hpp

@@ -15,6 +15,7 @@
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/EntryPointUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/IndexingUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
 #include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"

src/PIM/Common/Support/DebugDump.cpp

@@ -7,18 +7,26 @@
 
 namespace onnx_mlir {
 
-void dumpModule(mlir::ModuleOp moduleOp, const std::string& name) {
+std::fstream openDialectDumpFileWithExtension(const std::string& name, llvm::StringRef destination,  llvm::StringRef extension) {
   std::string outputDir = getOutputDir();
   if (outputDir.empty())
+    return {};
+
+  std::string dialectsDir = (outputDir + destination).str();
+  createDirectory(dialectsDir);
+  return std::fstream(dialectsDir + "/" + name + "." + extension.str(), std::ios::out);
+}
+
+void dumpModule(mlir::ModuleOp moduleOp, const std::string& name, bool assumeVerified) {
+  std::fstream file = openDialectDumpFileWithExtension(name, "/dialects", "mlir");
+  if (!file.is_open())
     return;
 
-  std::string dialectsDir = outputDir + "/dialects";
-  createDirectory(dialectsDir);
-
-  std::fstream file(dialectsDir + "/" + name + ".mlir", std::ios::out);
   llvm::raw_os_ostream os(file);
   mlir::OpPrintingFlags flags;
-  flags.elideLargeElementsAttrs().enableDebugInfo(true, false);
+  flags.elideLargeElementsAttrs().enableDebugInfo(false, false);
+  if (assumeVerified)
+    flags.assumeVerified();
   moduleOp.print(os, flags);
   os.flush();
   file.close();

src/PIM/Common/Support/DebugDump.hpp

@@ -1,13 +1,18 @@
 #pragma once
 
 #include "mlir/IR/BuiltinOps.h"
+#include "llvm/ADT/StringRef.h"
 
+#include <fstream>
 #include <string>
 
 namespace onnx_mlir {
 
 /// Emits a MLIR snapshot under the current compiler output
 /// directory for pass-level debugging.
-void dumpModule(mlir::ModuleOp moduleOp, const std::string& name);
+void dumpModule(mlir::ModuleOp moduleOp, const std::string& name, bool assumeVerified = false);
+
+/// Opens a file under the same dialect dump directory used by dumpModule.
+std::fstream openDialectDumpFileWithExtension(const std::string& name,llvm::StringRef destination = "/dialects",  llvm::StringRef extension = "mlir");
 
 } // namespace onnx_mlir

src/PIM/Common/Support/ReportUtils.cpp

@@ -17,6 +17,20 @@ std::fstream openReportFileWithExtension(const std::string& name, llvm::StringRe
 
 std::fstream openReportFile(const std::string& name) { return openReportFileWithExtension(name, "txt"); }
 
+std::fstream openAppendedReportFileWithExtension(const std::string& name, llvm::StringRef extension) {
+  std::string outputDir = getOutputDir();
+  if (outputDir.empty())
+    return {};
+
+  std::string reportsDir = outputDir + "/reports";
+  createDirectory(reportsDir);
+  return std::fstream(reportsDir + "/" + name + "." + extension.str(), std::ios::out | std::ios::app);
+}
+
+std::fstream openAppendedReportFile(const std::string& name) {
+  return openAppendedReportFileWithExtension(name, "txt");
+}
+
 std::string formatReportMemory(uint64_t bytes) {
   const char* units[] = {"B", "KB", "MB", "GB", "TB", "PB", "EB"};
   int i = 0;

src/PIM/Common/Support/ReportUtils.hpp

@@ -12,6 +12,8 @@ namespace onnx_mlir {
 
 std::fstream openReportFile(const std::string& name);
 std::fstream openReportFileWithExtension(const std::string& name, llvm::StringRef extension);
+std::fstream openAppendedReportFile(const std::string& name);
+std::fstream openAppendedReportFileWithExtension(const std::string& name, llvm::StringRef extension);
 std::string formatReportMemory(uint64_t bytes);
 
 struct ReportField {

src/PIM/Compiler/PimCodeGen.cpp

@@ -414,6 +414,9 @@ size_t PimAcceleratorMemory::getValueAddress(mlir::Value value,
                                              const StaticValueKnowledge& knowledge,
                                              std::optional<unsigned> lane) const {
   value = resolveCachedAlias(value, knowledge);
+
+  FailureOr<ResolvedContiguousAddress> resolvedAddress = resolveContiguousAddress(value, knowledge);
+  if (failed(resolvedAddress)) {
     auto compiledIt = compiledAddressExprs.find(value);
     if (compiledIt == compiledAddressExprs.end()) {
       auto compiledExpr = compileContiguousAddressExpr(value);
@@ -427,7 +430,7 @@ size_t PimAcceleratorMemory::getValueAddress(mlir::Value value,
       compiledIt = compiledAddressExprs.try_emplace(value, *compiledExpr).first;
     }
 
-  auto resolvedAddress = compiledIt->second.evaluate(knowledge, lane);
+    resolvedAddress = compiledIt->second.evaluate(knowledge, lane);
     if (failed(resolvedAddress)) {
       errs() << "Failed to evaluate contiguous address for value: ";
       value.print(errs());
@@ -440,6 +443,7 @@ size_t PimAcceleratorMemory::getValueAddress(mlir::Value value,
       }
       llvm_unreachable("Failed to resolve contiguous address");
     }
+  }
 
   MemoryValueKey key = getMemoryValueKey(resolvedAddress->base, lane);
   auto iter = memEntriesMap.find(key);
@@ -588,13 +592,37 @@ void PimCodeGen::emitInstruction(const pim_binary::InstructionRecord& instructio
   ++emittedInstructionCount;
   if (coreJsonStream)
     *coreJsonStream << json::Value(pim_binary::makeInstructionJson(instruction)) << ',';
+  updateScalarRegisterCache(instruction);
+}
+
+void PimCodeGen::updateScalarRegisterCache(const pim_binary::InstructionRecord& instruction) const {
+  switch (instruction.opcode) {
+  case pim_binary::Opcode::sldi:
+    scalarRegisterValues[instruction.rd] = instruction.r2OrImm;
+    break;
+  case pim_binary::Opcode::sld:
+  case pim_binary::Opcode::sadd:
+  case pim_binary::Opcode::ssub:
+  case pim_binary::Opcode::smul:
+  case pim_binary::Opcode::saddi:
+  case pim_binary::Opcode::smuli:
+    scalarRegisterValues[instruction.rd].reset();
+    break;
+  default:
+    break;
+  }
 }
 
 void PimCodeGen::genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate) const {
+  auto registerIndex = pim::checkedU8OrCrash(registerNumber, "register number");
+  auto immediateValue = pim::checkedI32OrCrash(immediate, "register immediate");
+  if (scalarRegisterValues[registerIndex] == immediateValue)
+    return;
+
   pim_binary::InstructionRecord instruction;
   instruction.opcode = pim_binary::Opcode::sldi;
-  instruction.rd = static_cast<uint8_t>(registerNumber);
-  instruction.r2OrImm = pim::checkedI32OrCrash(immediate, "register immediate");
+  instruction.rd = registerIndex;
+  instruction.r2OrImm = immediateValue;
   emitInstruction(instruction);
 }
 
@@ -1090,7 +1118,8 @@ enum class CompiledCoreOpKind : uint8_t {
 struct CompiledCoreNode {
   enum class Kind : uint8_t {
     Op,
-    Loop
+    Loop,
+    If
   };
 
   Kind kind = Kind::Op;
@@ -1099,7 +1128,10 @@ struct CompiledCoreNode {
   CompiledIndexExpr lowerBound;
   CompiledIndexExpr upperBound;
   CompiledIndexExpr step;
+  CompiledIndexExpr condition;
   std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> loopBody;
+  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> thenBody;
+  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> elseBody;
 };
 
 static FailureOr<CompiledCoreOpKind> classifyCompiledCoreOpKind(Operation& op) {
@@ -1177,6 +1209,28 @@ compileCoreEmissionPlan(Block& block, Operation* weightOwner, llvm::SmallVectorI
       continue;
     }
 
+    if (auto ifOp = dyn_cast<mlir::scf::IfOp>(op)) {
+      auto condition = compileIndexExpr(ifOp.getCondition());
+      if (failed(condition)) {
+        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen");
+        return failure();
+      }
+
+      CompiledCoreNode ifNode;
+      ifNode.kind = CompiledCoreNode::Kind::If;
+      ifNode.op = ifOp.getOperation();
+      ifNode.condition = *condition;
+      ifNode.thenBody = std::make_unique<llvm::SmallVector<CompiledCoreNode, 8>>();
+      if (failed(compileCoreEmissionPlan(ifOp.getThenRegion().front(), weightOwner, *ifNode.thenBody)))
+        return failure();
+      ifNode.elseBody = std::make_unique<llvm::SmallVector<CompiledCoreNode, 8>>();
+      if (!ifOp.getElseRegion().empty())
+        if (failed(compileCoreEmissionPlan(ifOp.getElseRegion().front(), weightOwner, *ifNode.elseBody)))
+          return failure();
+      plan.push_back(std::move(ifNode));
+      continue;
+    }
+
     auto opKind = classifyCompiledCoreOpKind(op);
     if (failed(opKind)) {
       InFlightDiagnostic diag = op.emitError() << "unsupported codegen for op '" << op.getName().getStringRef() << "'";
@@ -1239,6 +1293,26 @@ static LogicalResult executeCompiledCorePlan(
       continue;
     }
 
+    if (node.kind == CompiledCoreNode::Kind::If) {
+      auto condition = node.condition.evaluate(knowledge);
+      auto ifOp = cast<mlir::scf::IfOp>(node.op);
+      if (failed(condition)) {
+        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen");
+        return failure();
+      }
+
+      const auto& selectedBody = *condition != 0 ? node.thenBody : node.elseBody;
+      if (selectedBody && failed(executeCompiledCorePlan(*selectedBody,
+                                                         coreCodeGen,
+                                                         knowledge,
+                                                         resolveWeightSlot,
+                                                         processedOperations,
+                                                         batchLane,
+                                                         batchLaneCount)))
+        return failure();
+      continue;
+    }
+
     switch (node.opKind) {
     case CompiledCoreOpKind::Load:
       coreCodeGen.codeGenLoadOp(cast<pim::PimMemCopyHostToDevOp>(node.op), knowledge);

src/PIM/Compiler/PimCodeGen.hpp

@@ -9,6 +9,7 @@
 #include "llvm/Support/JSON.h"
 #include "llvm/Support/raw_os_ostream.h"
 
+#include <array>
 #include <fstream>
 #include <limits>
 #include <optional>
@@ -170,6 +171,7 @@ class PimCodeGen {
   const llvm::DenseMap<size_t, size_t>& emittedCoreIds;
   std::optional<unsigned> batchLane;
   mutable uint32_t emittedInstructionCount = 0;
+  mutable std::array<std::optional<int32_t>, 256> scalarRegisterValues = {};
 
   size_t addressOf(mlir::Value value, const StaticValueKnowledge& knowledge) const {
     return memory.getValueAddress(value, knowledge, batchLane);
@@ -177,6 +179,7 @@ class PimCodeGen {
   size_t remapCoreId(size_t coreId) const;
 
   void emitInstruction(const pim_binary::InstructionRecord& instruction) const;
+  void updateScalarRegisterCache(const pim_binary::InstructionRecord& instruction) const;
 
   void genSetRegisterImmediateUnsigned(size_t registerNumber, size_t immediate) const;
   void setupRd(size_t rdAddress, size_t rdOffset) const;

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -32,6 +32,43 @@ llvm::cl::opt<PimMemoryReportLevel> pimMemoryReport(
   llvm::cl::init(PimMemoryReportNone),
   llvm::cl::cat(OnnxMlirOptions));
 
+llvm::cl::opt<PimConvLoweringType> pimConvLowering(
+  "pim-conv-lowering",
+  llvm::cl::desc("Convolution lowering strategy for PIM"),
+  llvm::cl::values(clEnumValN(PimConvLoweringAuto, "auto", "Select the Conv lowering strategy automatically")),
+  llvm::cl::values(clEnumValN(PimConvLoweringLegacy, "legacy", "Use the legacy explicit-im2col Conv lowering")),
+  llvm::cl::values(clEnumValN(PimConvLoweringDepthwise, "depthwise", "Force the depthwise-specialized Conv lowering")),
+  llvm::cl::values(
+    clEnumValN(PimConvLoweringPackedIm2Col, "packed-im2col", "Use explicit im2col with packed multi-position GEMM")),
+  llvm::cl::values(clEnumValN(PimConvLoweringStreamedPatch,
+                              "streamed-patch",
+                              "Use streamed/chunked im2col rows without multi-position packing")),
+  llvm::cl::values(clEnumValN(PimConvLoweringStreamedPacked,
+                              "streamed-packed",
+                              "Use streamed/chunked im2col rows with packed multi-position GEMM")),
+  llvm::cl::values(clEnumValN(PimConvLoweringOutputChannelTiled,
+                              "output-channel-tiled",
+                              "Force Conv lowering that relies on Gemm output-channel tiling")),
+  llvm::cl::values(
+    clEnumValN(PimConvLoweringInputKTiled, "input-k-tiled", "Force Conv lowering that relies on Gemm K tiling")),
+  llvm::cl::values(clEnumValN(PimConvLoweringTiled2D,
+                              "tiled-2d",
+                              "Force Conv lowering that relies on Gemm 2D K/C tiling")),
+  llvm::cl::init(PimConvLoweringAuto),
+  llvm::cl::cat(OnnxMlirOptions));
+
+llvm::cl::opt<PimSpatialDataflowExportType> pimExportSpatialDataflow(
+  "pim-export-spatial-dataflow",
+  llvm::cl::desc("Emit Gephi-importable CSV dataflow reports around MergeComputeNodes materialization"),
+  llvm::cl::values(clEnumValN(SpatialDataflowExportNone, "none", "Do not emit Spatial dataflow CSV reports")),
+  llvm::cl::values(clEnumValN(SpatialDataflowExportPre, "pre", "Emit pre-materialization Spatial dataflow CSV reports")),
+  llvm::cl::values(
+    clEnumValN(SpatialDataflowExportPost, "post", "Emit post-materialization Spatial dataflow CSV reports")),
+  llvm::cl::values(
+    clEnumValN(SpatialDataflowExportBoth, "both", "Emit both pre- and post-materialization Spatial dataflow CSV reports")),
+  llvm::cl::init(SpatialDataflowExportNone),
+  llvm::cl::cat(OnnxMlirOptions));
+
 llvm::cl::opt<bool>
   pimOnlyCodegen("pim-only-codegen",
                  llvm::cl::desc("Only generate code for PIM (assume input is already in bufferized PIM IR)"),
@@ -49,11 +86,46 @@ llvm::cl::opt<bool> useExperimentalConvImpl("use-experimental-conv-impl",
                                             llvm::cl::init(false),
                                             llvm::cl::cat(OnnxMlirOptions));
 
+llvm::cl::opt<uint64_t> pimConvIm2colMaxElements(
+  "pim-conv-im2col-max-elements",
+  llvm::cl::desc("Maximum number of im2col elements to materialize globally for one Conv before streaming/chunking"),
+  llvm::cl::init(1ull << 20),
+  llvm::cl::cat(OnnxMlirOptions));
+
+llvm::cl::opt<uint64_t> pimConvStreamChunkPositions(
+  "pim-conv-stream-chunk-positions",
+  llvm::cl::desc("Maximum number of Conv output positions to materialize in one streamed chunk"),
+  llvm::cl::init(1024),
+  llvm::cl::cat(OnnxMlirOptions));
+
+llvm::cl::opt<bool> pimReportConvLowering("pim-report-conv-lowering",
+                                          llvm::cl::desc("Emit a bounded Conv lowering report"),
+                                          llvm::cl::init(true),
+                                          llvm::cl::cat(OnnxMlirOptions));
+
 llvm::cl::opt<bool> pimEmitJson("pim-emit-json",
                                 llvm::cl::desc("Also emit per-core JSON instruction files alongside binary .pim files"),
                                 llvm::cl::init(false),
                                 llvm::cl::cat(OnnxMlirOptions));
 
+llvm::cl::opt<bool> pimDetectCommunicationDeadlock(
+  "pim-detect-communication-deadlock",
+  llvm::cl::desc("Expensively simulate the statically expanded PIM send/receive order at verification time and fail if a blocking communication deadlock is found"),
+  llvm::cl::init(false),
+  llvm::cl::cat(OnnxMlirOptions));
+
+llvm::cl::opt<bool> pimMaterializeScalarFanoutGlobalOrder(
+  "pim-materialize-scalar-fanout-global-order",
+  llvm::cl::desc("Experimental expensive materializer mode: emit scalar-source fanout as globally ordered communication events instead of all-send fanout loops"),
+  llvm::cl::init(false),
+  llvm::cl::cat(OnnxMlirOptions));
+
+llvm::cl::opt<bool> pimTraceCommunicationMaterialization(
+  "pim-trace-communication-materialization",
+  llvm::cl::desc("Emit verbose materializer-time diagnostics and provenance attributes for every Spatial communication op"),
+  llvm::cl::init(false),
+  llvm::cl::cat(OnnxMlirOptions));
+
 llvm::cl::opt<size_t>
   crossbarSize("crossbar-size", llvm::cl::desc("Width and height of a single crossbar"), llvm::cl::init(2));
 

src/PIM/Compiler/PimCompilerOptions.hpp

@@ -30,19 +30,46 @@ typedef enum {
   PimMemoryReportFull = 2,
 } PimMemoryReportLevel;
 
+typedef enum {
+  PimConvLoweringAuto = 0,
+  PimConvLoweringLegacy = 1,
+  PimConvLoweringDepthwise = 2,
+  PimConvLoweringPackedIm2Col = 3,
+  PimConvLoweringStreamedPatch = 4,
+  PimConvLoweringStreamedPacked = 5,
+  PimConvLoweringOutputChannelTiled = 6,
+  PimConvLoweringInputKTiled = 7,
+  PimConvLoweringTiled2D = 8,
+} PimConvLoweringType;
+
+typedef enum {
+  SpatialDataflowExportNone = 0,
+  SpatialDataflowExportPre = 1,
+  SpatialDataflowExportPost = 2,
+  SpatialDataflowExportBoth = 3,
+} PimSpatialDataflowExportType;
+
 extern llvm::cl::OptionCategory OnnxMlirOptions;
 extern llvm::cl::opt<PimEmissionTargetType> pimEmissionTarget;
 extern llvm::cl::opt<PimMergeSchedulerType> pimMergeScheduler;
 extern llvm::cl::opt<PimMemoryReportLevel> pimMemoryReport;
+extern llvm::cl::opt<PimConvLoweringType> pimConvLowering;
+extern llvm::cl::opt<PimSpatialDataflowExportType> pimExportSpatialDataflow;
 
 extern llvm::cl::opt<bool> pimOnlyCodegen;
 extern llvm::cl::opt<bool> pimDisableMemoryCoalescing;
 extern llvm::cl::opt<bool> useExperimentalConvImpl;
 extern llvm::cl::opt<bool> pimEmitJson;
+extern llvm::cl::opt<bool> pimReportConvLowering;
+extern llvm::cl::opt<bool> pimDetectCommunicationDeadlock;
+extern llvm::cl::opt<bool> pimMaterializeScalarFanoutGlobalOrder;
+extern llvm::cl::opt<bool> pimTraceCommunicationMaterialization;
 
 extern llvm::cl::opt<size_t> crossbarSize;
 extern llvm::cl::opt<size_t> crossbarCountInCore;
 extern llvm::cl::opt<long> coresCount;
+extern llvm::cl::opt<uint64_t> pimConvIm2colMaxElements;
+extern llvm::cl::opt<uint64_t> pimConvStreamChunkPositions;
 
 bool hasExplicitPimCoreCount();
 void verifyExplicitPimCoreCount();

src/PIM/Compiler/PimCompilerUtils.cpp

@@ -29,6 +29,8 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
 
   if (pimEmissionTarget >= EmitSpatial) {
     pm.addPass(createONNXToSpatialPass());
+    pm.addPass(createSpatialLayoutPlanningPass());
+    pm.addPass(createLowerSpatialPlansPass());
     pm.addPass(createMergeComputeNodesPass());
     pm.addPass(createMessagePass("Onnx lowered to Spatial"));
   }

src/PIM/Conversion/CMakeLists.txt

@@ -1,3 +1,2 @@
 add_subdirectory(ONNXToSpatial)
-add_subdirectory(SpatialToGraphviz)
 add_subdirectory(SpatialToPim)

src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt

@@ -10,6 +10,7 @@ add_pim_library(OMONNXToSpatial
   Patterns/Post.cpp
   Patterns/GeneratedConversion.cpp
   Patterns/Math/Conv.cpp
+  Patterns/Math/ConvGeometry.cpp
   Patterns/Math/Elementwise.cpp
   Patterns/Math/Gemm.cpp
   Patterns/Math/MatMul.cpp
@@ -19,6 +20,7 @@ add_pim_library(OMONNXToSpatial
   Patterns/NN/Sigmoid.cpp
   Patterns/NN/Softmax.cpp
   Patterns/Tensor/Concat.cpp
+  Patterns/Tensor/Flatten.cpp
   Patterns/Tensor/Gather.cpp
   Patterns/Tensor/Resize.cpp
   Patterns/Tensor/Reshape.cpp
@@ -26,9 +28,13 @@ add_pim_library(OMONNXToSpatial
   Patterns/Tensor/Split.cpp
   Patterns/Tensor/Transpose.cpp
   ONNXToSpatialPass.cpp
+  SpatialLayoutPlanningPass.cpp
+  LowerSpatialPlansPass.cpp
   Common/AttributeUtils.cpp
+  Common/BiasAddUtils.cpp
   Common/ComputeRegionBuilder.cpp
-  Common/IndexingUtils.cpp
+  Common/MatrixProductLowering.cpp
+  Common/RowStripLayoutUtils.cpp
   Common/ShapeTilingUtils.cpp
   Common/WeightMaterialization.cpp
 

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

@@ -0,0 +1,112 @@
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+LogicalResult isSupportedBiasAddShape(RankedTensorType biasType, RankedTensorType resultType) {
+  if (!biasType || !resultType || !biasType.hasStaticShape() || !resultType.hasStaticShape())
+    return failure();
+  if (resultType.getRank() != 4)
+    return failure();
+  if (biasType.getElementType() != resultType.getElementType())
+    return failure();
+
+  const int64_t channels = resultType.getDimSize(1);
+  ArrayRef<int64_t> shape = biasType.getShape();
+  if (shape.empty())
+    return success();
+  if (shape.size() == 1)
+    return success(shape[0] == channels);
+  if (shape.size() == 2)
+    return success(shape[0] == 1 && shape[1] == channels);
+  if (shape.size() == 4)
+    return success(shape[0] == 1 && shape[1] == channels && shape[2] == 1 && shape[3] == 1);
+  return failure();
+}
+
+FailureOr<SmallVector<Attribute>> getBiasChannelValues(DenseElementsAttr denseAttr, RankedTensorType resultType) {
+  auto biasType = dyn_cast<RankedTensorType>(denseAttr.getType());
+  if (!biasType || failed(isSupportedBiasAddShape(biasType, resultType)))
+    return failure();
+
+  const int64_t channels = resultType.getDimSize(1);
+  if (denseAttr.isSplat()) {
+    return SmallVector<Attribute>(channels, denseAttr.getSplatValue<Attribute>());
+  }
+
+  SmallVector<Attribute> flattened(denseAttr.getValues<Attribute>());
+  if (biasType.getRank() == 1)
+    return flattened;
+  if (biasType.getRank() == 2)
+    return flattened;
+
+  SmallVector<Attribute> channelValues;
+  channelValues.reserve(channels);
+  const int64_t channelStride = biasType.getDimSize(2) * biasType.getDimSize(3);
+  for (int64_t channel = 0; channel < channels; ++channel)
+    channelValues.push_back(flattened[channel * channelStride]);
+  return channelValues;
+}
+
+bool isSupportedBiasAddValue(Value bias, RankedTensorType resultType, DenseElementsAttr* denseAttr) {
+  auto attr = getHostConstDenseElementsAttr(bias);
+  if (!attr)
+    return false;
+  auto biasType = dyn_cast<RankedTensorType>(attr.getType());
+  if (!biasType || failed(isSupportedBiasAddShape(biasType, resultType)))
+    return false;
+  if (failed(getBiasChannelValues(attr, resultType)))
+    return false;
+  if (denseAttr)
+    *denseAttr = attr;
+  return true;
+}
+
+FailureOr<BiasAddPlanCandidate> classifyBiasAddPlanCandidate(Value lhs, Value rhs, RankedTensorType resultType) {
+  auto lhsType = dyn_cast<RankedTensorType>(lhs.getType());
+  auto rhsType = dyn_cast<RankedTensorType>(rhs.getType());
+  if (!lhsType || !rhsType)
+    return failure();
+  if (lhsType == resultType && isSupportedBiasAddValue(rhs, resultType))
+    return BiasAddPlanCandidate {lhs, rhs};
+  if (rhsType == resultType && isSupportedBiasAddValue(lhs, resultType))
+    return BiasAddPlanCandidate {rhs, lhs};
+  return failure();
+}
+
+FailureOr<Value>
+materializeDenseBiasAddTensor(Value bias, RankedTensorType resultType, RewriterBase& rewriter, Location loc) {
+  DenseElementsAttr denseAttr;
+  if (!isSupportedBiasAddValue(bias, resultType, &denseAttr))
+    return failure();
+
+  FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, resultType);
+  if (failed(channelValues))
+    return failure();
+
+  SmallVector<Attribute> resultValues;
+  resultValues.reserve(resultType.getNumElements());
+  const int64_t batches = resultType.getDimSize(0);
+  const int64_t channels = resultType.getDimSize(1);
+  const int64_t height = resultType.getDimSize(2);
+  const int64_t width = resultType.getDimSize(3);
+  for (int64_t n = 0; n < batches; ++n)
+    for (int64_t c = 0; c < channels; ++c)
+      for (int64_t h = 0; h < height; ++h)
+        for (int64_t w = 0; w < width; ++w)
+          resultValues.push_back((*channelValues)[c]);
+
+  auto resultAttr = DenseElementsAttr::get(resultType, resultValues);
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), resultAttr, resultType);
+}
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,30 @@
+#pragma once
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Support/LogicalResult.h"
+
+namespace onnx_mlir {
+
+struct BiasAddPlanCandidate {
+  mlir::Value data;
+  mlir::Value bias;
+};
+
+mlir::LogicalResult isSupportedBiasAddShape(mlir::RankedTensorType biasType, mlir::RankedTensorType resultType);
+bool isSupportedBiasAddValue(mlir::Value bias,
+                             mlir::RankedTensorType resultType,
+                             mlir::DenseElementsAttr* denseAttr = nullptr);
+mlir::FailureOr<llvm::SmallVector<mlir::Attribute>>
+getBiasChannelValues(mlir::DenseElementsAttr denseAttr, mlir::RankedTensorType resultType);
+mlir::FailureOr<BiasAddPlanCandidate> classifyBiasAddPlanCandidate(mlir::Value lhs,
+                                                                   mlir::Value rhs,
+                                                                   mlir::RankedTensorType resultType);
+mlir::FailureOr<mlir::Value> materializeDenseBiasAddTensor(mlir::Value bias,
+                                                           mlir::RankedTensorType resultType,
+                                                           mlir::RewriterBase& rewriter,
+                                                           mlir::Location loc);
+
+} // namespace onnx_mlir

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

@@ -2,8 +2,9 @@
 
 #include "AttributeUtils.hpp"
 #include "ComputeRegionBuilder.hpp"
-#include "IndexingUtils.hpp"
+#include "MatrixProductLowering.hpp"
 #include "ShapeTilingUtils.hpp"
 #include "WeightMaterialization.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"

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

@@ -9,7 +9,7 @@ using namespace mlir;
 
 namespace onnx_mlir {
 
-Value sumTensors(ArrayRef<Value> tensors, ConversionPatternRewriter& rewriter) {
+Value sumTensors(ArrayRef<Value> tensors, PatternRewriter& rewriter) {
   if (tensors.size() == 1)
     return tensors[0];
 

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

@@ -60,6 +60,56 @@ struct SpatComputeBatchBodyArgs {
   mlir::ValueRange outputs;
 };
 
+inline mlir::SmallVector<mlir::Type> getGraphComputeBlockArgTypes(mlir::ValueRange weights, mlir::ValueRange inputs) {
+  mlir::SmallVector<mlir::Type> blockArgTypes;
+  blockArgTypes.reserve(weights.size() + inputs.size());
+  for (mlir::Value weight : weights)
+    blockArgTypes.push_back(weight.getType());
+  for (mlir::Value input : inputs)
+    blockArgTypes.push_back(input.getType());
+  return blockArgTypes;
+}
+
+inline mlir::SmallVector<mlir::Location> getGraphComputeBlockArgLocs(mlir::Location defaultLoc,
+                                                                     mlir::ValueRange weights,
+                                                                     mlir::ValueRange inputs) {
+  mlir::SmallVector<mlir::Location> blockArgLocs;
+  blockArgLocs.reserve(weights.size() + inputs.size());
+  for (mlir::Value weight : weights)
+    blockArgLocs.push_back(weight.getLoc());
+  for (mlir::Value input : inputs)
+    blockArgLocs.push_back(input.getLoc());
+  return blockArgLocs;
+}
+
+inline mlir::SmallVector<mlir::Type> getGraphComputeBatchBlockArgTypes(mlir::OpBuilder& builder,
+                                                                       mlir::TypeRange resultTypes,
+                                                                       mlir::ValueRange weights,
+                                                                       mlir::ValueRange inputs) {
+  mlir::SmallVector<mlir::Type> blockArgTypes {builder.getIndexType()};
+  blockArgTypes.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
+  for (mlir::Value weight : weights)
+    blockArgTypes.push_back(weight.getType());
+  for (mlir::Value input : inputs)
+    blockArgTypes.push_back(input.getType());
+  llvm::append_range(blockArgTypes, resultTypes);
+  return blockArgTypes;
+}
+
+inline mlir::SmallVector<mlir::Location> getGraphComputeBatchBlockArgLocs(mlir::Location defaultLoc,
+                                                                          mlir::TypeRange resultTypes,
+                                                                          mlir::ValueRange weights,
+                                                                          mlir::ValueRange inputs) {
+  mlir::SmallVector<mlir::Location> blockArgLocs {defaultLoc};
+  blockArgLocs.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
+  for (mlir::Value weight : weights)
+    blockArgLocs.push_back(weight.getLoc());
+  for (mlir::Value input : inputs)
+    blockArgLocs.push_back(input.getLoc());
+  blockArgLocs.append(resultTypes.size(), defaultLoc);
+  return blockArgLocs;
+}
+
 } // namespace detail
 
 template <typename RewriterT>
@@ -87,26 +137,43 @@ inline mlir::Value createSpatConcat(RewriterT& rewriter, mlir::Location loc, int
   return spatial::SpatConcatOp::create(rewriter, loc, outputType, rewriter.getI64IntegerAttr(axis), inputs).getOutput();
 }
 
-/// Builds a `spat.compute` with a fixed number of SSA inputs and erases it if
+template <typename RewriterT>
+spatial::SpatGraphCompute createEmptySpatGraphCompute(RewriterT& rewriter,
+                                                      mlir::Location loc,
+                                                      mlir::TypeRange resultTypes,
+                                                      mlir::ValueRange weights,
+                                                      mlir::ValueRange inputs,
+                                                      mlir::TypeRange blockArgTypes,
+                                                      llvm::ArrayRef<mlir::Location> blockArgLocs) {
+  auto computeOp = spatial::SpatGraphCompute::create(rewriter, loc, resultTypes, weights, inputs);
+  rewriter.createBlock(&computeOp.getBody(), computeOp.getBody().end(), blockArgTypes, blockArgLocs);
+  rewriter.setInsertionPointToStart(&computeOp.getBody().front());
+  return computeOp;
+}
+
+template <typename RewriterT>
+spatial::SpatGraphCompute createEmptySpatGraphCompute(RewriterT& rewriter,
+                                                      mlir::Location loc,
+                                                      mlir::TypeRange resultTypes,
+                                                      mlir::ValueRange weights,
+                                                      mlir::ValueRange inputs) {
+  auto blockArgTypes = detail::getGraphComputeBlockArgTypes(weights, inputs);
+  auto blockArgLocs = detail::getGraphComputeBlockArgLocs(loc, weights, inputs);
+  return createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs, blockArgTypes, blockArgLocs);
+}
+
+/// Builds a `spat.graph_compute` with a fixed number of SSA inputs and erases it if
 /// the body callback reports failure.
 template <size_t NumInputs, typename RewriterT, typename BodyFn>
-auto createSpatCompute(RewriterT& rewriter,
+auto createSpatGraphCompute(RewriterT& rewriter,
                             mlir::Location loc,
                             mlir::TypeRange resultTypes,
                             mlir::ValueRange weights,
                             mlir::ValueRange inputs,
                             BodyFn&& body) {
   assert(inputs.size() == NumInputs && "NumInputs must match the number of input values");
-  auto computeOp = spatial::SpatCompute::create(rewriter, loc, resultTypes, weights, inputs);
-
-  auto* block = new mlir::Block();
-  for (mlir::Value weight : weights)
-    block->addArgument(weight.getType(), loc);
-  for (mlir::Value input : inputs)
-    block->addArgument(input.getType(), loc);
-
-  computeOp.getBody().push_back(block);
-  rewriter.setInsertionPointToStart(block);
+  auto computeOp = createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs);
+  auto* block = &computeOp.getBody().front();
 
   using BodyResult = detail::InvokeWithBlockArgsResultT<std::decay_t<BodyFn>, std::make_index_sequence<NumInputs>>;
   if constexpr (std::is_same_v<BodyResult, void>) {
@@ -124,32 +191,24 @@ auto createSpatCompute(RewriterT& rewriter,
     if (mlir::failed(bodyResult)) {
       rewriter.setInsertionPointAfter(computeOp);
       rewriter.eraseOp(computeOp);
-      return mlir::FailureOr<spatial::SpatCompute>(mlir::failure());
+      return mlir::FailureOr<spatial::SpatGraphCompute>(mlir::failure());
     }
     rewriter.setInsertionPointAfter(computeOp);
-    return mlir::FailureOr<spatial::SpatCompute>(computeOp);
+    return mlir::FailureOr<spatial::SpatGraphCompute>(computeOp);
   }
 }
 
-/// Builds a `spat.compute` whose body consumes the block arguments as a single
+/// Builds a `spat.graph_compute` whose body consumes the block arguments as a single
 /// `ValueRange`, which is convenient for variadic reductions/concats.
 template <typename RewriterT, typename BodyFn>
-auto createSpatCompute(RewriterT& rewriter,
+auto createSpatGraphCompute(RewriterT& rewriter,
                             mlir::Location loc,
                             mlir::TypeRange resultTypes,
                             mlir::ValueRange weights,
                             mlir::ValueRange inputs,
                             BodyFn&& body) {
-  auto computeOp = spatial::SpatCompute::create(rewriter, loc, resultTypes, weights, inputs);
-
-  auto* block = new mlir::Block();
-  for (mlir::Value weight : weights)
-    block->addArgument(weight.getType(), loc);
-  for (mlir::Value input : inputs)
-    block->addArgument(input.getType(), loc);
-
-  computeOp.getBody().push_back(block);
-  rewriter.setInsertionPointToStart(block);
+  auto computeOp = createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs);
+  auto* block = &computeOp.getBody().front();
 
   using BodyResult = detail::InvokeWithValueRangeResultT<std::decay_t<BodyFn>>;
   if constexpr (std::is_same_v<BodyResult, void>) {
@@ -163,50 +222,60 @@ auto createSpatCompute(RewriterT& rewriter,
     if (mlir::failed(bodyResult)) {
       rewriter.setInsertionPointAfter(computeOp);
       rewriter.eraseOp(computeOp);
-      return mlir::FailureOr<spatial::SpatCompute>(mlir::failure());
+      return mlir::FailureOr<spatial::SpatGraphCompute>(mlir::failure());
     }
     rewriter.setInsertionPointAfter(computeOp);
-    return mlir::FailureOr<spatial::SpatCompute>(computeOp);
+    return mlir::FailureOr<spatial::SpatGraphCompute>(computeOp);
   }
 }
 
+template <typename RewriterT>
+auto createEmptySpatGraphComputeBatch(RewriterT& rewriter,
+                                      mlir::Location loc,
+                                      mlir::TypeRange resultTypes,
+                                      int64_t laneCount,
+                                      mlir::ValueRange weights,
+                                      mlir::ValueRange inputs,
+                                      mlir::TypeRange blockArgTypes,
+                                      llvm::ArrayRef<mlir::Location> blockArgLocs) {
+  if (laneCount <= 0 || laneCount > std::numeric_limits<int32_t>::max())
+    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
+
+  auto laneCountAttr = pim::getCheckedI32Attr(rewriter, loc, laneCount, "spatial compute_batch lane count");
+  if (mlir::failed(laneCountAttr))
+    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
+
+  auto batchOp = spatial::SpatGraphComputeBatch::create(rewriter, loc, resultTypes, *laneCountAttr, weights, inputs);
+  rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), blockArgTypes, blockArgLocs);
+  rewriter.setInsertionPointToStart(&batchOp.getBody().front());
+  return mlir::FailureOr<spatial::SpatGraphComputeBatch>(batchOp);
+}
+
+template <typename RewriterT>
+auto createEmptySpatGraphComputeBatch(RewriterT& rewriter,
+                                      mlir::Location loc,
+                                      mlir::TypeRange resultTypes,
+                                      int64_t laneCount,
+                                      mlir::ValueRange weights,
+                                      mlir::ValueRange inputs) {
+  auto blockArgTypes = detail::getGraphComputeBatchBlockArgTypes(rewriter, resultTypes, weights, inputs);
+  auto blockArgLocs = detail::getGraphComputeBatchBlockArgLocs(loc, resultTypes, weights, inputs);
+  return createEmptySpatGraphComputeBatch(
+    rewriter, loc, resultTypes, laneCount, weights, inputs, blockArgTypes, blockArgLocs);
+}
+
 template <typename RewriterT, typename BodyFn>
-auto createSpatComputeBatch(RewriterT& rewriter,
+auto createSpatGraphComputeBatch(RewriterT& rewriter,
                                  mlir::Location loc,
                                  mlir::TypeRange resultTypes,
                                  int64_t laneCount,
                                  mlir::ValueRange weights,
                                  mlir::ValueRange inputs,
                                  BodyFn&& body) {
-  if (laneCount <= 0 || laneCount > std::numeric_limits<int32_t>::max())
-    return mlir::FailureOr<spatial::SpatComputeBatch>(mlir::failure());
-
-  auto laneCountAttr = pim::getCheckedI32Attr(rewriter, loc, laneCount, "spatial compute_batch lane count");
-  if (mlir::failed(laneCountAttr))
-    return mlir::FailureOr<spatial::SpatComputeBatch>(mlir::failure());
-
-  auto batchOp = spatial::SpatComputeBatch::create(rewriter, loc, resultTypes, *laneCountAttr, weights, inputs);
-
-  mlir::SmallVector<mlir::Type> blockArgTypes {rewriter.getIndexType()};
-  mlir::SmallVector<mlir::Location> blockArgLocs {loc};
-  blockArgTypes.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
-  blockArgLocs.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
-  for (mlir::Value weight : weights) {
-    blockArgTypes.push_back(weight.getType());
-    blockArgLocs.push_back(weight.getLoc());
-  }
-  for (mlir::Value input : inputs) {
-    blockArgTypes.push_back(input.getType());
-    blockArgLocs.push_back(input.getLoc());
-  }
-  for (mlir::Type resultType : resultTypes) {
-    blockArgTypes.push_back(resultType);
-    blockArgLocs.push_back(loc);
-  }
-
-  auto* block =
-    rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), mlir::TypeRange(blockArgTypes), blockArgLocs);
-  rewriter.setInsertionPointToStart(block);
+  auto batchOp = createEmptySpatGraphComputeBatch(rewriter, loc, resultTypes, laneCount, weights, inputs);
+  if (failed(batchOp))
+    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
+  auto* block = &(*batchOp).getBody().front();
 
   detail::SpatComputeBatchBodyArgs args {
     block->getArgument(0),
@@ -217,21 +286,54 @@ auto createSpatComputeBatch(RewriterT& rewriter,
   using BodyResult = std::invoke_result_t<BodyFn, detail::SpatComputeBatchBodyArgs>;
   if constexpr (std::is_same_v<BodyResult, void>) {
     std::forward<BodyFn>(body)(args);
-    rewriter.setInsertionPointAfter(batchOp);
-    return mlir::FailureOr<spatial::SpatComputeBatch>(batchOp);
+    rewriter.setInsertionPointAfter(*batchOp);
+    return batchOp;
   }
   else {
     auto bodyResult = std::forward<BodyFn>(body)(args);
     if (mlir::failed(bodyResult)) {
-      rewriter.setInsertionPointAfter(batchOp);
-      rewriter.eraseOp(batchOp);
-      return mlir::FailureOr<spatial::SpatComputeBatch>(mlir::failure());
+      rewriter.setInsertionPointAfter(*batchOp);
+      rewriter.eraseOp(*batchOp);
+      return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
     }
-    rewriter.setInsertionPointAfter(batchOp);
-    return mlir::FailureOr<spatial::SpatComputeBatch>(batchOp);
+    rewriter.setInsertionPointAfter(*batchOp);
+    return batchOp;
   }
 }
 
+template <size_t NumInputs, typename RewriterT, typename BodyFn>
+auto createSpatCompute(RewriterT& rewriter,
+                       mlir::Location loc,
+                       mlir::TypeRange resultTypes,
+                       mlir::ValueRange weights,
+                       mlir::ValueRange inputs,
+                       BodyFn&& body) {
+  return createSpatGraphCompute<NumInputs>(
+    rewriter, loc, resultTypes, weights, inputs, std::forward<BodyFn>(body));
+}
+
+template <typename RewriterT, typename BodyFn>
+auto createSpatCompute(RewriterT& rewriter,
+                       mlir::Location loc,
+                       mlir::TypeRange resultTypes,
+                       mlir::ValueRange weights,
+                       mlir::ValueRange inputs,
+                       BodyFn&& body) {
+  return createSpatGraphCompute(rewriter, loc, resultTypes, weights, inputs, std::forward<BodyFn>(body));
+}
+
+template <typename RewriterT, typename BodyFn>
+auto createSpatComputeBatch(RewriterT& rewriter,
+                            mlir::Location loc,
+                            mlir::TypeRange resultTypes,
+                            int64_t laneCount,
+                            mlir::ValueRange weights,
+                            mlir::ValueRange inputs,
+                            BodyFn&& body) {
+  return createSpatGraphComputeBatch(
+    rewriter, loc, resultTypes, laneCount, weights, inputs, std::forward<BodyFn>(body));
+}
+
 inline void createParallelInsertSliceIntoBatchOutput(mlir::PatternRewriter& rewriter,
                                                      mlir::Location loc,
                                                      mlir::Value source,
@@ -262,6 +364,6 @@ mlir::Value materializeOrComputeUnary(mlir::Value input,
   return computeOp.getResult(0);
 }
 
-mlir::Value sumTensors(mlir::ArrayRef<mlir::Value> tensors, mlir::ConversionPatternRewriter& rewriter);
+mlir::Value sumTensors(mlir::ArrayRef<mlir::Value> tensors, mlir::PatternRewriter& rewriter);
 
 } // namespace onnx_mlir

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

@@ -0,0 +1,48 @@
+#include "MatrixProductLowering.hpp"
+
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+Value createZeroPaddedTensor(Value value, RankedTensorType resultType, PatternRewriter& rewriter, Location loc) {
+  auto sourceType = cast<RankedTensorType>(value.getType());
+  SmallVector<OpFoldResult> lowPads(sourceType.getRank(), rewriter.getIndexAttr(0));
+  SmallVector<OpFoldResult> highPads;
+  highPads.reserve(sourceType.getRank());
+  for (auto [sourceDim, resultDim] : llvm::zip(sourceType.getShape(), resultType.getShape()))
+    highPads.push_back(rewriter.getIndexAttr(resultDim - sourceDim));
+
+  auto padOp = tensor::PadOp::create(rewriter, loc, resultType, value, lowPads, highPads);
+  auto* padBlock = new Block();
+  for (int64_t i = 0; i < sourceType.getRank(); ++i)
+    padBlock->addArgument(rewriter.getIndexType(), loc);
+  padOp.getRegion().push_back(padBlock);
+  rewriter.setInsertionPointToStart(padBlock);
+  auto zero = getOrCreateConstant(
+    rewriter, padOp.getOperation(), rewriter.getZeroAttr(sourceType.getElementType()), sourceType.getElementType());
+  tensor::YieldOp::create(rewriter, loc, zero);
+  rewriter.setInsertionPointAfter(padOp);
+  return padOp.getResult();
+}
+
+Value createPaddedInputCompute(Value input,
+                               RankedTensorType paddedInputType,
+                               PatternRewriter& rewriter,
+                               Location loc) {
+  auto inputType = cast<RankedTensorType>(input.getType());
+  if (inputType == paddedInputType)
+    return input;
+
+  auto computeOp = createSpatCompute<1>(rewriter, loc, TypeRange {paddedInputType}, {}, input, [&](Value computeInput) {
+    Value paddedInput = createZeroPaddedTensor(computeInput, paddedInputType, rewriter, loc);
+    spatial::SpatYieldOp::create(rewriter, loc, paddedInput);
+  });
+  return computeOp.getResult(0);
+}
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,20 @@
+#pragma once
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Location.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+namespace onnx_mlir {
+
+mlir::Value createZeroPaddedTensor(mlir::Value value,
+                                   mlir::RankedTensorType resultType,
+                                   mlir::PatternRewriter& rewriter,
+                                   mlir::Location loc);
+
+mlir::Value createPaddedInputCompute(mlir::Value input,
+                                     mlir::RankedTensorType paddedInputType,
+                                     mlir::PatternRewriter& rewriter,
+                                     mlir::Location loc);
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,239 @@
+#include "llvm/ADT/SmallVector.h"
+
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) {
+  return RankedTensorType::get({logicalType.getDimSize(0), logicalType.getDimSize(1), 1, logicalType.getDimSize(3)},
+                               logicalType.getElementType(),
+                               logicalType.getEncoding());
+}
+
+RankedTensorType getRowStripStorageType(RankedTensorType logicalType) {
+  return RankedTensorType::get({logicalType.getDimSize(2), logicalType.getDimSize(1), 1, logicalType.getDimSize(3)},
+                               logicalType.getElementType(),
+                               logicalType.getEncoding());
+}
+
+std::pair<SmallVector<int64_t>, SmallVector<int64_t>> buildRowStripMetadata(RankedTensorType type) {
+  SmallVector<int64_t> offsets;
+  SmallVector<int64_t> sizes;
+  const int64_t channels = type.getDimSize(1);
+  const int64_t height = type.getDimSize(2);
+  const int64_t width = type.getDimSize(3);
+  offsets.reserve(height * 4);
+  sizes.reserve(height * 4);
+  for (int64_t row = 0; row < height; ++row) {
+    offsets.append({0, 0, row, 0});
+    sizes.append({1, channels, 1, width});
+  }
+  return {offsets, sizes};
+}
+
+SmallVector<OpFoldResult> buildRowStripFragmentOffsets(PatternRewriter& rewriter, OpFoldResult row) {
+  return {row, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+}
+
+SmallVector<OpFoldResult> buildRowStripFragmentSizes(PatternRewriter& rewriter, RankedTensorType logicalType) {
+  return {rewriter.getIndexAttr(1),
+          rewriter.getIndexAttr(logicalType.getDimSize(1)),
+          rewriter.getIndexAttr(1),
+          rewriter.getIndexAttr(logicalType.getDimSize(3))};
+}
+
+Value extractRowStripFragment(Value storage,
+                              RankedTensorType logicalType,
+                              OpFoldResult row,
+                              PatternRewriter& rewriter,
+                              Location loc) {
+  return tensor::ExtractSliceOp::create(rewriter,
+                                        loc,
+                                        getRowStripFragmentType(logicalType),
+                                        storage,
+                                        buildRowStripFragmentOffsets(rewriter, row),
+                                        buildRowStripFragmentSizes(rewriter, logicalType),
+                                        getUnitStrides(rewriter, 4));
+}
+
+void insertRowStripFragment(Value fragment,
+                            Value output,
+                            RankedTensorType logicalType,
+                            OpFoldResult row,
+                            PatternRewriter& rewriter,
+                            Location loc) {
+  createParallelInsertSliceIntoBatchOutput(rewriter,
+                                           loc,
+                                           fragment,
+                                           output,
+                                           buildRowStripFragmentOffsets(rewriter, row),
+                                           buildRowStripFragmentSizes(rewriter, logicalType),
+                                           getUnitStrides(rewriter, 4));
+}
+
+FailureOr<Value> createPerChannelConstantFragment(DenseElementsAttr denseAttr,
+                                                  RankedTensorType fragmentType,
+                                                  PatternRewriter& rewriter) {
+  FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, fragmentType);
+  if (failed(channelValues))
+    return failure();
+
+  SmallVector<Attribute> values;
+  values.reserve(fragmentType.getNumElements());
+  for (int64_t n = 0; n < fragmentType.getDimSize(0); ++n)
+    for (int64_t channel = 0; channel < fragmentType.getDimSize(1); ++channel)
+      for (int64_t h = 0; h < fragmentType.getDimSize(2); ++h)
+        for (int64_t w = 0; w < fragmentType.getDimSize(3); ++w)
+          values.push_back((*channelValues)[channel]);
+
+  auto attr = DenseElementsAttr::get(fragmentType, values);
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), attr, fragmentType);
+}
+
+FailureOr<Value> createRowStripStorageFromRows(Value rows,
+                                               RankedTensorType logicalType,
+                                               PatternRewriter& rewriter,
+                                               Location loc) {
+  auto rowsType = dyn_cast<RankedTensorType>(rows.getType());
+  if (!rowsType || !rowsType.hasStaticShape() || rowsType.getRank() != 2)
+    return failure();
+  if (!logicalType || !logicalType.hasStaticShape() || logicalType.getRank() != 4)
+    return failure();
+  if (logicalType.getDimSize(0) != 1)
+    return failure();
+  if (rowsType.getElementType() != logicalType.getElementType())
+    return failure();
+
+  const int64_t channels = logicalType.getDimSize(1);
+  const int64_t height = logicalType.getDimSize(2);
+  const int64_t width = logicalType.getDimSize(3);
+  if (rowsType.getDimSize(0) != height * width)
+    return failure();
+  if (rowsType.getDimSize(1) != channels)
+    return failure();
+
+  auto rowSliceType = RankedTensorType::get({width, channels}, logicalType.getElementType(), rowsType.getEncoding());
+  auto channelWidthType = RankedTensorType::get({channels, width}, logicalType.getElementType(), rowsType.getEncoding());
+  auto fragmentType = getRowStripFragmentType(logicalType);
+  auto storageType = getRowStripStorageType(logicalType);
+  auto batchOp = createSpatComputeBatch(
+    rewriter, loc, TypeRange {storageType}, height, {}, ValueRange {rows}, [&](detail::SpatComputeBatchBodyArgs args) {
+      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+      Value rowStart = affineMulConst(rewriter, loc, args.lane, width, anchorOp);
+      SmallVector<OpFoldResult> rowOffsets {rowStart, rewriter.getIndexAttr(0)};
+      SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(width), rewriter.getIndexAttr(channels)};
+      Value rowSlice = tensor::ExtractSliceOp::create(
+        rewriter, loc, rowSliceType, args.inputs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
+      Value channelWidth = ONNXTransposeOp::create(
+        rewriter, loc, channelWidthType, rowSlice, rewriter.getI64ArrayAttr({1, 0})).getResult();
+      Value fragment = tensor::ExpandShapeOp::create(
+        rewriter, loc, fragmentType, channelWidth, SmallVector<ReassociationIndices> {{0, 1}, {2, 3}});
+      insertRowStripFragment(fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
+      return success();
+    });
+  if (failed(batchOp))
+    return failure();
+  return batchOp->getResult(0);
+}
+
+FailureOr<Value>
+materializeRowStripStorageToDense(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
+  auto storageType = dyn_cast<RankedTensorType>(storage.getType());
+  if (!storageType || storageType != getRowStripStorageType(logicalType))
+    return failure();
+
+  auto batchOp = createSpatComputeBatch(
+    rewriter, loc, TypeRange {logicalType}, logicalType.getDimSize(2), {}, ValueRange {storage},
+    [&](detail::SpatComputeBatchBodyArgs args) {
+      Value fragment = extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
+      createParallelInsertSliceIntoBatchOutput(rewriter,
+                                               loc,
+                                               fragment,
+                                               args.outputs.front(),
+                                               SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
+                                                                          rewriter.getIndexAttr(0),
+                                                                          args.lane,
+                                                                          rewriter.getIndexAttr(0)},
+                                               buildRowStripFragmentSizes(rewriter, logicalType),
+                                               getUnitStrides(rewriter, 4));
+      return success();
+    });
+  if (failed(batchOp))
+    return failure();
+  return batchOp->getResult(0);
+}
+
+FailureOr<Value>
+applyRowStripRelu(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
+  auto fragmentType = getRowStripFragmentType(logicalType);
+  auto storageType = getRowStripStorageType(logicalType);
+  auto batchOp = createSpatComputeBatch(rewriter,
+                                        loc,
+                                        TypeRange {storageType},
+                                        logicalType.getDimSize(2),
+                                        {},
+                                        ValueRange {storage},
+                                        [&](detail::SpatComputeBatchBodyArgs args) {
+                                          Value fragment =
+                                            extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
+                                          fragment = spatial::SpatReluOp::create(rewriter, loc, fragmentType, fragment).getResult();
+                                          insertRowStripFragment(
+                                            fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
+                                          return success();
+                                        });
+  if (failed(batchOp))
+    return failure();
+  return batchOp->getResult(0);
+}
+
+FailureOr<Value>
+applyRowStripBiasAdd(Value storage, RankedTensorType logicalType, Value bias, PatternRewriter& rewriter, Location loc) {
+  DenseElementsAttr denseAttr;
+  if (!isSupportedBiasAddValue(bias, logicalType, &denseAttr))
+    return failure();
+  auto fragmentType = getRowStripFragmentType(logicalType);
+  auto storageType = getRowStripStorageType(logicalType);
+  auto batchOp = createSpatComputeBatch(rewriter,
+                                        loc,
+                                        TypeRange {storageType},
+                                        logicalType.getDimSize(2),
+                                        {},
+                                        ValueRange {storage},
+                                        [&](detail::SpatComputeBatchBodyArgs args) {
+                                          Value fragment =
+                                            extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
+                                          Value constant;
+                                          if (denseAttr.isSplat()) {
+                                            constant = getOrCreateConstant(
+                                              rewriter,
+                                              rewriter.getInsertionBlock()->getParentOp(),
+                                              DenseElementsAttr::get(fragmentType, denseAttr.getSplatValue<Attribute>()),
+                                              fragmentType);
+                                          }
+                                          else {
+                                            FailureOr<Value> perChannel =
+                                              createPerChannelConstantFragment(denseAttr, fragmentType, rewriter);
+                                            if (failed(perChannel))
+                                              return failure();
+                                            constant = *perChannel;
+                                          }
+                                          fragment =
+                                            spatial::SpatVAddOp::create(rewriter, loc, fragmentType, fragment, constant).getResult();
+                                          insertRowStripFragment(
+                                            fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
+                                          return success();
+                                        });
+  if (failed(batchOp))
+    return failure();
+  return batchOp->getResult(0);
+}
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,69 @@
+#pragma once
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/PatternMatch.h"
+
+namespace onnx_mlir {
+
+inline constexpr llvm::StringLiteral kRowStripIndexMap = "nchw_row_strip_fragments";
+
+struct RowStripPhysicalValue {
+  mlir::Value storage;
+  mlir::RankedTensorType logicalType;
+  llvm::SmallVector<int64_t, 16> fragmentOffsets;
+  llvm::SmallVector<int64_t, 16> fragmentSizes;
+};
+
+std::pair<llvm::SmallVector<int64_t>, llvm::SmallVector<int64_t>>
+buildRowStripMetadata(mlir::RankedTensorType type);
+
+mlir::RankedTensorType getRowStripFragmentType(mlir::RankedTensorType logicalType);
+
+mlir::RankedTensorType getRowStripStorageType(mlir::RankedTensorType logicalType);
+
+llvm::SmallVector<mlir::OpFoldResult> buildRowStripFragmentOffsets(mlir::PatternRewriter& rewriter,
+                                                                   mlir::OpFoldResult row);
+
+llvm::SmallVector<mlir::OpFoldResult> buildRowStripFragmentSizes(mlir::PatternRewriter& rewriter,
+                                                                 mlir::RankedTensorType logicalType);
+
+mlir::Value extractRowStripFragment(mlir::Value storage,
+                                    mlir::RankedTensorType logicalType,
+                                    mlir::OpFoldResult row,
+                                    mlir::PatternRewriter& rewriter,
+                                    mlir::Location loc);
+
+void insertRowStripFragment(mlir::Value fragment,
+                            mlir::Value output,
+                            mlir::RankedTensorType logicalType,
+                            mlir::OpFoldResult row,
+                            mlir::PatternRewriter& rewriter,
+                            mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> createPerChannelConstantFragment(mlir::DenseElementsAttr denseAttr,
+                                                              mlir::RankedTensorType fragmentType,
+                                                              mlir::PatternRewriter& rewriter);
+
+mlir::FailureOr<mlir::Value> createRowStripStorageFromRows(mlir::Value rows,
+                                                           mlir::RankedTensorType logicalType,
+                                                           mlir::PatternRewriter& rewriter,
+                                                           mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> materializeRowStripStorageToDense(mlir::Value storage,
+                                                               mlir::RankedTensorType logicalType,
+                                                               mlir::PatternRewriter& rewriter,
+                                                               mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> applyRowStripRelu(mlir::Value storage,
+                                               mlir::RankedTensorType logicalType,
+                                               mlir::PatternRewriter& rewriter,
+                                               mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> applyRowStripBiasAdd(mlir::Value storage,
+                                                  mlir::RankedTensorType logicalType,
+                                                  mlir::Value bias,
+                                                  mlir::PatternRewriter& rewriter,
+                                                  mlir::Location loc);
+
+} // namespace onnx_mlir

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

@@ -3,9 +3,6 @@
 
 #include "llvm/ADT/SmallVector.h"
 
-#include <functional>
-
-#include "IndexingUtils.hpp"
 #include "ShapeTilingUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
@@ -15,75 +12,8 @@ using namespace mlir;
 
 namespace onnx_mlir {
 
-bool hasStaticPositiveShape(ArrayRef<int64_t> shape) {
-  return llvm::all_of(shape, [](int64_t dim) { return dim > 0; });
-}
-
-bool hasStaticPositiveShape(RankedTensorType type) {
-  return type.hasStaticShape() && hasStaticPositiveShape(type.getShape());
-}
-
-int64_t getStaticShapeElementCount(ArrayRef<int64_t> shape) {
-  return std::accumulate(shape.begin(), shape.end(), int64_t {1}, std::multiplies<int64_t> {});
-}
-
-SmallVector<int64_t> permuteShape(ArrayRef<int64_t> shape, ArrayRef<int64_t> permutation) {
-  SmallVector<int64_t> permutedShape;
-  permutedShape.reserve(permutation.size());
-  for (int64_t axis : permutation)
-    permutedShape.push_back(shape[axis]);
-  return permutedShape;
-}
-
-SmallVector<int64_t> invertPermutation(ArrayRef<int64_t> permutation) {
-  SmallVector<int64_t> inversePermutation(permutation.size());
-  for (auto [newIndex, oldIndex] : llvm::enumerate(permutation))
-    inversePermutation[oldIndex] = static_cast<int64_t>(newIndex);
-  return inversePermutation;
-}
-
-FailureOr<SmallVector<int64_t>> getTransposePermutationChecked(std::optional<ArrayAttr> permAttr, int64_t rank) {
-  SmallVector<int64_t> permutation;
-  if (!permAttr) {
-    permutation.reserve(rank);
-    for (int64_t dim = rank - 1; dim >= 0; --dim)
-      permutation.push_back(dim);
-    return permutation;
-  }
-
-  if (static_cast<int64_t>(permAttr->size()) != rank)
-    return failure();
-
-  permutation.reserve(permAttr->size());
-  SmallVector<bool> seen(rank, false);
-  for (IntegerAttr attr : permAttr->getAsRange<IntegerAttr>()) {
-    int64_t axis = attr.getInt();
-    if (axis < 0 || axis >= rank || seen[axis])
-      return failure();
-    seen[axis] = true;
-    permutation.push_back(axis);
-  }
-  return permutation;
-}
-
-SmallVector<OpFoldResult> getUnitStrides(PatternRewriter& rewriter, int64_t rank) {
-  return SmallVector<OpFoldResult>(rank, rewriter.getIndexAttr(1));
-}
-
-SmallVector<OpFoldResult> getZeroOffsets(PatternRewriter& rewriter, int64_t rank) {
-  return SmallVector<OpFoldResult>(rank, rewriter.getIndexAttr(0));
-}
-
-SmallVector<OpFoldResult> getStaticSizes(PatternRewriter& rewriter, ArrayRef<int64_t> shape) {
-  SmallVector<OpFoldResult> sizes;
-  sizes.reserve(shape.size());
-  for (int64_t dim : shape)
-    sizes.push_back(rewriter.getIndexAttr(dim));
-  return sizes;
-}
-
 SmallVector<Value> sliceTensor(
-  const Value& tensorToSlice, size_t axis, int64_t sliceSize, ConversionPatternRewriter& rewriter, Location loc) {
+  const Value& tensorToSlice, size_t axis, int64_t sliceSize, PatternRewriter& rewriter, Location loc) {
   ArrayRef<long> shape = getTensorShape(tensorToSlice);
   assert("Invalid axis" && axis < shape.size());
 
@@ -129,7 +59,7 @@ SmallVector<Value> sliceTensor(
 }
 
 SmallVector<Value>
-sliceVector(const Value& vectorToSlice, int64_t sliceSize, ConversionPatternRewriter& rewriter, Location loc) {
+sliceVector(const Value& vectorToSlice, int64_t sliceSize, PatternRewriter& rewriter, Location loc) {
   ArrayRef<long> shape = getTensorShape(vectorToSlice);
   assert("Not a vector" && isVectorShape(shape));
   size_t axis = shape[0] != 1 ? 0 : 1;
@@ -137,7 +67,7 @@ sliceVector(const Value& vectorToSlice, int64_t sliceSize, ConversionPatternRewr
 }
 
 DenseMap<CoreId, SmallVector<Value>>
-sliceVectorPerCrossbarPerCore(const Value& vectorToSlice, ConversionPatternRewriter& rewriter, Location loc) {
+sliceVectorPerCrossbarPerCore(const Value& vectorToSlice, PatternRewriter& rewriter, Location loc) {
   SmallVector<Value> slices = sliceVector(vectorToSlice, crossbarSize, rewriter, loc);
   DenseMap<CoreId, SmallVector<Value>> slicesPerCore;
   for (size_t sliceId = 0; sliceId < slices.size(); sliceId++) {
@@ -147,33 +77,4 @@ sliceVectorPerCrossbarPerCore(const Value& vectorToSlice, ConversionPatternRewri
   return slicesPerCore;
 }
 
-Value extractAxisSlice(
-  PatternRewriter& rewriter, Location loc, Value source, int64_t axis, int64_t offset, int64_t size) {
-  auto sourceType = cast<RankedTensorType>(source.getType());
-  SmallVector<int64_t> resultShape(sourceType.getShape());
-  resultShape[axis] = size;
-  auto resultType = RankedTensorType::get(resultShape, sourceType.getElementType(), sourceType.getEncoding());
-
-  SmallVector<OpFoldResult> offsets = getZeroOffsets(rewriter, sourceType.getRank());
-  SmallVector<OpFoldResult> sizes = getStaticSizes(rewriter, sourceType.getShape());
-  offsets[axis] = rewriter.getIndexAttr(offset);
-  sizes[axis] = rewriter.getIndexAttr(size);
-  return tensor::ExtractSliceOp::create(
-           rewriter, loc, resultType, source, offsets, sizes, getUnitStrides(rewriter, sourceType.getRank()))
-    .getResult();
-}
-
-Value insertStaticSlice(
-  PatternRewriter& rewriter, Location loc, Value source, Value dest, ArrayRef<OpFoldResult> offsets) {
-  auto sourceType = cast<RankedTensorType>(source.getType());
-  return tensor::InsertSliceOp::create(rewriter,
-                                       loc,
-                                       source,
-                                       dest,
-                                       offsets,
-                                       getStaticSizes(rewriter, sourceType.getShape()),
-                                       getUnitStrides(rewriter, sourceType.getRank()))
-    .getResult();
-}
-
 } // namespace onnx_mlir

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

@@ -1,114 +1,31 @@
 #pragma once
 
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/IR/BuiltinTypes.h"
-#include "mlir/IR/Value.h"
 #include "mlir/IR/ValueRange.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 <optional>
-#include <type_traits>
-#include <utility>
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 
 namespace onnx_mlir {
 
-using HSliceId = size_t;
-using CoreId = size_t;
-
-template <class A, class B, class C = std::common_type_t<A, B>>
-constexpr C ceilIntegerDivide(A a, B b) {
-  static_assert(std::is_integral_v<A>, "A must be an integer type");
-  static_assert(std::is_integral_v<B>, "B must be an integer type");
-  C ac = static_cast<C>(a);
-  C bc = static_cast<C>(b);
-  return 1 + (ac - 1) / bc;
-}
-
-template <class A, class B, class C = std::common_type_t<A, B>>
-constexpr std::pair<C, C> ceilIntegerDivideWithRemainder(A a, B b) {
-  static_assert(std::is_integral_v<A>, "A must be an integer type");
-  static_assert(std::is_integral_v<B>, "B must be an integer type");
-  C ac = static_cast<C>(a);
-  C bc = static_cast<C>(b);
-  return {ceilIntegerDivide(ac, bc), ac % bc};
-}
-
-template <class T>
-bool isVectorShape(mlir::ArrayRef<T> shape) {
-  return shape.size() == 2 && (shape[0] == 1 || shape[1] == 1);
-}
-
-template <class T>
-bool isMatrixShape(mlir::ArrayRef<T> shape) {
-  return shape.size() == 2;
-}
-
-template <class T>
-bool isHVectorShape(mlir::ArrayRef<T> shape) {
-  return shape.size() == 2 && shape[0] == 1;
-}
-
-inline auto getTensorShape(mlir::Value tensor) {
-  return mlir::cast<mlir::RankedTensorType>(tensor.getType()).getShape();
-}
-
-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();
-}
-
-bool hasStaticPositiveShape(mlir::ArrayRef<int64_t> shape);
-
-bool hasStaticPositiveShape(mlir::RankedTensorType type);
-
-int64_t getStaticShapeElementCount(mlir::ArrayRef<int64_t> shape);
-
-llvm::SmallVector<int64_t> permuteShape(mlir::ArrayRef<int64_t> shape, mlir::ArrayRef<int64_t> permutation);
-
-llvm::SmallVector<int64_t> invertPermutation(mlir::ArrayRef<int64_t> permutation);
-
-mlir::FailureOr<llvm::SmallVector<int64_t>> getTransposePermutationChecked(std::optional<mlir::ArrayAttr> permAttr,
-                                                                           int64_t rank);
-
-llvm::SmallVector<mlir::OpFoldResult> getUnitStrides(mlir::PatternRewriter& rewriter, int64_t rank);
-
-llvm::SmallVector<mlir::OpFoldResult> getZeroOffsets(mlir::PatternRewriter& rewriter, int64_t rank);
-
-llvm::SmallVector<mlir::OpFoldResult> getStaticSizes(mlir::PatternRewriter& rewriter, mlir::ArrayRef<int64_t> shape);
-
 /// Slices a statically shaped tensor along one axis into contiguous pieces of
 /// at most `sliceSize` elements.
 llvm::SmallVector<mlir::Value> sliceTensor(const mlir::Value& tensorToSlice,
                                            size_t axis,
                                            int64_t sliceSize,
-                                           mlir::ConversionPatternRewriter& rewriter,
+                                           mlir::PatternRewriter& rewriter,
                                            mlir::Location loc);
 
 llvm::SmallVector<mlir::Value> sliceVector(const mlir::Value& vectorToSlice,
                                            int64_t sliceSize,
-                                           mlir::ConversionPatternRewriter& rewriter,
+                                           mlir::PatternRewriter& rewriter,
                                            mlir::Location loc);
 
 /// Partitions one logical vector into per-core crossbar-sized slices using the
 /// current PIM target geometry.
 llvm::DenseMap<CoreId, llvm::SmallVector<mlir::Value>> sliceVectorPerCrossbarPerCore(
-  const mlir::Value& vectorToSlice, mlir::ConversionPatternRewriter& rewriter, mlir::Location loc);
-
-mlir::Value extractAxisSlice(
-  mlir::PatternRewriter& rewriter, mlir::Location loc, mlir::Value source, int64_t axis, int64_t offset, int64_t size);
-
-mlir::Value insertStaticSlice(mlir::PatternRewriter& rewriter,
-                              mlir::Location loc,
-                              mlir::Value source,
-                              mlir::Value dest,
-                              llvm::ArrayRef<mlir::OpFoldResult> offsets);
+  const mlir::Value& vectorToSlice, mlir::PatternRewriter& rewriter, mlir::Location loc);
 
 } // namespace onnx_mlir

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

@@ -19,9 +19,11 @@ using namespace mlir;
 
 namespace onnx_mlir {
 
-bool isWeightLikeComputeOperand(Value value) {
+static bool isWeightMaterializationValue(Value value, bool requireMatrixShape) {
   auto rankedType = dyn_cast<RankedTensorType>(value.getType());
-  if (!rankedType || !isMatrixShape(rankedType.getShape()))
+  if (!rankedType)
+    return false;
+  if (requireMatrixShape && !isMatrixShape(rankedType.getShape()))
     return false;
 
   llvm::SmallPtrSet<Operation*, 8> visited;
@@ -29,8 +31,14 @@ bool isWeightLikeComputeOperand(Value value) {
   while (auto* definingOp = value.getDefiningOp()) {
     if (!visited.insert(definingOp).second)
       return false;
-    if (isa<arith::ConstantOp, ONNXConstantOp>(definingOp) || hasWeightAlways(definingOp))
+    if (isa<arith::ConstantOp, ONNXConstantOp>(definingOp) || hasWeightAlways(definingOp)) {
+      auto sourceType = dyn_cast<RankedTensorType>(value.getType());
+      if (!sourceType)
+        return false;
+      if (requireMatrixShape && !isMatrixShape(sourceType.getShape()))
+        return false;
       return true;
+    }
 
     if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(definingOp)) {
       value = extractSliceOp.getSource();
@@ -55,6 +63,8 @@ bool isWeightLikeComputeOperand(Value value) {
   return false;
 }
 
+bool isWeightLikeComputeOperand(Value value) { return isWeightMaterializationValue(value, /*requireMatrixShape=*/true); }
+
 FailureOr<Value> materializeWeightLikeValueInBlock(Value value, IRRewriter& rewriter, IRMapping& mapper) {
   if (auto mapped = mapper.lookupOrNull(value))
     return cast<Value>(mapped);
@@ -91,7 +101,7 @@ FailureOr<Value> materializeWeightLikeValueInBlock(Value value, IRRewriter& rewr
       continue;
     }
 
-    if (isWeightLikeComputeOperand(operand)) {
+    if (isWeightMaterializationValue(operand, /*requireMatrixShape=*/false)) {
       auto clonedOperand = materializeWeightLikeValueInBlock(operand, rewriter, mapper);
       if (failed(clonedOperand))
         return failure();

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -0,0 +1,400 @@
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallPtrSet.h"
+
+#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
+#include "mlir/Transforms/Passes.h"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Pass/PIMPasses.h"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+static constexpr StringLiteral kDenseLayout = "dense_nchw";
+static constexpr StringLiteral kRowStripLayout = "nchw_row_strip";
+
+static FailureOr<RowStripPhysicalValue> getRowStripValue(llvm::DenseMap<Value, RowStripPhysicalValue>& rowStripValues,
+                                                         Value value) {
+  auto it = rowStripValues.find(value);
+  if (it == rowStripValues.end())
+    return failure();
+  return it->second;
+}
+
+static FailureOr<RowStripPhysicalValue> buildRowStripValue(spatial::SpatBlueprintOp blueprint,
+                                                           Value storage) {
+  auto logicalType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
+  if (!logicalType)
+    return blueprint.emitOpError("requires ranked logical output type"), failure();
+  RowStripPhysicalValue value;
+  value.storage = storage;
+  value.logicalType = logicalType;
+  value.fragmentOffsets.append(blueprint.getFragmentOffsets().begin(), blueprint.getFragmentOffsets().end());
+  value.fragmentSizes.append(blueprint.getFragmentSizes().begin(), blueprint.getFragmentSizes().end());
+  if (blueprint.getIndexMap() != kRowStripIndexMap)
+    return blueprint.emitOpError("requires the canonical row-strip index map"), failure();
+  auto storageType = dyn_cast<RankedTensorType>(storage.getType());
+  if (!storageType || storageType != getRowStripStorageType(logicalType))
+    return blueprint.emitOpError("requires physical row-strip fragment storage"), failure();
+  return value;
+}
+
+static FailureOr<Value>
+lowerRowStripRelu(const RowStripPhysicalValue& input, spatial::SpatReluPlanOp planOp, PatternRewriter& rewriter) {
+  return applyRowStripRelu(input.storage, input.logicalType, rewriter, planOp.getLoc());
+}
+
+static FailureOr<Value> lowerRowStripBiasAdd(const RowStripPhysicalValue& input,
+                                             spatial::SpatBiasAddPlanOp planOp,
+                                             PatternRewriter& rewriter) {
+  return applyRowStripBiasAdd(input.storage, input.logicalType, planOp.getBias(), rewriter, planOp.getLoc());
+}
+
+static FailureOr<Value>
+materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) {
+  if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape())
+    return failure();
+  auto [expectedOffsets, expectedSizes] = buildRowStripMetadata(rowStripValue.logicalType);
+  if (!llvm::equal(rowStripValue.fragmentOffsets, expectedOffsets) || !llvm::equal(rowStripValue.fragmentSizes, expectedSizes))
+    return failure();
+  return materializeRowStripStorageToDense(rowStripValue.storage, rowStripValue.logicalType, rewriter, loc);
+}
+
+struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerSpatialPlansPass)
+
+  StringRef getArgument() const override { return "lower-spatial-plans"; }
+  StringRef getDescription() const override { return "Lower selected Spatial planning ops to low-level Spatial IR."; }
+
+  void runOnOperation() override {
+    ModuleOp moduleOp = getOperation();
+    MLIRContext* ctx = moduleOp.getContext();
+    auto entryFunc = getPimEntryFunc(moduleOp);
+    if (failed(entryFunc)) {
+      moduleOp.emitError("failed to locate the PIM entry function during LowerSpatialPlans");
+      signalPassFailure();
+      return;
+    }
+    func::FuncOp funcOp = *entryFunc;
+    PatternRewriter rewriter(ctx);
+    llvm::DenseMap<Value, RowStripPhysicalValue> rowStripValues;
+    llvm::SmallPtrSet<Operation*, 16> eraseAfterLowering;
+    auto verifyLogicalPhase = [&](StringRef stage) -> bool {
+      if (succeeded(verifyLogicalSpatialGraphInvariants(*entryFunc)))
+        return true;
+      moduleOp.emitError() << "logical Spatial graph verification failed " << stage;
+      signalPassFailure();
+      return false;
+    };
+
+    if (!verifyLogicalPhase("at the start of LowerSpatialPlans"))
+      return;
+    for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
+      if (auto planOp = dyn_cast<spatial::SpatConv2DPlanOp>(&op)) {
+        FailureOr<RowStripPhysicalValue> rowStripInput = getRowStripValue(rowStripValues, planOp.getInput());
+        auto rowStripBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
+          auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
+          return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
+        });
+        if (rowStripBlueprint != planOp.getResult().getUsers().end()) {
+          rewriter.setInsertionPoint(planOp);
+          FailureOr<Value> lowered = lowerSelectedConv2DPlan(
+            planOp,
+            succeeded(rowStripInput) ? std::optional<Value> {rowStripInput->storage} : std::nullopt,
+            /*emitRowStripLayout=*/true,
+            rewriter);
+          if (failed(lowered)) {
+            planOp.emitOpError("failed to lower selected row-strip Spatial Conv plan");
+            signalPassFailure();
+            return;
+          }
+          auto blueprint = cast<spatial::SpatBlueprintOp>(*rowStripBlueprint);
+          FailureOr<RowStripPhysicalValue> rowStripValue = buildRowStripValue(blueprint, *lowered);
+          if (failed(rowStripValue)) {
+            signalPassFailure();
+            return;
+          }
+          rowStripValues[blueprint.getResult()] = *rowStripValue;
+          eraseAfterLowering.insert(planOp);
+          eraseAfterLowering.insert(blueprint);
+          continue;
+        }
+        rewriter.setInsertionPoint(planOp);
+        FailureOr<Value> lowered =
+          lowerSelectedConv2DPlan(planOp, std::nullopt, /*emitRowStripLayout=*/false, rewriter);
+        if (failed(lowered)) {
+          planOp.emitOpError("failed to lower selected Spatial Conv plan");
+          signalPassFailure();
+          return;
+        }
+        rewriter.replaceOp(planOp, *lowered);
+        continue;
+      }
+
+      if (auto planOp = dyn_cast<spatial::SpatReluPlanOp>(&op)) {
+        if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) {
+          auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
+            auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
+            return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
+          });
+          if (outputBlueprint == planOp.getResult().getUsers().end()) {
+            planOp.emitOpError("row-strip Relu plan requires a row-strip blueprint result");
+            signalPassFailure();
+            return;
+          }
+
+          FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
+          rewriter.setInsertionPoint(planOp);
+          FailureOr<Value> lowered = lowerRowStripRelu(*input, planOp, rewriter);
+          if (failed(lowered)) {
+            planOp.emitOpError("failed to lower selected row-strip Spatial Relu plan");
+            signalPassFailure();
+            return;
+          }
+          auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
+          FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
+          if (failed(output)) {
+            signalPassFailure();
+            return;
+          }
+          rowStripValues[blueprint.getResult()] = *output;
+          eraseAfterLowering.insert(planOp);
+          eraseAfterLowering.insert(blueprint);
+          continue;
+        }
+
+        rewriter.setInsertionPoint(planOp);
+        auto computeOp = createSpatCompute<1>(
+          rewriter, planOp.getLoc(), planOp.getOutput().getType(), {}, planOp.getInput(), [&](Value x) {
+            auto relu = spatial::SpatReluOp::create(rewriter, planOp.getLoc(), planOp.getOutput().getType(), x);
+            spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), relu.getResult());
+          });
+        rewriter.replaceOp(planOp, computeOp.getResults());
+        continue;
+      }
+      if (auto planOp = dyn_cast<spatial::SpatBiasAddPlanOp>(&op)) {
+        if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) {
+          auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
+            auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
+            return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
+          });
+          if (outputBlueprint == planOp.getResult().getUsers().end()) {
+            planOp.emitOpError("row-strip bias_add plan requires a row-strip blueprint result");
+            signalPassFailure();
+            return;
+          }
+
+          FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
+          rewriter.setInsertionPoint(planOp);
+          FailureOr<Value> lowered = lowerRowStripBiasAdd(*input, planOp, rewriter);
+          if (failed(lowered)) {
+            planOp.emitOpError("failed to lower selected row-strip Spatial bias_add plan");
+            signalPassFailure();
+            return;
+          }
+          auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
+          FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
+          if (failed(output)) {
+            signalPassFailure();
+            return;
+          }
+          rowStripValues[blueprint.getResult()] = *output;
+          eraseAfterLowering.insert(planOp);
+          eraseAfterLowering.insert(blueprint);
+          continue;
+        }
+
+        auto resultType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+        if (!resultType) {
+          planOp.emitOpError("requires ranked output type");
+          signalPassFailure();
+          return;
+        }
+        rewriter.setInsertionPoint(planOp);
+        FailureOr<Value> denseBias = materializeDenseBiasAddTensor(planOp.getBias(), resultType, rewriter, planOp.getLoc());
+        if (failed(denseBias)) {
+          planOp.emitOpError("failed to materialize dense Conv-style bias");
+          signalPassFailure();
+          return;
+        }
+        auto computeOp = createSpatCompute<2>(rewriter,
+                                              planOp.getLoc(),
+                                              planOp.getOutput().getType(),
+                                              {},
+                                              ValueRange {planOp.getInput(), *denseBias},
+                                              [&](Value x, Value y) {
+                                                auto added = spatial::SpatVAddOp::create(
+                                                  rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, y);
+                                                spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added.getResult());
+                                              });
+        rewriter.replaceOp(planOp, computeOp.getResults());
+        continue;
+      }
+      if (auto materializeOp = dyn_cast<spatial::SpatMaterializeLayoutOp>(&op)) {
+        if (materializeOp.getSourcePhysicalLayout() == kDenseLayout
+            && materializeOp.getTargetPhysicalLayout() == kDenseLayout) {
+          rewriter.replaceOp(materializeOp, materializeOp.getInput());
+          continue;
+        }
+        if (materializeOp.getSourcePhysicalLayout() != kRowStripLayout
+            || materializeOp.getTargetPhysicalLayout() != kDenseLayout) {
+          materializeOp.emitOpError("non-dense materialize_layout lowering is not supported yet");
+          signalPassFailure();
+          return;
+        }
+        FailureOr<RowStripPhysicalValue> rowStripValue = getRowStripValue(rowStripValues, materializeOp.getInput());
+        if (failed(rowStripValue)) {
+          materializeOp.emitOpError("expected a row-strip blueprint input during row-strip materialization");
+          signalPassFailure();
+          return;
+        }
+        rewriter.setInsertionPoint(materializeOp);
+        FailureOr<Value> dense = materializeRowStripToDense(*rowStripValue, materializeOp.getLoc(), rewriter);
+        if (failed(dense)) {
+          materializeOp.emitOpError("failed to materialize selected row-strip layout back to dense NCHW");
+          signalPassFailure();
+          return;
+        }
+        rewriter.replaceOp(materializeOp, *dense);
+        continue;
+      }
+      if (auto blueprintOp = dyn_cast<spatial::SpatBlueprintOp>(&op)) {
+        if (blueprintOp.getPhysicalLayout() == kDenseLayout) {
+          rewriter.replaceOp(blueprintOp, blueprintOp.getInput());
+          continue;
+        }
+        if (blueprintOp.getPhysicalLayout() != kRowStripLayout) {
+          blueprintOp.emitOpError("non-dense blueprint lowering is not supported yet");
+          signalPassFailure();
+          return;
+        }
+        if (!eraseAfterLowering.contains(blueprintOp)) {
+          blueprintOp.emitOpError("unhandled row-strip blueprint remained during LowerSpatialPlans");
+          signalPassFailure();
+          return;
+        }
+      }
+    }
+    bool erasedAny = true;
+    while (erasedAny) {
+      erasedAny = false;
+      for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
+        if (!eraseAfterLowering.contains(&op))
+          continue;
+        if (!op.use_empty())
+          continue;
+        eraseAfterLowering.erase(&op);
+        rewriter.eraseOp(&op);
+        erasedAny = true;
+      }
+    }
+    if (!eraseAfterLowering.empty()) {
+      for (Operation& op : funcOp.getBody().front())
+        if (eraseAfterLowering.contains(&op))
+          op.emitOpError("selected row-strip planning op could not be fully eliminated during LowerSpatialPlans");
+      signalPassFailure();
+      return;
+    }
+    ConversionTarget helperTarget(*ctx);
+    helperTarget.addLegalDialect<spatial::SpatialDialect,
+                                 tensor::TensorDialect,
+                                 linalg::LinalgDialect,
+                                 affine::AffineDialect,
+                                 arith::ArithDialect,
+                                 scf::SCFDialect,
+                                 func::FuncDialect>();
+    helperTarget.addLegalOp<spatial::SpatGraphCompute, spatial::SpatGraphComputeBatch>();
+    helperTarget.addIllegalOp<ONNXGemmOp, ONNXTransposeOp>();
+    helperTarget.markOpRecursivelyLegal<spatial::SpatGraphCompute, spatial::SpatGraphComputeBatch>();
+
+    RewritePatternSet helperPatterns(ctx);
+    populateGemmPatterns(helperPatterns, ctx);
+    populateTransposePatterns(helperPatterns, ctx);
+    if (failed(applyPartialConversion(moduleOp, helperTarget, std::move(helperPatterns)))) {
+      moduleOp.emitError("failed to lower helper ONNX ops emitted by selected Spatial plan lowering");
+      signalPassFailure();
+      return;
+    }
+    FrozenRewritePatternSet nestedHelperPatterns([&] {
+      RewritePatternSet patterns(ctx);
+      populateGemmPatterns(patterns, ctx);
+      populateTransposePatterns(patterns, ctx);
+      return patterns;
+    }());
+    ConversionTarget nestedHelperTarget(*ctx);
+    nestedHelperTarget.addLegalDialect<spatial::SpatialDialect,
+                                       tensor::TensorDialect,
+                                       linalg::LinalgDialect,
+                                       affine::AffineDialect,
+                                       arith::ArithDialect,
+                                       scf::SCFDialect,
+                                       func::FuncDialect>();
+    nestedHelperTarget.addIllegalOp<ONNXGemmOp, ONNXTransposeOp>();
+    SmallVector<Operation*> computeLikeOps;
+    funcOp.walk([&](Operation* op) {
+      if (isa<spatial::SpatGraphCompute, spatial::SpatGraphComputeBatch>(op))
+        computeLikeOps.push_back(op);
+    });
+    for (Operation* op : computeLikeOps) {
+      if (failed(applyFullConversion(op, nestedHelperTarget, nestedHelperPatterns))) {
+        op->emitOpError("failed to lower nested helper ONNX ops emitted by selected Spatial plan lowering");
+        signalPassFailure();
+        return;
+      }
+    }
+    if (!verifyLogicalPhase("after nested helper conversions"))
+      return;
+    bool hasIllegalOps = false;
+    moduleOp.walk([&](Operation* op) {
+      if (isa<ONNXEntryPointOp>(op))
+        return;
+      if (isa<spatial::SpatConv2DPlanOp,
+              spatial::SpatBiasAddPlanOp,
+              spatial::SpatReluPlanOp,
+              spatial::SpatBlueprintOp,
+              spatial::SpatMaterializeLayoutOp>(op)
+          || op->getDialect()->getNamespace() == "onnx") {
+        op->emitOpError("operation must not remain after LowerSpatialPlans");
+        hasIllegalOps = true;
+      }
+    });
+
+    PassManager canonicalizationPM(ctx);
+    canonicalizationPM.addPass(createCanonicalizerPass());
+    if (failed(canonicalizationPM.run(moduleOp)))
+      moduleOp.emitWarning("failed to run LowerSpatialPlansPass canonicalization; continuing");
+
+    if (hasIllegalOps) {
+      signalPassFailure();
+    } else {
+      dumpModule(moduleOp, "spatial1_graph");
+    }
+
+    if (!verifyLogicalPhase("at the end of LowerSpatialPlans"))
+      return;
+  }
+};
+
+} // namespace
+
+std::unique_ptr<Pass> createLowerSpatialPlansPass() { return std::make_unique<LowerSpatialPlansPass>(); }
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -13,11 +13,13 @@
 
 #include "Common/Common.hpp"
 #include "Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
+#include "ONNXToSpatialVerifier.hpp"
 
 using namespace mlir;
 
@@ -41,10 +43,17 @@ struct ONNXToSpatialPass : PassWrapper<ONNXToSpatialPass, OperationPass<ModuleOp
 static void populateEmptyFunction(func::FuncOp funcOp) {
   IRRewriter rewriter(funcOp.getContext());
   IRMapping mapper;
-  SmallVector<spatial::SpatCompute> computes(funcOp.getOps<spatial::SpatCompute>());
-  SmallVector<spatial::SpatComputeBatch> computeBatches(funcOp.getOps<spatial::SpatComputeBatch>());
-  if (!computes.empty() || !computeBatches.empty())
+  SmallVector<spatial::SpatGraphCompute> computes(funcOp.getOps<spatial::SpatGraphCompute>());
+  SmallVector<spatial::SpatGraphComputeBatch> computeBatches(funcOp.getOps<spatial::SpatGraphComputeBatch>());
+  SmallVector<spatial::SpatConv2DPlanOp> convPlans(funcOp.getOps<spatial::SpatConv2DPlanOp>());
+  SmallVector<spatial::SpatBiasAddPlanOp> biasAddPlans(funcOp.getOps<spatial::SpatBiasAddPlanOp>());
+  SmallVector<spatial::SpatReluPlanOp> reluPlans(funcOp.getOps<spatial::SpatReluPlanOp>());
+  SmallVector<spatial::SpatBlueprintOp> blueprints(funcOp.getOps<spatial::SpatBlueprintOp>());
+  SmallVector<spatial::SpatMaterializeLayoutOp> materializers(funcOp.getOps<spatial::SpatMaterializeLayoutOp>());
+  if (!computes.empty() || !computeBatches.empty() || !convPlans.empty() || !biasAddPlans.empty() || !reluPlans.empty()
+      || !blueprints.empty() || !materializers.empty()) {
     return;
+  }
 
   auto returnOp = cast<func::ReturnOp>(funcOp.getFunctionBody().front().getTerminator());
   rewriter.setInsertionPoint(returnOp);
@@ -58,16 +67,16 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
     sourceLocs.push_back(source.getLoc());
   }
 
-  auto newCompute = spatial::SpatCompute::create(
-    rewriter, returnOp.getLoc(), returnOp.getOperandTypes(), funcOp.getArguments(), {}, {});
-  auto* newBlock = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), sourceTypes, sourceLocs);
+  auto newCompute = createEmptySpatGraphCompute(
+    rewriter, returnOp.getLoc(), returnOp.getOperandTypes(), {}, funcOp.getArguments(), sourceTypes, sourceLocs);
+  auto* newBlock = &newCompute.getBody().front();
   for (auto [blockArg, computeArg] : llvm::zip(newBlock->getArguments(), newCompute.getOperands()))
     mapper.map(computeArg, blockArg);
   newCompute.getProperties().setOperandSegmentSizes({0, static_cast<int>(sourceTypes.size())});
 
   rewriter.setInsertionPointToEnd(newBlock);
   for (Operation& op : funcOp.getOps())
-    if (!isa<spatial::SpatCompute, func::ReturnOp>(&op))
+    if (!isa<spatial::SpatGraphCompute, func::ReturnOp>(&op))
       rewriter.clone(op, mapper);
 
   auto yield = spatial::SpatYieldOp::create(rewriter, funcOp.getLoc(), returnOp.getOperands());
@@ -75,7 +84,7 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
     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)) {
+    if (!isa<spatial::SpatGraphCompute, func::ReturnOp>(&op)) {
       op.dropAllUses();
       rewriter.eraseOp(&op);
     }
@@ -96,7 +105,7 @@ void ONNXToSpatialPass::runOnOperation() {
                             affine::AffineDialect,
                             arith::ArithDialect,
                             scf::SCFDialect>();
-  preTarget.addIllegalOp<ONNXConstantOp, ONNXFlattenOp>();
+  preTarget.addIllegalOp<ONNXConstantOp>();
 
   RewritePatternSet prePatterns(ctx);
   populatePrePatterns(prePatterns, ctx);
@@ -135,6 +144,7 @@ void ONNXToSpatialPass::runOnOperation() {
   target.addIllegalOp<ONNXSigmoidOp>();
   target.addIllegalOp<ONNXSoftmaxOp>();
   target.addIllegalOp<ONNXConcatOp>();
+  target.addIllegalOp<ONNXFlattenOp>();
   target.addIllegalOp<ONNXGatherOp>();
   target.addIllegalOp<ONNXReshapeOp>();
   target.addIllegalOp<ONNXResizeOp>();
@@ -152,6 +162,11 @@ void ONNXToSpatialPass::runOnOperation() {
     return;
   }
 
+    if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
+      moduleOp.emitError("logical Spatial graph verification failed after ONNX conversion");
+      signalPassFailure();
+      return;
+    }
   ConversionTarget earlyPostTarget(*ctx);
   earlyPostTarget.addLegalDialect<spatial::SpatialDialect,
                                   ONNXDialect,
@@ -161,13 +176,13 @@ void ONNXToSpatialPass::runOnOperation() {
                                   arith::ArithDialect,
                                   scf::SCFDialect>();
 
-  PassManager cleanupPM(ctx);
-  cleanupPM.addPass(createCanonicalizerPass());
-  if (failed(cleanupPM.run(moduleOp)))
-    moduleOp.emitWarning("failed to run ONNX-to-Spatial canonicalization cleanup; continuing");
-
   annotateWeightsConstants(*entryFunc);
 
+    if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
+      moduleOp.emitError("logical Spatial graph verification failed after weight annotation");
+      signalPassFailure();
+      return;
+    }
   ConversionTarget postTarget(*ctx);
   postTarget.addLegalDialect<spatial::SpatialDialect,
                              ONNXDialect,
@@ -176,11 +191,16 @@ void ONNXToSpatialPass::runOnOperation() {
                              affine::AffineDialect,
                              arith::ArithDialect,
                              scf::SCFDialect>();
-  postTarget.addDynamicallyLegalOp<spatial::SpatCompute>(
-    [](spatial::SpatCompute computeOp) { return !requiresPostRewrite(computeOp); });
-  postTarget.addDynamicallyLegalOp<spatial::SpatComputeBatch>(
-    [](spatial::SpatComputeBatch computeOp) { return !requiresPostRewrite(computeOp); });
+  postTarget.addDynamicallyLegalOp<spatial::SpatGraphCompute>(
+    [](spatial::SpatGraphCompute computeOp) { return !requiresPostRewrite(computeOp); });
+  postTarget.addDynamicallyLegalOp<spatial::SpatGraphComputeBatch>(
+    [](spatial::SpatGraphComputeBatch computeOp) { return !requiresPostRewrite(computeOp); });
 
+    if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
+      moduleOp.emitError("logical Spatial graph verification failed before post rewrites");
+      signalPassFailure();
+      return;
+    }
   RewritePatternSet postPatterns(ctx);
   populatePostPatterns(postPatterns, ctx);
   if (failed(applyPartialConversion(*entryFunc, postTarget, std::move(postPatterns)))) {
@@ -191,8 +211,18 @@ void ONNXToSpatialPass::runOnOperation() {
 
   populateEmptyFunction(*entryFunc);
 
+  PassManager canonicalizationPM(ctx);
+  canonicalizationPM.addPass(createCanonicalizerPass());
+  if (failed(canonicalizationPM.run(moduleOp)))
+    moduleOp.emitWarning("failed to run ONNXToSpatial canonicalization; continuing");
+
   dumpModule(moduleOp, "spatial0");
 
+  if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
+      moduleOp.emitError("logical Spatial graph verification failed after ONNX-to-Spatial");
+      signalPassFailure();
+      return;
+    }
   if (failed(verifyONNXToSpatial(*entryFunc))) {
     moduleOp.emitError("ONNX-to-Spatial host legality verification failed");
     signalPassFailure();

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp

@@ -1,4 +1,6 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/Diagnostics.h"
 #include "mlir/Support/LLVM.h"
 
 #include "Common/IR/WeightUtils.hpp"
@@ -13,6 +15,8 @@ namespace onnx_mlir {
 
 namespace {
 
+constexpr StringLiteral kPhaseMarker = "phase-check";
+
 void checkWeightUseChains(func::FuncOp func, pim::CappedDiagnosticReporter& diagnostics) {
   func.walk([&](Operation* op) {
     if (!hasWeightAlways(op))
@@ -23,134 +27,205 @@ void checkWeightUseChains(func::FuncOp func, pim::CappedDiagnosticReporter& diag
         continue;
 
       diagnostics.report(op, [&](Operation* illegalOp) {
-        illegalOp->emitOpError(
-          "weight-marked values may only flow through static view/slice helper chains into Spatial VMM weights");
+        illegalOp->emitOpError()
+          << kPhaseMarker
+          << " weight-marked values may only flow through static view/slice helper chains into Spatial VMM weights";
       });
       return;
     }
   });
 }
 
-Region* getParentRegion(Value value) {
-  if (auto blockArg = dyn_cast<BlockArgument>(value))
-    return blockArg.getOwner()->getParent();
-  if (Operation* definingOp = value.getDefiningOp())
-    return definingOp->getParentRegion();
-  return nullptr;
+bool isRegionOrAncestorOf(Region& region, Region* candidate) {
+  return candidate && (&region == candidate || region.isAncestor(candidate));
 }
 
-bool isDefinedInsideRegion(Value value, Region& region) {
-  Region* parentRegion = getParentRegion(value);
-  return parentRegion && (&region == parentRegion || region.isAncestor(parentRegion));
+bool isValueDefinedInsideRegion(Value value, Region& region) {
+  if (auto blockArg = dyn_cast<BlockArgument>(value))
+    return isRegionOrAncestorOf(region, blockArg.getOwner()->getParent());
+  if (Operation* definingOp = value.getDefiningOp())
+    return isRegionOrAncestorOf(region, definingOp->getParentRegion());
+  return false;
+}
+
+bool isLegalExternalCapture(Value value, Region& region) {
+  if (isValueDefinedInsideRegion(value, region))
+    return true;
+
+  Operation* definingOp = value.getDefiningOp();
+  return definingOp && definingOp->hasTrait<OpTrait::ConstantLike>();
+}
+
+bool isRecordedDeferredCommunicationSource(Operation* op, Value value) {
+  auto transfer = dyn_cast<spatial::SpatDeferredCommunicationOp>(op);
+  return transfer && llvm::is_contained(transfer.getSources(), value);
+}
+
+template <typename ComputeOpTy>
+void verifyComputeBodyCaptures(ComputeOpTy compute, StringRef kind, pim::CappedDiagnosticReporter& diagnostics) {
+  Region& body = compute.getBody();
+  body.walk([&](Operation* nestedOp) {
+    for (OpOperand& operand : nestedOp->getOpOperands()) {
+      Value value = operand.get();
+      if (isLegalExternalCapture(value, body) || isRecordedDeferredCommunicationSource(nestedOp, value))
+        continue;
+
+      Operation* definingOp = value.getDefiningOp();
+      diagnostics.report(compute.getOperation(), [&](Operation* illegalOp) {
+        InFlightDiagnostic diag =
+          illegalOp->emitOpError() << kPhaseMarker << " " << kind << " body captures non-constant external operand #"
+                                   << operand.getOperandNumber() << " used by " << nestedOp->getName().getStringRef();
+        diag << " (type " << value.getType() << ")";
+        if (definingOp)
+          diag.attachNote(definingOp->getLoc()) << "defining op is " << definingOp->getName().getStringRef();
+        else if (auto blockArg = dyn_cast<BlockArgument>(value)) {
+          if (Operation* owner = blockArg.getOwner()->getParentOp())
+            diag.attachNote(owner->getLoc())
+              << "external block argument belongs to " << owner->getName().getStringRef();
+        }
+      });
+    }
+  });
 }
 
 bool isLegalHostBackedValue(Value value) {
   Operation* definingOp = value.getDefiningOp();
   if (!definingOp)
     return isa<BlockArgument>(value);
-
-  if (isa<spatial::SpatChannelReceiveOp>(definingOp))
-    return false;
-
   return definingOp->getDialect()->getNamespace() != "spat";
 }
 
-LogicalResult verifyComputeLikeInputs(Operation* computeLikeOp,
-                                      ValueRange inputs,
+bool isScheduledPhase1Value(Value value) {
+  Operation* definingOp = value.getDefiningOp();
+  return isa_and_nonnull<spatial::SpatScheduledCompute, spatial::SpatScheduledComputeBatch>(definingOp);
+}
+
+template <typename ComputeOpTy>
+void verifyScheduledInputs(ComputeOpTy compute,
                            bool allowChannelReceiveInputs,
                            StringRef kind,
                            pim::CappedDiagnosticReporter& diagnostics) {
-  for (auto [inputIndex, input] : llvm::enumerate(inputs)) {
-    unsigned currentInputIndex = inputIndex;
+  for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs())) {
+    size_t currentInputIndex = inputIndex;
     Operation* definingOp = input.getDefiningOp();
     if (allowChannelReceiveInputs && isa_and_nonnull<spatial::SpatChannelReceiveOp>(definingOp))
       continue;
+    if (isScheduledPhase1Value(input))
+      continue;
     if (isLegalHostBackedValue(input))
       continue;
 
-    diagnostics.report(computeLikeOp, [&](Operation* illegalOp) {
-      InFlightDiagnostic diagnostic = illegalOp->emitOpError()
-                                   << kind << " input #" << currentInputIndex
-                                   << (allowChannelReceiveInputs ? " must come from the host or an explicit "
-                                                                   "spat.channel_receive"
+    diagnostics.report(compute.getOperation(), [&](Operation* illegalOp) {
+      InFlightDiagnostic diag = illegalOp->emitOpError()
+                                << kPhaseMarker << " " << kind << " input #" << currentInputIndex
+                                << (allowChannelReceiveInputs ? " must come from the host or explicit spat.channel_receive"
                                                               : " must come from the host");
       if (definingOp)
-        diagnostic.attachNote(definingOp->getLoc()) << "illegal Spatial producer is " << definingOp->getName();
+        diag.attachNote(definingOp->getLoc()) << "illegal producer is " << definingOp->getName().getStringRef();
     });
-    return failure();
   }
-  return success();
 }
 
-void verifyNoExternalTensorCaptures(Operation* ownerOp,
-                                    Region& region,
-                                    StringRef kind,
+template <typename ComputeOpTy>
+void verifyNoNestedFragmentAssemblyBlueprints(ComputeOpTy compute,
                                                   pim::CappedDiagnosticReporter& diagnostics) {
-  region.walk([&](Operation* op) {
-    for (OpOperand& operand : op->getOpOperands()) {
-      Value value = operand.get();
-      if (!isa<TensorType>(value.getType()))
-        continue;
-      if (isDefinedInsideRegion(value, region) || isa<BlockArgument>(value))
-        continue;
-
-      Operation* definingOp = value.getDefiningOp();
-      if (definingOp && definingOp->hasTrait<OpTrait::ConstantLike>())
-        continue;
-
-      diagnostics.report(ownerOp, [&](Operation* illegalOp) {
-        InFlightDiagnostic diagnostic = illegalOp->emitOpError() << kind << " body may not capture external tensor "
-                                                                 << "values";
-        diagnostic.attachNote(op->getLoc())
-          << "tensor operand #" << operand.getOperandNumber() << " is defined outside the compute body by "
-          << (definingOp ? definingOp->getName().getStringRef() : StringRef("<block argument>"));
+  compute.getBody().walk([&](spatial::SpatBlueprintOp blueprint) {
+    std::optional<StringRef> mode = blueprint.getMode();
+    if (!mode || *mode != "fragment_assembly")
+      return;
+    diagnostics.report(blueprint.getOperation(), [&](Operation* illegalOp) {
+      illegalOp->emitOpError("fragment assembly blueprint must be host-level after merge materialization");
+    });
   });
 }
+
+void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter& diagnostics) {
+  for (Operation& op : funcOp.getOps()) {
+    if (isa<func::ReturnOp,
+            spatial::SpatGraphCompute,
+            spatial::SpatGraphComputeBatch,
+            spatial::SpatConv2DPlanOp,
+            spatial::SpatBiasAddPlanOp,
+            spatial::SpatReluPlanOp,
+            spatial::SpatBlueprintOp,
+            spatial::SpatMaterializeLayoutOp>(&op)) {
+      continue;
+    }
+    if (isa<spatial::SpatScheduledCompute, spatial::SpatScheduledComputeBatch>(&op)) {
+      diagnostics.report(&op, [&](Operation* illegalOp) {
+        illegalOp->emitOpError() << kPhaseMarker << " scheduled Spatial compute op is not allowed in logical graph phase";
       });
+      continue;
+    }
+    if (isa<spatial::SpatChannelReceiveOp, spatial::SpatChannelSendOp>(&op)) {
+      diagnostics.report(&op, [&](Operation* illegalOp) {
+        illegalOp->emitOpError() << kPhaseMarker
+                                 << " explicit channel communication is not expected before merge materialization";
+      });
+      continue;
+    }
+    if (isCompileTimeOp(&op))
+      continue;
+
+    diagnostics.report(&op, [&](Operation* illegalOp) {
+      illegalOp->emitOpError()
+        << kPhaseMarker << " non-foldable top-level runtime op remains in logical Spatial graph; lower it inside spat.graph_compute";
+    });
+  }
+}
+
+void verifyScheduledTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter& diagnostics) {
+  for (Operation& op : funcOp.getOps()) {
+    if (isa<spatial::SpatChannelSendOp, spatial::SpatChannelReceiveOp>(&op)) {
+      diagnostics.report(&op, [&](Operation* illegalOp) {
+        illegalOp->emitOpError() << kPhaseMarker << " real channel communication is not allowed in scheduled phase 1";
+      });
+    }
+  }
 }
 
 } // namespace
 
-LogicalResult verifyONNXToSpatial(func::FuncOp funcOp) {
+LogicalResult verifyNoComputeBodyCaptures(func::FuncOp funcOp) {
   pim::CappedDiagnosticReporter diagnostics;
-
-  for (Operation& op : funcOp.getOps()) {
-    if (isa<func::ReturnOp, spatial::SpatCompute, spatial::SpatComputeBatch>(&op))
-      continue;
-    if (isCompileTimeOp(&op))
-      continue;
-
-    diagnostics.report(&op, [](Operation* illegalOp) {
-      illegalOp->emitOpError(
-        "non-foldable top-level runtime op remains after ONNX-to-Spatial; lower it inside spat.compute");
-    });
-  }
-  checkWeightUseChains(funcOp, diagnostics);
-  diagnostics.emitSuppressedSummary(funcOp, "ONNX-to-Spatial verification failed");
-
+  for (auto compute : funcOp.getOps<spatial::SpatGraphCompute>())
+    verifyComputeBodyCaptures(compute, "graph_compute", diagnostics);
+  for (auto batch : funcOp.getOps<spatial::SpatGraphComputeBatch>())
+    verifyComputeBodyCaptures(batch, "graph_compute_batch", diagnostics);
+  for (auto compute : funcOp.getOps<spatial::SpatScheduledCompute>())
+    verifyComputeBodyCaptures(compute, "scheduled_compute", diagnostics);
+  for (auto batch : funcOp.getOps<spatial::SpatScheduledComputeBatch>())
+    verifyComputeBodyCaptures(batch, "scheduled_compute_batch", diagnostics);
+  diagnostics.emitSuppressedSummary(funcOp, "compute body capture verification failed");
   return success(!diagnostics.hasFailure());
 }
 
-LogicalResult verifySpatialCommunicationInvariants(func::FuncOp funcOp) {
+LogicalResult verifyONNXToSpatial(func::FuncOp funcOp) { return verifyLogicalSpatialGraphInvariants(funcOp); }
+
+LogicalResult verifyLogicalSpatialGraphInvariants(func::FuncOp funcOp) {
   pim::CappedDiagnosticReporter diagnostics;
-
-  for (auto computeOp : funcOp.getOps<spatial::SpatCompute>()) {
-    (void) verifyComputeLikeInputs(
-      computeOp.getOperation(), computeOp.getInputs(), /*allowChannelReceiveInputs=*/true, "spat.compute", diagnostics);
-    verifyNoExternalTensorCaptures(computeOp.getOperation(), computeOp.getBody(), "spat.compute", diagnostics);
+  verifyLogicalTopLevelOps(funcOp, diagnostics);
+  checkWeightUseChains(funcOp, diagnostics);
+  if (failed(verifyNoComputeBodyCaptures(funcOp)))
+    return failure();
+  diagnostics.emitSuppressedSummary(funcOp, "logical Spatial graph verification failed");
+  return success(!diagnostics.hasFailure());
 }
 
-  for (auto computeBatchOp : funcOp.getOps<spatial::SpatComputeBatch>()) {
-    (void) verifyComputeLikeInputs(computeBatchOp.getOperation(),
-                                   computeBatchOp.getInputs(),
-                                   /*allowChannelReceiveInputs=*/false,
-                                   "spat.compute_batch",
-                                   diagnostics);
-    verifyNoExternalTensorCaptures(
-      computeBatchOp.getOperation(), computeBatchOp.getBody(), "spat.compute_batch", diagnostics);
+LogicalResult verifyScheduledSpatialInvariants(func::FuncOp funcOp) {
+  pim::CappedDiagnosticReporter diagnostics;
+  verifyScheduledTopLevelOps(funcOp, diagnostics);
+  for (auto compute : funcOp.getOps<spatial::SpatScheduledCompute>()) {
+    verifyScheduledInputs(compute, /*allowChannelReceiveInputs=*/true, "spat.scheduled_compute", diagnostics);
+    verifyNoNestedFragmentAssemblyBlueprints(compute, diagnostics);
   }
-
-  diagnostics.emitSuppressedSummary(funcOp, "Spatial communication invariant verification failed");
+  for (auto batch : funcOp.getOps<spatial::SpatScheduledComputeBatch>()) {
+    verifyScheduledInputs(batch, /*allowChannelReceiveInputs=*/false, "spat.scheduled_compute_batch", diagnostics);
+    verifyNoNestedFragmentAssemblyBlueprints(batch, diagnostics);
+  }
+  if (failed(verifyNoComputeBodyCaptures(funcOp)))
+    return failure();
+  diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial verification failed");
   return success(!diagnostics.hasFailure());
 }
 

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp

@@ -6,6 +6,8 @@
 namespace onnx_mlir {
 
 mlir::LogicalResult verifyONNXToSpatial(mlir::func::FuncOp funcOp);
-mlir::LogicalResult verifySpatialCommunicationInvariants(mlir::func::FuncOp funcOp);
+mlir::LogicalResult verifyNoComputeBodyCaptures(mlir::func::FuncOp funcOp);
+mlir::LogicalResult verifyLogicalSpatialGraphInvariants(mlir::func::FuncOp funcOp);
+mlir::LogicalResult verifyScheduledSpatialInvariants(mlir::func::FuncOp funcOp);
 
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Patterns.cpp

@@ -19,6 +19,7 @@ void populateConversionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
   populateSigmoidPatterns(patterns, ctx);
   populateSoftmaxPatterns(patterns, ctx);
   populateConcatPatterns(patterns, ctx);
+  populateFlattenPatterns(patterns, ctx);
   populateGatherPatterns(patterns, ctx);
   populateResizePatterns(patterns, ctx);
   populateReshapePatterns(patterns, ctx);

src/PIM/Conversion/ONNXToSpatial/Patterns.hpp

@@ -26,6 +26,7 @@ void populateReluPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext*
 void populateSigmoidPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateSoftmaxPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateConcatPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populateFlattenPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateGatherPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateResizePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateReshapePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
@@ -33,8 +34,8 @@ void populateSlicePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext*
 void populateSplitPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateTransposePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 
-bool requiresPostRewrite(spatial::SpatCompute computeOp);
-bool requiresPostRewrite(spatial::SpatComputeBatch computeOp);
+bool requiresPostRewrite(spatial::SpatGraphCompute computeOp);
+bool requiresPostRewrite(spatial::SpatGraphComputeBatch computeOp);
 void annotateWeightsConstants(mlir::func::FuncOp funcOp);
 
 } // namespace onnx_mlir

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

@@ -0,0 +1,77 @@
+#include "ConvGeometry.hpp"
+
+#include <algorithm>
+
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+
+namespace onnx_mlir {
+
+bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOut, int64_t numChannelsInPerGroup) {
+  return group == numChannelsIn && numChannelsInPerGroup == 1 && numChannelsOut % group == 0;
+}
+
+ConvGeometry buildConvGeometry(const ConvLoweringState& state) {
+  ConvGeometry geo {
+    state.batchSize,
+    state.numChannelsIn,
+    state.xHeight,
+    state.xWidth,
+    state.numChannelsOut,
+    state.wHeight,
+    state.wWidth,
+    state.outHeight,
+    state.outWidth,
+    state.group,
+    state.numChannelsInPerGroup,
+    state.numChannelsOutPerGroup,
+    state.numChannelsInPerGroup * state.wHeight * state.wWidth,
+    state.numChannelsOutPerGroup,
+    state.batchSize * state.outHeight * state.outWidth,
+    static_cast<int64_t>(crossbarSize.getValue()),
+    1,
+    0,
+    state.hasBias,
+    isDepthwiseConv(state.group, state.numChannelsIn, state.numChannelsOut, state.numChannelsInPerGroup),
+  };
+  geo.pack = std::max<int64_t>(1, geo.xbarSize / std::max<int64_t>(geo.k, geo.c));
+  geo.im2colElements = static_cast<uint64_t>(std::max<int64_t>(0, geo.p)) * static_cast<uint64_t>(std::max<int64_t>(0, geo.k));
+  return geo;
+}
+
+uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, int64_t packFactor) {
+  const uint64_t patchElements = static_cast<uint64_t>(std::max<int64_t>(1, geo.k));
+  uint64_t chunkPositions = std::max<uint64_t>(1, pimConvIm2colMaxElements / patchElements);
+  chunkPositions = std::min<uint64_t>(chunkPositions, static_cast<uint64_t>(std::max<int64_t>(1, geo.p)));
+  chunkPositions = std::min<uint64_t>(chunkPositions, std::max<uint64_t>(1, pimConvStreamChunkPositions));
+
+  if (packFactor > 1 && chunkPositions > static_cast<uint64_t>(packFactor)) {
+    chunkPositions -= chunkPositions % static_cast<uint64_t>(packFactor);
+    chunkPositions = std::max<uint64_t>(chunkPositions, static_cast<uint64_t>(packFactor));
+  }
+  return std::max<uint64_t>(1, chunkPositions);
+}
+
+RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvLoweringState& state) {
+  const int64_t rawBegin = outputRows.begin * state.strideHeight - state.padHeightBegin;
+  const int64_t rawEnd =
+    (outputRows.end - 1) * state.strideHeight - state.padHeightBegin + state.dilationHeight * (state.wHeight - 1) + 1;
+  return {std::max<int64_t>(0, rawBegin), std::min<int64_t>(state.xHeight, rawEnd)};
+}
+
+ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvLoweringState& state) {
+  ConvRowDemand demand;
+  demand.outputRows = outputRows;
+  demand.neededInputRows = computeConvInputRowsForOutputRows(outputRows, state);
+  demand.acquiredInputRows = demand.neededInputRows;
+
+  const int64_t rawBegin = outputRows.begin * state.strideHeight - state.padHeightBegin;
+  const int64_t rawEnd =
+    (outputRows.end - 1) * state.strideHeight - state.padHeightBegin + state.dilationHeight * (state.wHeight - 1) + 1;
+  demand.topHaloRows = std::max<int64_t>(0, -rawBegin);
+  demand.bottomHaloRows = std::max<int64_t>(0, rawEnd - state.xHeight);
+  demand.acquiredInputRows = demand.neededInputRows;
+  return demand;
+}
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/Patterns/Math/ConvGeometry.hpp

@@ -0,0 +1,86 @@
+#pragma once
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Value.h"
+
+#include <cstdint>
+
+namespace onnx_mlir {
+
+struct ConvLoweringState {
+  mlir::Value x;
+  mlir::Value w;
+  mlir::Value b;
+  mlir::RankedTensorType xType;
+  mlir::RankedTensorType wType;
+  mlir::RankedTensorType outType;
+  int64_t batchSize;
+  int64_t numChannelsIn;
+  int64_t xHeight;
+  int64_t xWidth;
+  int64_t numChannelsOut;
+  int64_t wHeight;
+  int64_t wWidth;
+  int64_t outHeight;
+  int64_t outWidth;
+  int64_t group;
+  int64_t numChannelsInPerGroup;
+  int64_t numChannelsOutPerGroup;
+  int64_t padHeightBegin;
+  int64_t padHeightEnd;
+  int64_t padWidthBegin;
+  int64_t padWidthEnd;
+  int64_t strideHeight;
+  int64_t strideWidth;
+  int64_t dilationHeight;
+  int64_t dilationWidth;
+  bool hasBias;
+};
+
+struct ConvGeometry {
+  int64_t batchSize;
+  int64_t numChannelsIn;
+  int64_t xHeight;
+  int64_t xWidth;
+  int64_t numChannelsOut;
+  int64_t wHeight;
+  int64_t wWidth;
+  int64_t outHeight;
+  int64_t outWidth;
+  int64_t group;
+  int64_t numChannelsInPerGroup;
+  int64_t numChannelsOutPerGroup;
+  int64_t k;
+  int64_t c;
+  int64_t p;
+  int64_t xbarSize;
+  int64_t pack;
+  uint64_t im2colElements;
+  bool hasBias;
+  bool isDepthwise;
+};
+
+struct RowInterval {
+  int64_t begin = 0;
+  int64_t end = 0;
+};
+
+struct ConvRowDemand {
+  RowInterval outputRows;
+  RowInterval neededInputRows;
+  RowInterval acquiredInputRows;
+  int64_t topHaloRows = 0;
+  int64_t bottomHaloRows = 0;
+};
+
+bool isDepthwiseConv(int64_t group, int64_t numChannelsIn, int64_t numChannelsOut, int64_t numChannelsInPerGroup);
+
+ConvGeometry buildConvGeometry(const ConvLoweringState& state);
+
+uint64_t chooseStreamChunkPositions(const ConvGeometry& geo, int64_t packFactor);
+
+RowInterval computeConvInputRowsForOutputRows(RowInterval outputRows, const ConvLoweringState& state);
+
+ConvRowDemand buildConvRowDemand(RowInterval outputRows, const ConvLoweringState& state);
+
+} // namespace onnx_mlir

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

@@ -5,7 +5,7 @@
 
 #include "llvm/ADT/SmallVector.h"
 
-#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -47,38 +47,28 @@ static FailureOr<Value> materializeBroadcastedConstantTensor(Value value,
     return failure();
 
   const int64_t rankOffset = static_cast<int64_t>(resultShape.size() - sourceShape.size());
-  for (int64_t i = 0; i < static_cast<int64_t>(resultShape.size()); ++i) {
-    const int64_t sourceIndex = i - rankOffset;
-    const int64_t sourceDim = sourceIndex < 0 ? 1 : sourceShape[sourceIndex];
-    const int64_t resultDim = resultShape[i];
-    if (sourceDim != 1 && sourceDim != resultDim)
-      return failure();
-  }
-
-  SmallVector<Attribute> sourceValues(denseAttr.getValues<Attribute>());
   SmallVector<int64_t> sourceStrides = computeRowMajorStrides(sourceShape);
   SmallVector<int64_t> resultStrides = computeRowMajorStrides(resultShape);
-
+  SmallVector<Attribute> sourceValues(denseAttr.getValues<Attribute>());
   SmallVector<Attribute> resultValues;
   resultValues.reserve(resultType.getNumElements());
 
   for (int64_t flatIndex = 0; flatIndex < resultType.getNumElements(); ++flatIndex) {
     int64_t remaining = flatIndex;
     int64_t sourceFlatIndex = 0;
-
     for (int64_t i = 0; i < static_cast<int64_t>(resultShape.size()); ++i) {
       const int64_t resultIndex = resultStrides.empty() ? 0 : remaining / resultStrides[i];
       remaining = resultStrides.empty() ? 0 : remaining % resultStrides[i];
-
       const int64_t sourceIndex = i - rankOffset;
       if (sourceIndex < 0)
         continue;
-
       const int64_t sourceDim = sourceShape[sourceIndex];
+      const int64_t resultDim = resultShape[i];
+      if (sourceDim != 1 && sourceDim != resultDim)
+        return failure();
       const int64_t mappedIndex = sourceDim == 1 ? 0 : resultIndex;
       sourceFlatIndex += mappedIndex * sourceStrides[sourceIndex];
     }
-
     resultValues.push_back(sourceValues[sourceFlatIndex]);
   }
 
@@ -106,7 +96,7 @@ static FailureOr<Value> materializeReciprocalTensor(Value value,
   if (failed(broadcastedValue))
     return failure();
 
-  auto denseAttr = dyn_cast<DenseFPElementsAttr>(getDenseConstantAttr(*broadcastedValue));
+  auto denseAttr = dyn_cast<DenseFPElementsAttr>(getHostConstDenseElementsAttr(*broadcastedValue));
   if (!denseAttr)
     return failure();
 
@@ -185,10 +175,45 @@ struct DivToSpatialCompute : OpConversionPattern<ONNXDivOp> {
   }
 };
 
+struct AddToSpatialCompute : OpConversionPattern<ONNXAddOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(ONNXAddOp op, ONNXAddOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const override {
+    auto resultType = dyn_cast<RankedTensorType>(op.getResult().getType());
+    if (!resultType || !resultType.hasStaticShape())
+      return failure();
+
+    FailureOr<BiasAddPlanCandidate> candidate =
+      classifyBiasAddPlanCandidate(adaptor.getA(), adaptor.getB(), resultType);
+    if (succeeded(candidate)) {
+      auto plan = spatial::SpatBiasAddPlanOp::create(
+        rewriter, op.getLoc(), resultType, candidate->data, candidate->bias, rewriter.getStringAttr("nchw"));
+      rewriter.replaceOp(op, plan.getResult());
+      return success();
+    }
+
+    auto lhs = prepareElementwiseOperand(adaptor.getA(), resultType, rewriter, op.getLoc());
+    if (failed(lhs))
+      return failure();
+    auto rhs = prepareElementwiseOperand(adaptor.getB(), resultType, rewriter, op.getLoc());
+    if (failed(rhs))
+      return failure();
+
+    auto computeOp =
+      createSpatCompute<2>(rewriter, op.getLoc(), resultType, {}, ValueRange {*lhs, *rhs}, [&](Value x, Value y) {
+        auto loweredOp = spatial::SpatVAddOp::create(rewriter, op.getLoc(), resultType, x, y);
+        spatial::SpatYieldOp::create(rewriter, op.getLoc(), loweredOp.getResult());
+      });
+    rewriter.replaceOp(op, computeOp);
+    return success();
+  }
+};
+
 } // namespace
 
 void populateElementwisePatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
-  patterns.add<BinaryElementwiseToSpatialCompute<ONNXAddOp, spatial::SpatVAddOp>>(ctx);
+  patterns.add<AddToSpatialCompute>(ctx);
   patterns.add<BinaryElementwiseToSpatialCompute<ONNXSubOp, spatial::SpatVSubOp>>(ctx);
   patterns.add<BinaryElementwiseToSpatialCompute<ONNXMulOp, spatial::SpatVMulOp>>(ctx);
   patterns.add<DivToSpatialCompute>(ctx);

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

@@ -87,28 +87,6 @@ static Value createGemmBatchHOffset(Value lane,
                                  rewriter.getInsertionBlock()->getParentOp());
 }
 
-static Value
-createZeroPaddedTensor(Value value, RankedTensorType resultType, ConversionPatternRewriter& rewriter, Location loc) {
-  auto sourceType = cast<RankedTensorType>(value.getType());
-  SmallVector<OpFoldResult> lowPads(sourceType.getRank(), rewriter.getIndexAttr(0));
-  SmallVector<OpFoldResult> highPads;
-  highPads.reserve(sourceType.getRank());
-  for (auto [sourceDim, resultDim] : llvm::zip(sourceType.getShape(), resultType.getShape()))
-    highPads.push_back(rewriter.getIndexAttr(resultDim - sourceDim));
-
-  auto padOp = tensor::PadOp::create(rewriter, loc, resultType, value, lowPads, highPads);
-  auto* padBlock = new Block();
-  for (int64_t i = 0; i < sourceType.getRank(); ++i)
-    padBlock->addArgument(rewriter.getIndexType(), loc);
-  padOp.getRegion().push_back(padBlock);
-  rewriter.setInsertionPointToStart(padBlock);
-  auto zero = getOrCreateConstant(
-    rewriter, padOp.getOperation(), rewriter.getZeroAttr(sourceType.getElementType()), sourceType.getElementType());
-  tensor::YieldOp::create(rewriter, loc, zero);
-  rewriter.setInsertionPointAfter(padOp);
-  return padOp.getResult();
-}
-
 static FailureOr<Value> materializePaddedConstantMatrix(Value value,
                                                         RankedTensorType resultType,
                                                         ConversionPatternRewriter& rewriter,
@@ -232,22 +210,6 @@ static Value extractATile(
   return tensor::ExtractSliceOp::create(rewriter, loc, aTileType, a, offsets, sizes, strides).getResult();
 }
 
-static Value createPaddedInputCompute(Value input,
-                                      RankedTensorType paddedInputType,
-                                      ConversionPatternRewriter& rewriter,
-                                      Location loc) {
-  auto inputType = cast<RankedTensorType>(input.getType());
-  if (inputType == paddedInputType)
-    return input;
-
-  auto computeOp = createSpatCompute<1>(rewriter, loc, TypeRange {paddedInputType}, {}, input, [&](Value computeInput) {
-    Value paddedInput = createZeroPaddedTensor(computeInput, paddedInputType, rewriter, loc);
-    spatial::SpatYieldOp::create(rewriter, loc, paddedInput);
-  });
-
-  return computeOp.getResult(0);
-}
-
 static FailureOr<spatial::SpatComputeBatch> createVmmBatch(Value a,
                                                            Value b,
                                                            RankedTensorType aType,
@@ -285,9 +247,8 @@ static FailureOr<spatial::SpatComputeBatch> createVmmBatch(Value a,
       SmallVector<OpFoldResult> bSizes {rewriter.getIndexAttr(crossbarSize.getValue()),
                                         rewriter.getIndexAttr(crossbarSize.getValue())};
       SmallVector<OpFoldResult> unitStrides = getUnitStrides(rewriter, 2);
-      Value bTile =
-        tensor::ExtractSliceOp::create(rewriter, loc, bTileType, args.weights.front(), bOffsets, bSizes, unitStrides)
-          .getResult();
+      Value bTile = extractStaticSliceOrIdentity(
+        rewriter, loc, args.weights.front(), bTileType, bOffsets, bSizes, unitStrides);
       Value piece = spatial::SpatVMMOp::create(rewriter, loc, pieceType, bTile, aTile).getResult();
 
       SmallVector<OpFoldResult> pieceOffsets {args.lane, rewriter.getIndexAttr(0)};

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

@@ -255,42 +255,6 @@ static Value transposeLastTwoDims(Value value, PatternRewriter& rewriter, Locati
   return createONNXTranspose(resultType, {0, 2, 1});
 }
 
-static Value createZeroPaddedTensor(Value value, RankedTensorType resultType, PatternRewriter& rewriter, Location loc) {
-  auto sourceType = cast<RankedTensorType>(value.getType());
-  SmallVector<OpFoldResult> lowPads(sourceType.getRank(), rewriter.getIndexAttr(0));
-  SmallVector<OpFoldResult> highPads;
-  highPads.reserve(sourceType.getRank());
-  for (auto [sourceDim, resultDim] : llvm::zip(sourceType.getShape(), resultType.getShape()))
-    highPads.push_back(rewriter.getIndexAttr(resultDim - sourceDim));
-
-  auto padOp = tensor::PadOp::create(rewriter, loc, resultType, value, lowPads, highPads);
-  auto* padBlock = new Block();
-  for (int64_t i = 0; i < sourceType.getRank(); ++i)
-    padBlock->addArgument(rewriter.getIndexType(), loc);
-  padOp.getRegion().push_back(padBlock);
-  rewriter.setInsertionPointToStart(padBlock);
-  auto zero = getOrCreateConstant(
-    rewriter, padOp.getOperation(), rewriter.getZeroAttr(sourceType.getElementType()), sourceType.getElementType());
-  tensor::YieldOp::create(rewriter, loc, zero);
-  rewriter.setInsertionPointAfter(padOp);
-  return padOp.getResult();
-}
-
-static Value createPaddedBatchedInputCompute(Value input,
-                                             RankedTensorType paddedInputType,
-                                             PatternRewriter& rewriter,
-                                             Location loc) {
-  auto inputType = cast<RankedTensorType>(input.getType());
-  if (inputType == paddedInputType)
-    return input;
-
-  auto computeOp = createSpatCompute<1>(rewriter, loc, TypeRange {paddedInputType}, {}, input, [&](Value computeInput) {
-    Value paddedInput = createZeroPaddedTensor(computeInput, paddedInputType, rewriter, loc);
-    spatial::SpatYieldOp::create(rewriter, loc, paddedInput);
-  });
-  return computeOp.getResult(0);
-}
-
 static FailureOr<Value> materializePaddedBatchedWeight(Value value,
                                                        ArrayRef<int64_t> sourceBatchShape,
                                                        ArrayRef<int64_t> targetBatchShape,
@@ -950,7 +914,10 @@ struct MatMulToGemm : OpRewritePattern<ONNXMatMulOp> {
 
   LogicalResult matchAndRewrite(ONNXMatMulOp matmulOp, PatternRewriter& rewriter) const override {
     auto shapeInfo = analyzeMatMulShape(matmulOp);
-    if (failed(shapeInfo) || shapeInfo->lhsWasVector || shapeInfo->rhsWasVector || !shapeInfo->outputBatchShape.empty())
+    if (failed(shapeInfo) || shapeInfo->lhsWasVector || shapeInfo->rhsWasVector)
+      return failure();
+
+    if (!shapeInfo->outputBatchShape.empty())
       return failure();
 
     Location loc = matmulOp.getLoc();
@@ -991,9 +958,19 @@ struct MatMulToGemm : OpRewritePattern<ONNXMatMulOp> {
       gemmResult =
         ONNXTransposeOp::create(rewriter, loc, shapeInfo->outType, gemmResult, rewriter.getI64ArrayAttr({1, 0}))
           .getResult();
+
+    if (shapeInfo->outputBatchShape.empty()) {
       rewriter.replaceOp(matmulOp, gemmResult);
       return success();
     }
+
+    auto directOutType =
+      RankedTensorType::get({1, shapeInfo->m, shapeInfo->n}, shapeInfo->outType.getElementType(), shapeInfo->outType.getEncoding());
+    Value batchedResult = ensureBatchedTensor(gemmResult, /*batchSize=*/1, shapeInfo->m, shapeInfo->n, rewriter, loc);
+    Value finalResult = finalizeNormalizedMatMulResult(batchedResult, directOutType, *shapeInfo, rewriter, loc);
+    rewriter.replaceOp(matmulOp, finalResult);
+    return success();
+  }
 };
 
 struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
@@ -1040,7 +1017,7 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
       auto paddedRhs =
         materializePaddedBatchedWeight(plan.rhs, plan.rhsBatchShape, plan.outputBatchShape, paddedRhsType, rewriter);
       if (succeeded(paddedRhs)) {
-        Value paddedLhs = createPaddedBatchedInputCompute(plan.lhs, paddedLhsType, rewriter, loc);
+        Value paddedLhs = createPaddedInputCompute(plan.lhs, paddedLhsType, rewriter, loc);
         const int64_t laneCount = plan.batch * plan.m * numKSlices * numOutHSlices;
         auto partialPiecesType = RankedTensorType::get({laneCount, static_cast<int64_t>(crossbarSize.getValue())},
                                                        shapeInfo->outType.getElementType());

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

@@ -16,12 +16,9 @@ struct ReluToSpatialCompute : OpConversionPattern<ONNXReluOp> {
   matchAndRewrite(ONNXReluOp reluOp, ONNXReluOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const override {
     Location loc = reluOp.getLoc();
     Type resultType = reluOp.getResult().getType();
-    constexpr size_t numInputs = 1;
-    auto computeOp = createSpatCompute<numInputs>(rewriter, loc, resultType, {}, adaptor.getX(), [&](Value x) {
-      auto spatReluOp = spatial::SpatReluOp::create(rewriter, loc, resultType, x);
-      spatial::SpatYieldOp::create(rewriter, loc, spatReluOp.getResult());
-    });
-    rewriter.replaceOp(reluOp, computeOp);
+    auto reluPlan = spatial::SpatReluPlanOp::create(
+      rewriter, loc, resultType, adaptor.getX(), rewriter.getStringAttr("nchw"));
+    rewriter.replaceOp(reluOp, reluPlan.getResult());
     return success();
   }
 };

src/PIM/Conversion/ONNXToSpatial/Patterns/Post.cpp

@@ -10,6 +10,7 @@
 
 #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/WeightMaterialization.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -118,18 +119,16 @@ static LogicalResult mapPromotedInputArguments(ComputeOpTy compute,
 }
 
 // Promotes foldable helper chains from runtime inputs to weights to avoid artificial compute inputs.
-struct PromoteWeightLikeComputeInputsPattern : OpRewritePattern<spatial::SpatCompute> {
-  using OpRewritePattern<spatial::SpatCompute>::OpRewritePattern;
+struct PromoteWeightLikeComputeInputsPattern : OpRewritePattern<spatial::SpatGraphCompute> {
+  using OpRewritePattern<spatial::SpatGraphCompute>::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(spatial::SpatCompute compute, PatternRewriter& rewriter) const override {
+  LogicalResult matchAndRewrite(spatial::SpatGraphCompute compute, PatternRewriter& rewriter) const override {
     auto promoted = computePromotedOperands(compute);
     if (failed(promoted))
       return rewriter.notifyMatchFailure(compute, "no weight-like inputs to promote");
     Block& oldBlock = compute.getBody().front();
 
     rewriter.setInsertionPointAfter(compute);
-    auto newCompute = spatial::SpatCompute::create(
-      rewriter, compute.getLoc(), compute.getResultTypes(), promoted->newWeights, promoted->newInputs);
     SmallVector<Type> newBlockArgTypes;
     SmallVector<Location> newBlockArgLocs;
     for (Value weight : promoted->newWeights) {
@@ -138,10 +137,14 @@ struct PromoteWeightLikeComputeInputsPattern : OpRewritePattern<spatial::SpatCom
     }
     llvm::append_range(newBlockArgTypes, promoted->newInputTypes);
     llvm::append_range(newBlockArgLocs, promoted->newInputLocs);
-    auto* newBlock = rewriter.createBlock(
-      &newCompute.getBody(), newCompute.getBody().end(), TypeRange(newBlockArgTypes), newBlockArgLocs);
-    newCompute.getProperties().setOperandSegmentSizes(
-      {static_cast<int>(promoted->newWeights.size()), static_cast<int>(promoted->newInputs.size())});
+    auto newCompute = createEmptySpatGraphCompute(rewriter,
+                                                  compute.getLoc(),
+                                                  compute.getResultTypes(),
+                                                  promoted->newWeights,
+                                                  promoted->newInputs,
+                                                  TypeRange(newBlockArgTypes),
+                                                  newBlockArgLocs);
+    auto* newBlock = &newCompute.getBody().front();
     rewriter.setInsertionPointToStart(newBlock);
 
     IRRewriter bodyRewriter(rewriter.getContext());
@@ -182,10 +185,10 @@ struct PromoteWeightLikeComputeInputsPattern : OpRewritePattern<spatial::SpatCom
 };
 
 // Promotes foldable batch helper chains to weights while preserving compact compute_batch IR.
-struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::SpatComputeBatch> {
-  using OpRewritePattern<spatial::SpatComputeBatch>::OpRewritePattern;
+struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::SpatGraphComputeBatch> {
+  using OpRewritePattern<spatial::SpatGraphComputeBatch>::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(spatial::SpatComputeBatch compute, PatternRewriter& rewriter) const override {
+  LogicalResult matchAndRewrite(spatial::SpatGraphComputeBatch compute, PatternRewriter& rewriter) const override {
     auto promoted = computePromotedOperands(compute);
     if (failed(promoted))
       return rewriter.notifyMatchFailure(compute, "no weight-like batch inputs to promote");
@@ -193,12 +196,6 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::Sp
 
     rewriter.setInsertionPointAfter(compute);
 
-    auto laneCountAttr = pim::getCheckedI32Attr(
-      rewriter, compute, static_cast<uint64_t>(compute.getLaneCount()), "promoted compute_batch lane count");
-    if (failed(laneCountAttr))
-      return failure();
-    auto newCompute = spatial::SpatComputeBatch::create(
-      rewriter, compute.getLoc(), compute.getResultTypes(), *laneCountAttr, promoted->newWeights, promoted->newInputs);
     auto laneArg = compute.getLaneArgument();
     if (!laneArg)
       return rewriter.notifyMatchFailure(compute, "missing compute_batch lane block argument");
@@ -223,23 +220,30 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::Sp
       newBlockArgLocs.push_back(outputArg->getLoc());
     }
 
-    auto* newBlock = rewriter.createBlock(
-      &newCompute.getBody(), newCompute.getBody().end(), TypeRange(newBlockArgTypes), newBlockArgLocs);
-    newCompute.getProperties().setOperandSegmentSizes(
-      {static_cast<int>(promoted->newWeights.size()), static_cast<int>(promoted->newInputs.size())});
+    auto newCompute = createEmptySpatGraphComputeBatch(rewriter,
+                                                       compute.getLoc(),
+                                                       compute.getResultTypes(),
+                                                       compute.getLaneCount(),
+                                                       promoted->newWeights,
+                                                       promoted->newInputs,
+                                                       TypeRange(newBlockArgTypes),
+                                                       newBlockArgLocs);
+    if (failed(newCompute))
+      return failure();
+    auto* newBlock = &(*newCompute).getBody().front();
     rewriter.setInsertionPointToStart(newBlock);
 
     IRRewriter bodyRewriter(rewriter.getContext());
     bodyRewriter.setInsertionPointToStart(newBlock);
 
     IRMapping mapper;
-    auto newLaneArg = newCompute.getLaneArgument();
+    auto newLaneArg = (*newCompute).getLaneArgument();
     if (!newLaneArg)
       return rewriter.notifyMatchFailure(compute, "missing rewritten compute_batch lane block argument");
     mapper.map(*laneArg, *newLaneArg);
     for (auto [weightIndex, weight] : llvm::enumerate(compute.getWeights())) {
       auto oldWeightArg = compute.getWeightArgument(weightIndex);
-      auto newWeightArg = newCompute.getWeightArgument(weightIndex);
+      auto newWeightArg = (*newCompute).getWeightArgument(weightIndex);
       if (!oldWeightArg || !newWeightArg)
         return rewriter.notifyMatchFailure(compute, "missing compute_batch weight block argument during rewrite");
       mapper.map(*oldWeightArg, *newWeightArg);
@@ -249,7 +253,7 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::Sp
           *promoted,
           bodyRewriter,
           mapper,
-          [&](size_t index) { return newCompute.getInputArgument(index); },
+          [&](size_t index) { return (*newCompute).getInputArgument(index); },
           rewriter)))
       return failure();
     for (auto resultIndex : llvm::seq<size_t>(0, compute.getNumResults())) {
@@ -263,7 +267,7 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::Sp
     for (Operation& op : oldBlock)
       rewriter.clone(op, mapper);
 
-    rewriter.replaceOp(compute, newCompute.getResults());
+    rewriter.replaceOp(compute, (*newCompute).getResults());
     return success();
   }
 };
@@ -281,8 +285,8 @@ void annotateWeightsConstants(func::FuncOp funcOp) {
   });
 }
 
-bool requiresPostRewrite(spatial::SpatCompute computeOp) { return hasPromotableWeightLikeInputs(computeOp); }
+bool requiresPostRewrite(spatial::SpatGraphCompute computeOp) { return hasPromotableWeightLikeInputs(computeOp); }
 
-bool requiresPostRewrite(spatial::SpatComputeBatch computeOp) { return hasPromotableWeightLikeInputs(computeOp); }
+bool requiresPostRewrite(spatial::SpatGraphComputeBatch computeOp) { return hasPromotableWeightLikeInputs(computeOp); }
 
 } // namespace onnx_mlir

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

@@ -0,0 +1,112 @@
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Transforms/DialectConversion.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/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+static FailureOr<int64_t> normalizeFlattenAxis(int64_t axis, int64_t rank) {
+  int64_t normalizedAxis = axis < 0 ? rank + axis : axis;
+  if (normalizedAxis < 0 || normalizedAxis > rank)
+    return failure();
+  return normalizedAxis;
+}
+
+static int64_t product(ArrayRef<int64_t> values) {
+  int64_t result = 1;
+  for (int64_t value : values)
+    result *= value;
+  return result;
+}
+
+static SmallVector<ReassociationIndices> getCollapseTo1DReassociation(int64_t rank) {
+  SmallVector<ReassociationIndices> reassociation(1);
+  reassociation.front().reserve(rank);
+  for (int64_t dim = 0; dim < rank; ++dim)
+    reassociation.front().push_back(dim);
+  return reassociation;
+}
+
+static SmallVector<ReassociationIndices> getExpandFrom1DReassociation(int64_t rank) {
+  SmallVector<ReassociationIndices> reassociation(1);
+  reassociation.front().reserve(rank);
+  for (int64_t dim = 0; dim < rank; ++dim)
+    reassociation.front().push_back(dim);
+  return reassociation;
+}
+
+static Value buildFlatten(Value input,
+                          RankedTensorType sourceType,
+                          RankedTensorType resultType,
+                          int64_t axis,
+                          ConversionPatternRewriter& rewriter,
+                          Location loc) {
+  if (sourceType == resultType)
+    return input;
+
+  if (axis > 0 && axis < sourceType.getRank()) {
+    SmallVector<ReassociationIndices> reassociation(2);
+    for (int64_t dim = 0; dim < axis; ++dim)
+      reassociation[0].push_back(dim);
+    for (int64_t dim = axis; dim < sourceType.getRank(); ++dim)
+      reassociation[1].push_back(dim);
+    return tensor::CollapseShapeOp::create(rewriter, loc, resultType, input, reassociation);
+  }
+
+  Value flattened = input;
+  if (sourceType.getRank() != 1) {
+    auto flatType = RankedTensorType::get({sourceType.getNumElements()}, sourceType.getElementType());
+    flattened = tensor::CollapseShapeOp::create(
+      rewriter, loc, flatType, flattened, getCollapseTo1DReassociation(sourceType.getRank()));
+  }
+  return tensor::ExpandShapeOp::create(
+    rewriter, loc, resultType, flattened, getExpandFrom1DReassociation(resultType.getRank()));
+}
+
+struct Flatten : OpConversionPattern<ONNXFlattenOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(ONNXFlattenOp flattenOp,
+                                ONNXFlattenOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    auto sourceType = dyn_cast<RankedTensorType>(adaptor.getInput().getType());
+    auto resultType = dyn_cast<RankedTensorType>(flattenOp.getOperation()->getResult(0).getType());
+    if (!sourceType || !resultType || !sourceType.hasStaticShape() || !resultType.hasStaticShape())
+      return failure();
+    if (!hasStaticPositiveShape(sourceType) || !hasStaticPositiveShape(resultType) || resultType.getRank() != 2)
+      return failure();
+
+    auto axis = normalizeFlattenAxis(flattenOp.getAxis(), sourceType.getRank());
+    if (failed(axis))
+      return failure();
+
+    int64_t outerDim = product(sourceType.getShape().take_front(*axis));
+    int64_t innerDim = product(sourceType.getShape().drop_front(*axis));
+    if (resultType.getShape()[0] != outerDim || resultType.getShape()[1] != innerDim)
+      return failure();
+
+    auto replaceWithFlatten = [&](auto build) -> LogicalResult {
+      Value flattened = materializeOrComputeUnary(adaptor.getInput(), resultType, rewriter, flattenOp.getLoc(), build);
+      rewriter.replaceOp(flattenOp, flattened);
+      return success();
+    };
+
+    return replaceWithFlatten([&](Value input) {
+      return buildFlatten(input, sourceType, resultType, *axis, rewriter, flattenOp.getLoc());
+    });
+  }
+};
+
+} // namespace
+
+void populateFlattenPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.add<Flatten>(ctx); }
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp

@@ -0,0 +1,21 @@
+#pragma once
+
+#include <optional>
+
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Support/LogicalResult.h"
+
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+namespace onnx_mlir {
+
+mlir::FailureOr<mlir::Value>
+lowerSelectedConv2DPlan(spatial::SpatConv2DPlanOp planOp,
+                        std::optional<mlir::Value> rowStripInput,
+                        bool emitRowStripLayout,
+                        mlir::PatternRewriter& rewriter);
+
+mlir::LogicalResult canLowerConvPlanToRowStrip(spatial::SpatConv2DPlanOp planOp);
+mlir::LogicalResult canConsumeAndProduceRowStrip(spatial::SpatConv2DPlanOp planOp);
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -0,0 +1,243 @@
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/Pass/Pass.h"
+
+#include "llvm/ADT/DenseMap.h"
+
+#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Pass/PIMPasses.h"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+static constexpr StringLiteral kLogicalLayout = "nchw";
+static constexpr StringLiteral kDenseLayout = "dense_nchw";
+static constexpr StringLiteral kRowStripLayout = "nchw_row_strip";
+
+enum class SelectedLayout {
+  DenseNchw,
+  NchwRowStrip,
+};
+
+static SelectedLayout getSelectedLayout(llvm::DenseMap<Value, SelectedLayout>& layouts, Value value) {
+  auto it = layouts.find(value);
+  return it == layouts.end() ? SelectedLayout::DenseNchw : it->second;
+}
+
+static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap<Value, SelectedLayout>& layouts) {
+  if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(user))
+    return getSelectedLayout(layouts, reluPlan.getResult()) == SelectedLayout::NchwRowStrip;
+  if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user))
+    return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::NchwRowStrip;
+  if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
+    return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::NchwRowStrip;
+  return false;
+}
+
+static bool allUsersCanHandleRowStrip(Value value, llvm::DenseMap<Value, SelectedLayout>& layouts) {
+  for (Operation* user : value.getUsers()) {
+    if (usesSelectedRowStrip(user, layouts))
+      continue;
+    // Dense-only users must be materialized explicitly.
+    continue;
+  }
+  return true;
+}
+
+static bool canConsumeRowStripAsUser(Operation* user) {
+  if (isa<spatial::SpatReluPlanOp>(user))
+    return true;
+  if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user)) {
+    auto resultType = dyn_cast<RankedTensorType>(biasAddPlan.getOutput().getType());
+    return resultType && isSupportedBiasAddValue(biasAddPlan.getBias(), resultType);
+  }
+  if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
+    return succeeded(canConsumeAndProduceRowStrip(convPlan));
+  return false;
+}
+
+static bool hasRowStripConsumer(Value value) {
+  for (Operation* user : value.getUsers())
+    if (canConsumeRowStripAsUser(user))
+      return true;
+  return false;
+}
+
+
+static bool canSelectConvRowStrip(spatial::SpatConv2DPlanOp convPlan,
+                                  llvm::DenseMap<Value, SelectedLayout>& layouts) {
+  SelectedLayout inputLayout = getSelectedLayout(layouts, convPlan.getInput());
+  if (inputLayout == SelectedLayout::NchwRowStrip)
+    return succeeded(canConsumeAndProduceRowStrip(convPlan));
+  return succeeded(canLowerConvPlanToRowStrip(convPlan));
+}
+
+static SelectedLayout chooseConvLayout(spatial::SpatConv2DPlanOp convPlan,
+                                       llvm::DenseMap<Value, SelectedLayout>& layouts) {
+  if (!canSelectConvRowStrip(convPlan, layouts))
+    return SelectedLayout::DenseNchw;
+  if (getSelectedLayout(layouts, convPlan.getInput()) != SelectedLayout::NchwRowStrip
+      && !hasRowStripConsumer(convPlan.getResult()))
+    return SelectedLayout::DenseNchw;
+  if (!allUsersCanHandleRowStrip(convPlan.getResult(), layouts))
+    return SelectedLayout::DenseNchw;
+  return SelectedLayout::NchwRowStrip;
+}
+
+static SelectedLayout chooseReluLayout(spatial::SpatReluPlanOp reluPlan,
+                                       llvm::DenseMap<Value, SelectedLayout>& layouts) {
+  if (getSelectedLayout(layouts, reluPlan.getInput()) != SelectedLayout::NchwRowStrip)
+    return SelectedLayout::DenseNchw;
+  if (!hasRowStripConsumer(reluPlan.getResult()))
+    return SelectedLayout::DenseNchw;
+  if (!allUsersCanHandleRowStrip(reluPlan.getResult(), layouts))
+    return SelectedLayout::DenseNchw;
+  return SelectedLayout::NchwRowStrip;
+}
+
+static SelectedLayout chooseBiasAddLayout(spatial::SpatBiasAddPlanOp biasAddPlan,
+                                          llvm::DenseMap<Value, SelectedLayout>& layouts) {
+  if (getSelectedLayout(layouts, biasAddPlan.getInput()) != SelectedLayout::NchwRowStrip)
+    return SelectedLayout::DenseNchw;
+  auto resultType = dyn_cast<RankedTensorType>(biasAddPlan.getOutput().getType());
+  if (!resultType || !isSupportedBiasAddValue(biasAddPlan.getBias(), resultType))
+    return SelectedLayout::DenseNchw;
+  if (!hasRowStripConsumer(biasAddPlan.getResult()))
+    return SelectedLayout::DenseNchw;
+  if (!allUsersCanHandleRowStrip(biasAddPlan.getResult(), layouts))
+    return SelectedLayout::DenseNchw;
+  return SelectedLayout::NchwRowStrip;
+}
+
+static spatial::SpatBlueprintOp insertRowStripBlueprint(IRRewriter& rewriter, Value value) {
+  auto outputType = cast<RankedTensorType>(value.getType());
+  auto [offsets, sizes] = buildRowStripMetadata(outputType);
+  return spatial::SpatBlueprintOp::create(rewriter,
+                                              value.getLoc(),
+                                              outputType,
+                                              value,
+                                              ValueRange {},
+                                              rewriter.getStringAttr(kLogicalLayout),
+                                              rewriter.getStringAttr(kRowStripLayout),
+                                              rewriter.getDenseI64ArrayAttr(offsets),
+                                              rewriter.getDenseI64ArrayAttr(sizes),
+                                              rewriter.getStringAttr(kRowStripIndexMap),
+                                              nullptr,
+                                              nullptr,
+                                              nullptr,
+                                              nullptr,
+                                              nullptr,
+                                              nullptr);
+}
+
+static void materializeDenseUses(IRRewriter& rewriter,
+                                 Value layoutValue,
+                                 llvm::DenseMap<Value, SelectedLayout>& layouts) {
+  SmallVector<OpOperand*> denseUses;
+  for (OpOperand& use : layoutValue.getUses()) {
+    if (usesSelectedRowStrip(use.getOwner(), layouts))
+      continue;
+    denseUses.push_back(&use);
+  }
+
+  for (OpOperand* use : denseUses) {
+    Operation* owner = use->getOwner();
+    rewriter.setInsertionPoint(owner);
+    auto materialized = spatial::SpatMaterializeLayoutOp::create(rewriter,
+                                                                 owner->getLoc(),
+                                                                 use->get().getType(),
+                                                                 use->get(),
+                                                                 rewriter.getStringAttr(kLogicalLayout),
+                                                                 rewriter.getStringAttr(kRowStripLayout),
+                                                                 rewriter.getStringAttr(kDenseLayout));
+    use->set(materialized.getResult());
+  }
+}
+
+struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SpatialLayoutPlanningPass)
+
+  StringRef getArgument() const override { return "spatial-layout-planning"; }
+  StringRef getDescription() const override { return "Select conservative Spatial layouts and insert reconciliation barriers."; }
+
+  void runOnOperation() override {
+    auto entryFunc = getPimEntryFunc(getOperation());
+    if (failed(entryFunc)) {
+      getOperation().emitError("failed to locate the PIM entry function during Spatial layout planning");
+      signalPassFailure();
+      return;
+    }
+
+    func::FuncOp funcOp = *entryFunc;
+    IRRewriter rewriter(&getContext());
+    llvm::DenseMap<Value, SelectedLayout> layouts;
+
+    bool changed = true;
+    while (changed) {
+      changed = false;
+      for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
+        if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(&op)) {
+          SelectedLayout selected = chooseConvLayout(convPlan, layouts);
+          if (layouts[convPlan.getResult()] != selected) {
+            layouts[convPlan.getResult()] = selected;
+            changed = true;
+          }
+          continue;
+        }
+        if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op)) {
+          SelectedLayout selected = chooseReluLayout(reluPlan, layouts);
+          if (layouts[reluPlan.getResult()] != selected) {
+            layouts[reluPlan.getResult()] = selected;
+            changed = true;
+          }
+          continue;
+        }
+        if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(&op)) {
+          SelectedLayout selected = chooseBiasAddLayout(biasAddPlan, layouts);
+          if (layouts[biasAddPlan.getResult()] != selected) {
+            layouts[biasAddPlan.getResult()] = selected;
+            changed = true;
+          }
+          continue;
+        }
+      }
+    }
+
+    for (Operation& op : llvm::make_early_inc_range(funcOp.getBody().front())) {
+      Value producedValue;
+      if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(&op))
+        producedValue = convPlan.getResult();
+      else if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(&op))
+        producedValue = biasAddPlan.getResult();
+      else if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op))
+        producedValue = reluPlan.getResult();
+      else
+        continue;
+
+      if (getSelectedLayout(layouts, producedValue) != SelectedLayout::NchwRowStrip)
+        continue;
+
+      rewriter.setInsertionPointAfter(&op);
+      auto blueprint = insertRowStripBlueprint(rewriter, producedValue);
+      rewriter.replaceAllUsesExcept(producedValue, blueprint.getResult(), blueprint);
+      materializeDenseUses(rewriter, blueprint.getResult(), layouts);
+    }
+    if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
+      getOperation().emitError("logical Spatial graph verification failed after SpatialLayoutPlanning");
+      signalPassFailure();
+    }
+  }
+};
+
+} // namespace
+
+std::unique_ptr<Pass> createSpatialLayoutPlanningPass() { return std::make_unique<SpatialLayoutPlanningPass>(); }
+
+} // namespace onnx_mlir

src/PIM/Conversion/SpatialToGraphviz/CMakeLists.txt

@@ -1,17 +0,0 @@
-add_onnx_mlir_rewriter(SpatialToGraphviz)
-
-add_pim_library(OMSpatialToGraphviz
-  SpatialToGraphviz.cpp
-
-  EXCLUDE_FROM_OM_LIBS
-
-  LINK_LIBS PUBLIC
-  MLIRTosaDialect
-  OMCompilerOptions
-  OMPimCommon
-  OMONNXOps
-  SpatialOps
-
-  ACCEL_INCLUDE_DIRS PRIVATE
-  ${PIM_GENERATED_INCLUDE_DIRS}
-)

src/PIM/Conversion/SpatialToGraphviz/SpatialToGraphviz.cpp

@@ -1,259 +0,0 @@
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/Dialect/Tosa/IR/TosaOps.h"
-#include "mlir/IR/Block.h"
-#include "mlir/IR/Diagnostics.h"
-#include "mlir/IR/Value.h"
-#include "mlir/Pass/Pass.h"
-#include "mlir/Support/LLVM.h"
-#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-
-#include "llvm/Support/Casting.h"
-#include "llvm/Support/Format.h"
-
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
-#include "src/Dialect/ONNX/ONNXOps.hpp"
-
-#define FORMAT_OPERATION(op) 'x' << llvm::format_hex_no_prefix(reinterpret_cast<size_t>(op), 0)
-#define FORMAT_ARGUMENT(computeOpPointer, argumentNum) llvm::format("Arg_%p_%u", computeOpPointer, argumentNum)
-
-using namespace mlir;
-
-namespace onnx_mlir {
-
-namespace {
-
-struct SpatialToGraphvizPass : public PassWrapper<SpatialToGraphvizPass, OperationPass<ModuleOp>> {
-
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SpatialToGraphvizPass)
-
-  StringRef getArgument() const override { return "convert-spatial-to-graphviz"; }
-
-  StringRef getDescription() const override { return "Lower ONNX ops to Spatial ops."; }
-
-  SpatialToGraphvizPass(raw_ostream& os = llvm::errs())
-  : os(os) {}
-  SpatialToGraphvizPass(const SpatialToGraphvizPass& pass)
-  : SpatialToGraphvizPass(pass.os) {}
-  void runOnOperation() final;
-
-private:
-  raw_ostream& os;
-
-  /**
-   * Draws the subgraph for a given spatial::SpatCompute, including:
-   * 1. Input nodes (block arguments)
-   * 2. Operations
-   * 3. Edges between yield (output) and its users
-   *
-   * @param op The spatial::SpatCompute to draw the subgraph for.
-   * @param computeNum The number of the compute operation.
-   */
-  void drawComputeOpSubgraph(spatial::SpatCompute op, size_t computeNum) {
-    os << "\tsubgraph cluster" << computeNum << " {\n\t\tlabel=\"Compute" << computeNum << "\";\n"
-       << "\t\tstyle=filled;\n"
-       << "\t\tcolor=lightblue;\n";
-
-    Block& block = op.getBody().front();
-
-    // Inputs
-    size_t inputNum = 0;
-    for (BlockArgument& input : block.getArguments()) {
-
-      auto fromOp = FORMAT_ARGUMENT(op.getOperation(), inputNum);
-
-      os << "\t\t" << fromOp << " [label=\"Arg" << inputNum << "\",shape=box];\n";
-      for (auto userOp : input.getUsers())
-        os << "\t\t" << fromOp << " -> " << FORMAT_OPERATION(userOp) << ";\n";
-      inputNum++;
-    }
-
-    // Iterate operations
-    for (auto& childOp : block.getOperations()) {
-      os << "\t\t" << FORMAT_OPERATION(&childOp) << " [label=\"" << childOp.getName() << "\"];\n";
-
-      drawEdgesFromOpToItsUsers(&childOp);
-    }
-
-    os << "\t}\n";
-
-    // Draw edges from the yield to the users of this computeOp
-    Operation* yieldOp = block.getTerminator();
-    if (!isa<spatial::SpatYieldOp>(yieldOp)) {
-      yieldOp->emitError("Terminator of block must be YieldOp ???");
-      signalPassFailure();
-      return;
-    }
-
-    for (auto computeOpResult : op->getResults()) {
-      for (auto& computeOpUse : computeOpResult.getUses()) {
-        auto toOp = FORMAT_ARGUMENT(computeOpUse.getOwner(), computeOpUse.getOperandNumber());
-        os << "\t" << FORMAT_OPERATION(yieldOp) << " -> " << toOp << ";\n";
-      }
-    }
-  }
-
-  /**
-   * @brief Draws the subgraph for a concatOp.
-   *
-   * This function draws a subgraph for a concatOp. The subgraph consists of a
-   * node for each input of the concatOp, as well as an output node. Edges are
-   * created from the output node to each user of the concatOp.
-   *
-   * @param concatOp The concatOp for which the subgraph is drawn.
-   * @param concatOpNum The number of the concatOp.
-   */
-  void drawConcatOpSubgraph(Operation* concatOp, size_t concatOpNum) {
-    os << "\tsubgraph clusterconcat" << concatOpNum << " {\n\t\tlabel=\"ConcatOp" << concatOpNum << "\";\n"
-       << "\t\tstyle=filled;\n"
-       << "\t\tcolor=orange;\n";
-
-    // Inputs
-    size_t inputNum = 0;
-    for (Value input : concatOp->getOperands()) {
-      auto fromOp = FORMAT_ARGUMENT(concatOp, inputNum);
-
-      os << "\t\t" << fromOp << " [label=\"Input" << inputNum << "\"];\n";
-      for (auto userOp : input.getUsers())
-        os << "\t\t" << fromOp << " -> " << FORMAT_OPERATION(userOp) << ";\n";
-      inputNum++;
-    }
-
-    // Output
-    os << "\t\t" << FORMAT_OPERATION(concatOp) << " [label=Out];\n";
-
-    os << "\t}\n";
-
-    // Edges from output to users
-
-    for (auto& computeOpUse : concatOp->getResult(0).getUses()) {
-      os << "\t" << FORMAT_OPERATION(concatOp) << " -> "
-         << FORMAT_ARGUMENT(computeOpUse.getOwner(), computeOpUse.getOperandNumber()) << ";\n";
-    }
-  }
-
-  /**
-   * Draws the ExtractSliceOp in the graph visualization.
-   *
-   * This function takes a tensor::ExtractSliceOp and adds the corresponding
-   * node and edges to the graph visualization. It creates a node with the
-   * label as the static offsets attribute of the sliceOp, and connects it to
-   * the compute operations that use the result of the sliceOp.
-   *
-   * @param sliceOp The tensor::ExtractSliceOp to be drawn in the graph
-   * visualization.
-   */
-  void drawExtractSliceOp(tensor::ExtractSliceOp sliceOp) {
-    auto nodeId = FORMAT_ARGUMENT(sliceOp.getOperation(), 0);
-    os << "\t" << nodeId << " [label=\"Slice: ";
-    sliceOp.getStaticOffsetsAttr().print(os);
-    os << "\",color=lawngreen];\n";
-
-    for (auto& computeOpUse : sliceOp.getResult().getUses()) {
-      os << "\t" << nodeId << " -> " << FORMAT_ARGUMENT(computeOpUse.getOwner(), computeOpUse.getOperandNumber())
-         << ";\n";
-    }
-  }
-
-  void drawBiasTileOp(tensor::ExtractSliceOp sliceOp) {
-    auto nodeId = FORMAT_ARGUMENT(sliceOp.getOperation(), 0);
-    os << "\t" << nodeId << " [label=\"Bias: ";
-    sliceOp.getStaticOffsetsAttr().print(os);
-    os << "\",color=lightpink];\n";
-
-    for (auto user : sliceOp.getResult().getUsers())
-      os << "\t" << nodeId << " -> " << FORMAT_OPERATION(user) << ";\n";
-  }
-
-  /**
-   * Draws edges from the given operation to its users.
-   *
-   * @param fromOp The operation from which the edges are drawn.
-   */
-  void drawEdgesFromOpToItsUsers(mlir::Operation* fromOp) {
-    for (auto result : fromOp->getResults())
-      for (auto userOp : result.getUsers())
-        os << "\t\t" << FORMAT_OPERATION(fromOp) << " -> " << FORMAT_OPERATION(userOp) << ";\n";
-  }
-
-  /**
-   * Draws input node and edges for the given `funcOp`.
-   *
-   * @param funcOp The `funcOp` for which to draw input nodes and edges.
-   */
-  void drawInputNodesAndEdges(func::FuncOp& funcOp) {
-    os << "\tinput [label=\"Module Input\",color=green];\n";
-
-    size_t funcOpArgNum = 0;
-    for (BlockArgument& arg : funcOp.getArguments()) {
-
-      for (auto& useOp : arg.getUses()) {
-        os << "\tinput -> " << FORMAT_ARGUMENT(useOp.getOwner(), useOp.getOperandNumber()) << "[label=" << funcOpArgNum
-           << "];\n";
-      }
-      funcOpArgNum++;
-    }
-  }
-};
-
-void SpatialToGraphvizPass::runOnOperation() {
-  ModuleOp module = getOperation();
-
-  auto entryFunc = getPimEntryFunc(module);
-  if (failed(entryFunc)) {
-    module.emitError("failed to locate the PIM entry function for Spatial graph visualization");
-    signalPassFailure();
-    return;
-  }
-  func::FuncOp func = *entryFunc;
-
-  os << "digraph G {\n"
-     << "\tnode [style=filled,color=white];\n";
-
-  size_t computeNum = 0;
-  size_t concatNum = 0;
-
-  // Iterate over the ComputeOps within FuncOp:
-  // 1. Print their subgraph
-  // 2. Print the edges from its inputs to its outputs
-  for (Operation& op : func.getOps()) {
-    if (auto computeOp = dyn_cast<spatial::SpatCompute>(op)) {
-      drawComputeOpSubgraph(computeOp, computeNum++);
-    }
-    else if (auto concatOp = dyn_cast<tensor::ConcatOp>(op)) {
-      drawConcatOpSubgraph(concatOp, concatNum++);
-    }
-    else if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(op)) {
-      auto producerOp = extractSliceOp->getOperand(0).getDefiningOp();
-      if (producerOp) {
-        // Skip extractSliceOp if producer is constant weights (ONNXConstantOp)
-        if (llvm::isa<ONNXConstantOp>(producerOp))
-          continue;
-        // If produced by tosa::ReshapeOp (i.e. it is a bias tile) connect
-        // directly to its user, which is not a ComputeOp argument.
-        if (llvm::isa<tosa::ReshapeOp>(producerOp)) {
-          drawBiasTileOp(extractSliceOp);
-          continue;
-        }
-      }
-
-      drawExtractSliceOp(extractSliceOp);
-    }
-  }
-
-  // Draw input node, and edges to it users
-  drawInputNodesAndEdges(func);
-
-  // Draw output node (use the return Operation - argument number=0 - as nodeId)
-  auto returnOp = func.getBody().front().getTerminator();
-  os << '\t' << FORMAT_ARGUMENT(returnOp, 0) << " [label=\"Module Output\",color=green];\n";
-
-  os << "}\n";
-}
-
-} // namespace
-
-std::unique_ptr<Pass> createSpatialToGraphvizPass() { return std::make_unique<SpatialToGraphvizPass>(); }
-
-} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp

@@ -2,13 +2,16 @@
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/IRMapping.h"
 #include "mlir/IR/Matchers.h"
+#include <limits>
 
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Conversion/SpatialToPim/SpatialToPimPass.hpp"
 #include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
@@ -26,24 +29,6 @@ static bool isUsedOnlyAsExplicitHostOperand(Value value) {
   });
 }
 
-static FailureOr<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) {
-    auto checkedCoreId =
-      pim::checkedI32(static_cast<uint64_t>(fallbackCoreId), computeBatchOp, "fallback spatial compute_batch core id");
-    if (failed(checkedCoreId))
-      return failure();
-    coreIds.push_back(*checkedCoreId);
-    ++fallbackCoreId;
-  }
-  return coreIds;
-}
-
 static FailureOr<unsigned> getDirectReturnOperandIndex(OpResult result) {
   if (!result.hasOneUse())
     return failure();
@@ -54,6 +39,183 @@ static FailureOr<unsigned> getDirectReturnOperandIndex(OpResult result) {
   return result.getUses().begin()->getOperandNumber();
 }
 
+static FailureOr<SmallVector<FragmentAssemblyCopy, 8>>
+collectFragmentAssemblyCopiesFromBlueprint(spatial::SpatBlueprintOp blueprint,
+                                          IRMapping& mapper,
+                                          int64_t lane,
+                                          unsigned hostTargetIndex,
+                                          Value fixedSource = {}) {
+  SmallVector<FragmentAssemblyCopy, 8> copies;
+  auto resultType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
+  if (!resultType || !resultType.hasStaticShape())
+    return blueprint.emitOpError("fragment assembly lowering requires static ranked tensor results");
+
+  std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
+  std::optional<ArrayRef<int64_t>> fragmentStridesAttr = blueprint.getFragmentStrides();
+  if (!operandIndicesAttr || !fragmentStridesAttr)
+    return blueprint.emitOpError(
+      "fragment assembly lowering requires explicit operand indices and unit strides");
+
+  ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
+  std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
+  if (!sourceOffsetsAttr)
+    return blueprint.emitOpError("fragment assembly lowering requires explicit source offsets");
+  ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
+  ArrayRef<int64_t> flatOffsets = blueprint.getFragmentOffsets();
+  ArrayRef<int64_t> flatSizes = blueprint.getFragmentSizes();
+  ArrayRef<int64_t> flatStrides = *fragmentStridesAttr;
+  int64_t rank = resultType.getRank();
+
+  SmallVector<Value> fragmentOperands {blueprint.getInput()};
+  llvm::append_range(fragmentOperands, blueprint.getFragments());
+  if (failed(validateFragmentAssemblyMetadata(blueprint,
+                                              rank,
+                                              fragmentOperands.size(),
+                                              operandIndices,
+                                              sourceOffsets,
+                                              flatOffsets,
+                                              flatSizes,
+                                              flatStrides)))
+    return failure();
+
+  SmallVector<int64_t> hostStrides = computeRowMajorStrides(resultType.getShape());
+  for (int64_t fragmentIndex = 0; fragmentIndex < static_cast<int64_t>(operandIndices.size()); ++fragmentIndex) {
+    Value source = fixedSource ? fixedSource : mapper.lookupOrDefault(fragmentOperands[operandIndices[fragmentIndex]]);
+    auto sourceType = dyn_cast<RankedTensorType>(source.getType());
+    if (!sourceType || !sourceType.hasStaticShape())
+      return blueprint.emitOpError("fragment assembly lowering requires static ranked tensor operands");
+
+    size_t elementSize = getElementTypeSizeInBytes(sourceType.getElementType());
+    SmallVector<int64_t, 4> fragmentOffsets;
+    SmallVector<int64_t, 4> fragmentSizes;
+    for (int64_t dim = 0; dim < rank; ++dim) {
+      int64_t flatIndex = fragmentIndex * rank + dim;
+      if (flatStrides[flatIndex] != 1)
+        return blueprint.emitOpError("fragment assembly lowering only supports unit strides");
+      fragmentOffsets.push_back(flatOffsets[flatIndex]);
+      fragmentSizes.push_back(flatSizes[flatIndex]);
+    }
+
+    if (failed(forEachContiguousDestinationChunk(
+          resultType.getShape(),
+          fragmentOffsets,
+          fragmentSizes,
+          [&](ArrayRef<int64_t> chunkOffsets, int64_t relativeSourceOffset, int64_t chunkElements) -> LogicalResult {
+            int64_t hostElementOffset = 0;
+            for (auto [dim, offset] : llvm::enumerate(chunkOffsets))
+              hostElementOffset += offset * hostStrides[dim];
+
+            FragmentAssemblyCopy copy;
+            copy.source = source;
+            copy.sourceType = sourceType;
+            copy.hostTargetIndex = hostTargetIndex;
+            copy.lane = lane;
+            copy.sourceByteOffset = (sourceOffsets[fragmentIndex] + relativeSourceOffset) * static_cast<int64_t>(elementSize);
+            copy.hostByteOffset = hostElementOffset * static_cast<int64_t>(elementSize);
+            copy.byteSize = chunkElements * static_cast<int64_t>(elementSize);
+            copies.push_back(copy);
+            return success();
+          })))
+      return failure();
+  }
+
+  return copies;
+}
+
+static FailureOr<SmallVector<FragmentAssemblyCopy, 8>>
+collectTopLevelFragmentAssemblyCopies(OpResult result, RankedTensorType packedResultType, uint32_t laneCount) {
+  SmallVector<FragmentAssemblyCopy, 8> copies;
+  if (!packedResultType.hasStaticShape() || laneCount == 0)
+    return failure();
+
+  int64_t packedElementCount = packedResultType.getNumElements();
+  if (packedElementCount % static_cast<int64_t>(laneCount) != 0)
+    return failure();
+  int64_t payloadElementCount = packedElementCount / static_cast<int64_t>(laneCount);
+  size_t elementSize = getElementTypeSizeInBytes(packedResultType.getElementType());
+
+  for (OpOperand& use : result.getUses()) {
+    auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(use.getOwner());
+    if (!blueprint || blueprint->getParentOp() != blueprint->getParentOfType<func::FuncOp>())
+      return failure();
+    std::optional<StringRef> mode = blueprint.getMode();
+    std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
+    std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
+    if (!mode || *mode != "fragment_assembly" || !operandIndicesAttr || !sourceOffsetsAttr)
+      return failure();
+    if (!blueprint.getOutput().hasOneUse() || !isa<func::ReturnOp>(*blueprint.getOutput().getUsers().begin()))
+      return failure();
+
+    auto hostResultType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
+    std::optional<ArrayRef<int64_t>> stridesAttr = blueprint.getFragmentStrides();
+    if (!hostResultType || !hostResultType.hasStaticShape() || !stridesAttr)
+      return failure();
+
+    ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
+    ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
+    ArrayRef<int64_t> flatOffsets = blueprint.getFragmentOffsets();
+    ArrayRef<int64_t> flatSizes = blueprint.getFragmentSizes();
+    ArrayRef<int64_t> flatStrides = *stridesAttr;
+    int64_t rank = hostResultType.getRank();
+    unsigned returnIndex = blueprint.getOutput().getUses().begin()->getOperandNumber();
+    SmallVector<Value> fragmentOperands {blueprint.getInput()};
+    llvm::append_range(fragmentOperands, blueprint.getFragments());
+    if (failed(validateFragmentAssemblyMetadata(blueprint,
+                                                rank,
+                                                fragmentOperands.size(),
+                                                operandIndices,
+                                                sourceOffsets,
+                                                flatOffsets,
+                                                flatSizes,
+                                                flatStrides)))
+      return failure();
+    SmallVector<int64_t> hostStrides = computeRowMajorStrides(hostResultType.getShape());
+    for (int64_t fragmentIndex = 0; fragmentIndex < static_cast<int64_t>(operandIndices.size()); ++fragmentIndex) {
+      if (operandIndices[fragmentIndex] != static_cast<int64_t>(use.getOperandNumber()))
+        continue;
+
+      int64_t sourceElementOffset = sourceOffsets[fragmentIndex];
+      int64_t lane = sourceElementOffset / payloadElementCount;
+      if (lane < 0 || lane >= static_cast<int64_t>(laneCount))
+        return failure();
+      SmallVector<int64_t, 4> fragmentOffsets;
+      SmallVector<int64_t, 4> fragmentSizes;
+      for (int64_t dim = 0; dim < rank; ++dim) {
+        int64_t flatIndex = fragmentIndex * rank + dim;
+        if (flatStrides[flatIndex] != 1)
+          return failure();
+        fragmentOffsets.push_back(flatOffsets[flatIndex]);
+        fragmentSizes.push_back(flatSizes[flatIndex]);
+      }
+
+      if (failed(forEachContiguousDestinationChunk(
+            hostResultType.getShape(),
+            fragmentOffsets,
+            fragmentSizes,
+            [&](ArrayRef<int64_t> chunkOffsets, int64_t relativeSourceOffset, int64_t chunkElements) -> LogicalResult {
+              int64_t hostElementOffset = 0;
+              for (auto [dim, offset] : llvm::enumerate(chunkOffsets))
+                hostElementOffset += offset * hostStrides[dim];
+
+              FragmentAssemblyCopy copy;
+              copy.source = result;
+              copy.sourceType = packedResultType;
+              copy.hostTargetIndex = returnIndex;
+              copy.lane = lane;
+              copy.sourceByteOffset =
+                ((sourceElementOffset % payloadElementCount) + relativeSourceOffset) * static_cast<int64_t>(elementSize);
+              copy.hostByteOffset = hostElementOffset * static_cast<int64_t>(elementSize);
+              copy.byteSize = chunkElements * static_cast<int64_t>(elementSize);
+              copies.push_back(copy);
+              return success();
+            })))
+        return failure();
+    }
+  }
+
+  return copies;
+}
+
 static Value createScaledIndexValue(IRRewriter& rewriter, Location loc, Value base, int64_t scale) {
   if (scale == 1)
     return base;
@@ -63,25 +225,32 @@ static Value createScaledIndexValue(IRRewriter& rewriter, Location loc, Value ba
 }
 
 static Value createHostTargetOffset(IRRewriter& rewriter,
-                                    tensor::ParallelInsertSliceOp insertSlice,
+                                    Location loc,
                                     ShapedType destinationType,
+                                    ArrayRef<OpFoldResult> mixedOffsets,
+                                    ArrayRef<int64_t> additionalOffsets,
                                     IRMapping& mapper) {
   int64_t elementBytes = static_cast<int64_t>(getElementTypeSizeInBytes(destinationType.getElementType()));
   SmallVector<int64_t> strides = computeRowMajorStrides(destinationType.getShape());
 
   Value totalOffset;
-  Location loc = insertSlice.getLoc();
-  for (auto [dim, offset] : llvm::enumerate(insertSlice.getMixedOffsets())) {
+  for (auto [dim, offset] : llvm::enumerate(mixedOffsets)) {
     int64_t scale = strides[dim] * elementBytes;
     Value scaledOffset;
     if (auto attr = dyn_cast<Attribute>(offset)) {
       auto intAttr = dyn_cast<IntegerAttr>(attr);
       assert(intAttr && "expected integer offset attribute");
-      scaledOffset =
-        getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), intAttr.getInt() * scale);
-    }
-    else {
+      scaledOffset = getOrCreateIndexConstant(rewriter,
+                                              rewriter.getInsertionBlock()->getParentOp(),
+                                              (intAttr.getInt() + additionalOffsets[dim]) * scale);
+    } else {
       scaledOffset = createScaledIndexValue(rewriter, loc, mapper.lookupOrDefault(cast<Value>(offset)), scale);
+      if (additionalOffsets[dim] != 0) {
+        Value staticOffset = getOrCreateIndexConstant(rewriter,
+                                                      rewriter.getInsertionBlock()->getParentOp(),
+                                                      additionalOffsets[dim] * scale);
+        scaledOffset = arith::AddIOp::create(rewriter, loc, scaledOffset, staticOffset).getResult();
+      }
     }
 
     totalOffset =
@@ -93,9 +262,22 @@ static Value createHostTargetOffset(IRRewriter& rewriter,
   return totalOffset;
 }
 
+static Value createHostTargetOffset(IRRewriter& rewriter,
+                                    tensor::ParallelInsertSliceOp insertSlice,
+                                    ShapedType destinationType,
+                                    IRMapping& mapper) {
+  SmallVector<int64_t> zeroOffsets(destinationType.getRank(), 0);
+  return createHostTargetOffset(rewriter,
+                                insertSlice.getLoc(),
+                                destinationType,
+                                insertSlice.getMixedOffsets(),
+                                zeroOffsets,
+                                mapper);
+}
+
 } // namespace
 
-LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatComputeBatch computeBatchOp,
+LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatScheduledComputeBatch computeBatchOp,
                                                             IRRewriter& rewriter) {
   Location loc = computeBatchOp.getLoc();
   Block& oldBlock = computeBatchOp.getBody().front();
@@ -110,7 +292,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
       "resultful compute_batch lowering currently requires a spat.in_parallel terminator");
   }
 
-  auto coreIds = getPimCoreIdsForBatchOp(computeBatchOp, coreId);
+  auto coreIds = getRequiredScheduledBatchCoreIds(computeBatchOp, "spatial compute_batch core id");
   if (failed(coreIds))
     return failure();
   SmallVector<Value> batchWeights(computeBatchOp.getWeights().begin(), computeBatchOp.getWeights().end());
@@ -130,14 +312,32 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
   coreBatchOp->setAttr(onnx_mlir::kCoreIdsAttrName, rewriter.getDenseI32ArrayAttr(*coreIds));
 
   SmallVector<unsigned> returnOperandIndices;
+  SmallVector<SmallVector<FragmentAssemblyCopyRun, 1>, 4> fragmentAssemblyRunsByResult;
   if (computeBatchOp.getNumResults() != 0) {
-    returnOperandIndices.resize(computeBatchOp.getNumResults());
+    returnOperandIndices.resize(computeBatchOp.getNumResults(), std::numeric_limits<unsigned>::max());
+    fragmentAssemblyRunsByResult.resize(computeBatchOp.getNumResults());
     for (auto [resultIndex, result] : llvm::enumerate(computeBatchOp.getResults())) {
+      if (result.use_empty())
+        continue;
       FailureOr<unsigned> returnOperandIndex = getDirectReturnOperandIndex(cast<OpResult>(result));
-      if (failed(returnOperandIndex))
-        return computeBatchOp.emitOpError(
-          "resultful compute_batch lowering currently requires each result to be used directly by func.return");
+      if (succeeded(returnOperandIndex)) {
         returnOperandIndices[resultIndex] = *returnOperandIndex;
+        continue;
+      }
+
+      auto resultType = dyn_cast<RankedTensorType>(result.getType());
+      if (!resultType || !resultType.hasStaticShape())
+        return computeBatchOp.emitOpError(
+          "resultful compute_batch publication lowering requires static ranked tensor results");
+      FailureOr<SmallVector<FragmentAssemblyCopy, 8>> fragmentAssemblyCopies =
+        collectTopLevelFragmentAssemblyCopies(cast<OpResult>(result), resultType, computeBatchOp.getLaneCount());
+      if (failed(fragmentAssemblyCopies))
+        return computeBatchOp.emitOpError("failed to collect top-level fragment assembly copies for compute_batch result");
+      FailureOr<SmallVector<FragmentAssemblyCopyRun, 8>> fragmentAssemblyRuns =
+        groupFragmentAssemblyCopyRuns(*fragmentAssemblyCopies, computeBatchOp.getLaneCount());
+      if (failed(fragmentAssemblyRuns))
+        return computeBatchOp.emitOpError("failed to group top-level fragment assembly copies into regular runs");
+      fragmentAssemblyRunsByResult[resultIndex].assign(fragmentAssemblyRuns->begin(), fragmentAssemblyRuns->end());
     }
   }
 
@@ -195,6 +395,18 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
     if (isa<spatial::SpatYieldOp>(op))
       continue;
 
+    if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
+      std::optional<StringRef> modeAttr = blueprint.getMode();
+      if (modeAttr && *modeAttr == "fragment_assembly") {
+        for (Operation* user : blueprint.getOutput().getUsers()) {
+          if (!isa<tensor::ParallelInsertSliceOp>(user))
+            return blueprint.emitOpError(
+              "fragment assembly blueprint lowering expects only tensor.parallel_insert_slice users");
+        }
+        continue;
+      }
+    }
+
     if (auto parallelOp = dyn_cast<spatial::SpatInParallelOp>(op)) {
       auto firstOutputArg = computeBatchOp.getOutputArgument(0);
       if (!firstOutputArg)
@@ -211,10 +423,75 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
         unsigned resultIndex = outputArg.getArgNumber() - firstOutputArg->getArgNumber();
         if (resultIndex >= returnOperandIndices.size())
           return insertSlice.emitOpError("result index out of range while lowering host batch output");
+        bool hasDirectReturn = returnOperandIndices[resultIndex] != std::numeric_limits<unsigned>::max();
+        bool hasFragmentAssembly = resultIndex < fragmentAssemblyRunsByResult.size()
+                                   && !fragmentAssemblyRunsByResult[resultIndex].empty();
+        if (!hasDirectReturn && !hasFragmentAssembly)
+          continue;
 
         Value mappedSource = mapper.lookup(insertSlice.getSource());
+
+        if (hasFragmentAssembly) {
+          BlockArgument laneArg = coreBatchOp.getLaneArgument();
+          auto mappedSourceType = dyn_cast<ShapedType>(mappedSource.getType());
+          if (!mappedSourceType || !mappedSourceType.hasStaticShape())
+            return insertSlice.emitOpError("fragment assembly batch lowering requires a static ranked lane-local source");
+          DenseMap<unsigned, Value> updatedOutputs;
+          for (const FragmentAssemblyCopyRun& run : fragmentAssemblyRunsByResult[resultIndex]) {
+            Value outputTensor = updatedOutputs.lookup(run.hostTargetIndex);
+            if (!outputTensor)
+              outputTensor = outputTensors[run.hostTargetIndex](rewriter, insertSlice.getLoc());
+            FragmentAssemblyCopyRun mappedRun = run;
+            mappedRun.source = mappedSource;
+            FailureOr<Value> updated =
+              emitFragmentAssemblyCopyRuns(rewriter,
+                                           insertSlice.getLoc(),
+                                           ArrayRef<FragmentAssemblyCopyRun> {mappedRun},
+                                           outputTensor,
+                                           coreBatchOp.getOperation(),
+                                           laneArg);
+            if (failed(updated))
+              return failure();
+            updatedOutputs[run.hostTargetIndex] = *updated;
+          }
+          continue;
+        }
+
         Value hostTarget = getOrCreateHostOutputTensor(resultIndex, insertSlice.getLoc());
         auto hostTargetType = cast<ShapedType>(hostTarget.getType());
+        if (auto blueprint =
+              insertSlice.getSource().getDefiningOp<spatial::SpatBlueprintOp>()) {
+          std::optional<StringRef> modeAttr = blueprint.getMode();
+          if (modeAttr && *modeAttr == "fragment_assembly") {
+            FailureOr<SmallVector<FragmentAssemblyCopy, 8>> fragmentAssemblyCopies =
+              collectFragmentAssemblyCopiesFromBlueprint(blueprint, mapper, /*lane=*/0, /*hostTargetIndex=*/0);
+            if (failed(fragmentAssemblyCopies))
+              return failure();
+            FailureOr<SmallVector<FragmentAssemblyCopyRun, 8>> fragmentAssemblyRuns =
+              groupFragmentAssemblyCopyRuns(*fragmentAssemblyCopies, /*laneCount=*/1);
+            if (failed(fragmentAssemblyRuns))
+              return failure();
+            SmallVector<int64_t> zeroOffsets(hostTargetType.getRank(), 0);
+            Value baseHostOffset = createHostTargetOffset(rewriter,
+                                                         blueprint.getLoc(),
+                                                         hostTargetType,
+                                                         insertSlice.getMixedOffsets(),
+                                                         zeroOffsets,
+                                                         mapper);
+            FailureOr<Value> updatedHostTarget = emitFragmentAssemblyCopyRuns(rewriter,
+                                                                              blueprint.getLoc(),
+                                                                              *fragmentAssemblyRuns,
+                                                                              hostTarget,
+                                                                              coreBatchOp.getOperation(),
+                                                                              std::nullopt,
+                                                                              baseHostOffset);
+            if (failed(updatedHostTarget))
+              return failure();
+            hostOutputTensors[resultIndex] = *updatedHostTarget;
+            continue;
+          }
+        }
+
         Value hostTargetOffset = createHostTargetOffset(rewriter, insertSlice, hostTargetType, mapper);
         Value zeroOffset = getOrCreateIndexConstant(rewriter, coreBatchOp.getOperation(), 0);
         auto sizeAttr = getTensorSizeInBytesAttr(rewriter, coreBatchOp.getOperation(), mappedSource);
@@ -264,9 +541,15 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatCompute
       Operation* definingOp = operand.getDefiningOp();
       if (definingOp && definingOp->getBlock() == &oldBlock)
         continue;
+      if (definingOp && definingOp->hasTrait<OpTrait::ConstantLike>())
+        continue;
 
-      return computeBatchOp.emitOpError(
-        "expected external tensor communication to be materialized in Spatial before batch lowering");
+      InFlightDiagnostic diagnostic =
+        computeBatchOp.emitOpError("expected external tensor communication to be materialized in Spatial before batch lowering");
+      diagnostic << " while cloning nested op '" << op.getName() << "' tensor operand #" << operandIndex;
+      if (definingOp)
+        diagnostic << " from external producer '" << definingOp->getName() << "'";
+      return diagnostic;
     }
 
     Operation* cloned = rewriter.clone(op, mapper);

src/PIM/Conversion/SpatialToPim/Common.cpp

@@ -1,11 +1,17 @@
 #include "mlir/IR/ValueRange.h"
 
+#include "mlir/Dialect/Arith/IR/Arith.h"
+
 #include "llvm/ADT/STLExtras.h"
 
 #include <cassert>
+#include <limits>
 
 #include "Common.hpp"
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
+#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
 using namespace llvm;
 using namespace mlir;
@@ -72,4 +78,488 @@ mlir::Value getBestOutputTensorFromOperandsOrAllocate(RewriterBase& rewriter, Op
     rewriter, operation->getLoc(), resultShapedType.getShape(), resultShapedType.getElementType());
 }
 
+LogicalResult validateFragmentAssemblyMetadata(spatial::SpatBlueprintOp blueprint,
+                                               int64_t resultRank,
+                                               size_t operandCount,
+                                               ArrayRef<int64_t> operandIndices,
+                                               ArrayRef<int64_t> sourceOffsets,
+                                               ArrayRef<int64_t> flatOffsets,
+                                               ArrayRef<int64_t> flatSizes,
+                                               ArrayRef<int64_t> flatStrides) {
+  if (operandIndices.size() != sourceOffsets.size())
+    return blueprint.emitOpError("fragment assembly operand index and source offset counts must match");
+  if (flatOffsets.size() != flatSizes.size())
+    return blueprint.emitOpError("fragment assembly offset and size arrays must have matching lengths");
+  if (flatStrides.size() != flatOffsets.size())
+    return blueprint.emitOpError("fragment assembly stride and offset arrays must have matching lengths");
+  if (flatOffsets.size() != operandIndices.size() * static_cast<size_t>(resultRank))
+    return blueprint.emitOpError("fragment assembly metadata must provide one rank-sized offset/size/stride tuple per fragment");
+
+  for (auto [fragmentIndex, operandIndex] : llvm::enumerate(operandIndices)) {
+    if (operandIndex < 0 || operandIndex >= static_cast<int64_t>(operandCount))
+      return blueprint.emitOpError("fragment assembly operand index is out of range");
+    if (sourceOffsets[fragmentIndex] < 0)
+      return blueprint.emitOpError("fragment assembly source offsets must be nonnegative");
+  }
+
+  return success();
+}
+
+static SmallVector<int64_t, 4> expandFlatElementIndex(int64_t flatIndex, ArrayRef<int64_t> shape) {
+  SmallVector<int64_t, 4> 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;
+}
+
+FailureOr<SmallVector<int64_t, 4>>
+getStaticSliceOffsetsForElementOffset(Operation* anchor,
+                                      ShapedType sourceType,
+                                      ArrayRef<int64_t> fragmentShape,
+                                      int64_t sourceElementOffset,
+                                      StringRef fieldName) {
+  if (!sourceType.hasStaticShape())
+    return (anchor->emitOpError() << fieldName << " requires a static source shape"), failure();
+  if (sourceElementOffset < 0)
+    return (anchor->emitOpError() << fieldName << " requires a nonnegative source element offset"), failure();
+  if (sourceType.getRank() != static_cast<int64_t>(fragmentShape.size()))
+    return (anchor->emitOpError() << fieldName << " requires fragment rank to match source rank"), failure();
+
+  int64_t sourceElementCount = sourceType.getNumElements();
+  int64_t fragmentElementCount = 1;
+  for (int64_t dim = 0; dim < sourceType.getRank(); ++dim) {
+    if (fragmentShape[dim] < 0)
+      return (anchor->emitOpError() << fieldName << " requires nonnegative fragment sizes"), failure();
+    fragmentElementCount *= fragmentShape[dim];
+  }
+  if (sourceElementOffset + fragmentElementCount > sourceElementCount)
+    return (anchor->emitOpError() << fieldName << " exceeds the source tensor bounds"), failure();
+
+  SmallVector<int64_t, 4> sliceOffsets = expandFlatElementIndex(sourceElementOffset, sourceType.getShape());
+  for (int64_t dim = 0; dim < sourceType.getRank(); ++dim) {
+    if (sliceOffsets[dim] + fragmentShape[dim] > sourceType.getDimSize(dim))
+      return (anchor->emitOpError() << fieldName << " does not describe a valid unit-stride slice"), failure();
+  }
+  return sliceOffsets;
+}
+
+LogicalResult
+forEachContiguousDestinationChunk(ArrayRef<int64_t> destShape,
+                                  ArrayRef<int64_t> baseOffsets,
+                                  ArrayRef<int64_t> sizes,
+                                  llvm::function_ref<LogicalResult(ArrayRef<int64_t>, int64_t, int64_t)> callback) {
+  int64_t rank = static_cast<int64_t>(sizes.size());
+  int64_t suffixStart = rank - 1;
+  while (suffixStart > 0 && sizes[suffixStart] == destShape[suffixStart])
+    --suffixStart;
+  if (sizes[suffixStart] == destShape[suffixStart] && suffixStart == 0)
+    suffixStart = 0;
+  else
+    ++suffixStart;
+
+  int64_t chunkElements = 1;
+  for (int64_t dim = suffixStart; dim < rank; ++dim)
+    chunkElements *= sizes[dim];
+
+  SmallVector<int64_t, 4> prefixExtents(sizes.begin(), sizes.begin() + suffixStart);
+  SmallVector<int64_t, 4> current(prefixExtents.size(), 0);
+  int64_t sourceChunkOrdinal = 0;
+
+  auto visit = [&](auto&& visit, int64_t dim) -> LogicalResult {
+    if (dim == static_cast<int64_t>(prefixExtents.size())) {
+      SmallVector<int64_t, 4> chunkOffsets(baseOffsets.begin(), baseOffsets.end());
+      for (int64_t prefixDim = 0; prefixDim < static_cast<int64_t>(current.size()); ++prefixDim)
+        chunkOffsets[prefixDim] += current[prefixDim];
+      if (failed(callback(chunkOffsets, sourceChunkOrdinal * chunkElements, chunkElements)))
+        return failure();
+      ++sourceChunkOrdinal;
+      return success();
+    }
+
+    for (int64_t index = 0; index < prefixExtents[dim]; ++index) {
+      current[dim] = index;
+      if (failed(visit(visit, dim + 1)))
+        return failure();
+    }
+    return success();
+  };
+
+  if (prefixExtents.empty())
+    return callback(baseOffsets, 0, chunkElements);
+  return visit(visit, 0);
+}
+
+static mlir::Value
+createSteppedOffset(OpBuilder& builder, Location loc, mlir::Value start, mlir::Value index, int64_t stepBytes) {
+  if (stepBytes == 0)
+    return start;
+  mlir::Value step = arith::ConstantIndexOp::create(builder, loc, stepBytes);
+  mlir::Value scaled = arith::MulIOp::create(builder, loc, index, step).getResult();
+  return arith::AddIOp::create(builder, loc, start, scaled).getResult();
+}
+
+static mlir::Value createIndexedOffset(OpBuilder& builder,
+                                       Location loc,
+                                       mlir::Value indexArg,
+                                       ArrayRef<int64_t> values) {
+  assert(!values.empty() && "expected lane-indexed values");
+  if (llvm::all_of(values.drop_front(), [&](int64_t value) { return value == values.front(); }))
+    return arith::ConstantIndexOp::create(builder, loc, values.front());
+
+  if (values.size() >= 2) {
+    int64_t step = values[1] - values[0];
+    bool arithmetic = llvm::all_of(llvm::seq<size_t>(2, values.size()), [&](size_t index) {
+      return values[index] == values.front() + static_cast<int64_t>(index) * step;
+    });
+    if (arithmetic) {
+      mlir::Value base = arith::ConstantIndexOp::create(builder, loc, values.front());
+      mlir::Value stepValue = arith::ConstantIndexOp::create(builder, loc, step);
+      mlir::Value scaledIndex = arith::MulIOp::create(builder, loc, indexArg, stepValue).getResult();
+      return arith::AddIOp::create(builder, loc, base, scaledIndex).getResult();
+    }
+  }
+
+  mlir::Value selected = arith::ConstantIndexOp::create(builder, loc, values.front());
+  for (auto [lane, value] : llvm::enumerate(values.drop_front())) {
+    mlir::Value indexValue = arith::ConstantIndexOp::create(builder, loc, static_cast<int64_t>(lane + 1));
+    mlir::Value cmp = arith::CmpIOp::create(builder, loc, arith::CmpIPredicate::eq, indexArg, indexValue);
+    mlir::Value candidate = arith::ConstantIndexOp::create(builder, loc, value);
+    selected = arith::SelectOp::create(builder, loc, cmp, candidate, selected);
+  }
+  return selected;
+}
+
+struct FragmentAssemblyCopyRunFamily {
+  FragmentAssemblyCopyRun prototype;
+  SmallVector<int64_t, 8> sourceRunStartDeltas;
+  SmallVector<int64_t, 8> hostRunStartDeltas;
+};
+
+static bool computeUniformRunStartDelta(ArrayRef<int64_t> prototypeStarts,
+                                        ArrayRef<int64_t> runStarts,
+                                        int64_t& delta) {
+  if (prototypeStarts.size() != runStarts.size() || prototypeStarts.empty())
+    return false;
+
+  delta = runStarts.front() - prototypeStarts.front();
+  return llvm::all_of(llvm::zip_equal(prototypeStarts, runStarts), [&](auto pair) {
+    auto [prototypeStart, runStart] = pair;
+    return runStart - prototypeStart == delta;
+  });
+}
+
+static bool canMergeFragmentAssemblyCopyRunIntoFamily(const FragmentAssemblyCopyRunFamily& family,
+                                                      const FragmentAssemblyCopyRun& run,
+                                                      int64_t& sourceRunStartDelta,
+                                                      int64_t& hostRunStartDelta) {
+  const FragmentAssemblyCopyRun& prototype = family.prototype;
+  if (prototype.source != run.source || prototype.sourceType != run.sourceType
+      || prototype.hostTargetIndex != run.hostTargetIndex || prototype.count != run.count
+      || prototype.sourceStepBytes != run.sourceStepBytes || prototype.hostStepBytes != run.hostStepBytes
+      || prototype.byteSize != run.byteSize)
+    return false;
+
+  if (!computeUniformRunStartDelta(prototype.sourceStartBytesByLane, run.sourceStartBytesByLane, sourceRunStartDelta))
+    return false;
+  return computeUniformRunStartDelta(prototype.hostStartBytesByLane, run.hostStartBytesByLane, hostRunStartDelta);
+}
+
+static SmallVector<FragmentAssemblyCopyRunFamily, 8>
+groupFragmentAssemblyCopyRunFamilies(ArrayRef<FragmentAssemblyCopyRun> runs) {
+  auto compareRunStarts = [](ArrayRef<int64_t> lhs, ArrayRef<int64_t> rhs) {
+    return std::lexicographical_compare(lhs.begin(), lhs.end(), rhs.begin(), rhs.end());
+  };
+
+  SmallVector<FragmentAssemblyCopyRun, 8> sortedRuns(runs.begin(), runs.end());
+  llvm::sort(sortedRuns, [&](const FragmentAssemblyCopyRun& lhs, const FragmentAssemblyCopyRun& rhs) {
+    if (lhs.hostTargetIndex != rhs.hostTargetIndex)
+      return lhs.hostTargetIndex < rhs.hostTargetIndex;
+    if (lhs.source != rhs.source)
+      return lhs.source.getAsOpaquePointer() < rhs.source.getAsOpaquePointer();
+    if (lhs.byteSize != rhs.byteSize)
+      return lhs.byteSize < rhs.byteSize;
+    if (lhs.count != rhs.count)
+      return lhs.count < rhs.count;
+    if (lhs.sourceStepBytes != rhs.sourceStepBytes)
+      return lhs.sourceStepBytes < rhs.sourceStepBytes;
+    if (lhs.hostStepBytes != rhs.hostStepBytes)
+      return lhs.hostStepBytes < rhs.hostStepBytes;
+    if (compareRunStarts(lhs.sourceStartBytesByLane, rhs.sourceStartBytesByLane))
+      return true;
+    if (compareRunStarts(rhs.sourceStartBytesByLane, lhs.sourceStartBytesByLane))
+      return false;
+    return compareRunStarts(lhs.hostStartBytesByLane, rhs.hostStartBytesByLane);
+  });
+
+  SmallVector<FragmentAssemblyCopyRunFamily, 8> families;
+  for (const FragmentAssemblyCopyRun& run : sortedRuns) {
+    int64_t sourceRunStartDelta = 0;
+    int64_t hostRunStartDelta = 0;
+    if (!families.empty()
+        && canMergeFragmentAssemblyCopyRunIntoFamily(
+          families.back(), run, sourceRunStartDelta, hostRunStartDelta)) {
+      families.back().sourceRunStartDeltas.push_back(sourceRunStartDelta);
+      families.back().hostRunStartDeltas.push_back(hostRunStartDelta);
+      continue;
+    }
+
+    FragmentAssemblyCopyRunFamily family;
+    family.prototype = run;
+    family.sourceRunStartDeltas.push_back(0);
+    family.hostRunStartDeltas.push_back(0);
+    families.push_back(std::move(family));
+  }
+
+  return families;
+}
+
+FailureOr<SmallVector<FragmentAssemblyCopyRun, 8>>
+groupFragmentAssemblyCopyRuns(ArrayRef<FragmentAssemblyCopy> copies, uint32_t laneCount) {
+  if (laneCount == 0)
+    return failure();
+
+  struct LaneLocalCopyRun {
+    FragmentAssemblyCopyRun run;
+    int64_t lane = 0;
+  };
+
+  SmallVector<FragmentAssemblyCopy, 8> sortedCopies(copies.begin(), copies.end());
+  llvm::sort(sortedCopies, [](const FragmentAssemblyCopy& lhs, const FragmentAssemblyCopy& rhs) {
+    if (lhs.hostTargetIndex != rhs.hostTargetIndex)
+      return lhs.hostTargetIndex < rhs.hostTargetIndex;
+    if (lhs.source != rhs.source)
+      return lhs.source.getAsOpaquePointer() < rhs.source.getAsOpaquePointer();
+    if (lhs.lane != rhs.lane)
+      return lhs.lane < rhs.lane;
+    if (lhs.byteSize != rhs.byteSize)
+      return lhs.byteSize < rhs.byteSize;
+    if (lhs.sourceByteOffset != rhs.sourceByteOffset)
+      return lhs.sourceByteOffset < rhs.sourceByteOffset;
+    return lhs.hostByteOffset < rhs.hostByteOffset;
+  });
+
+  SmallVector<LaneLocalCopyRun, 8> laneRuns;
+  for (const FragmentAssemblyCopy& copy : sortedCopies) {
+    if (copy.lane < 0 || copy.lane >= static_cast<int64_t>(laneCount))
+      return failure();
+
+    if (!laneRuns.empty()) {
+      LaneLocalCopyRun& laneRun = laneRuns.back();
+      FragmentAssemblyCopyRun& run = laneRun.run;
+      if (run.source == copy.source && run.sourceType == copy.sourceType
+          && run.hostTargetIndex == copy.hostTargetIndex && laneRun.lane == copy.lane && run.byteSize == copy.byteSize
+          && run.sourceStartBytesByLane.size() == 1 && run.hostStartBytesByLane.size() == 1) {
+        int64_t previousSourceOffset = run.sourceStartBytesByLane.front() + (run.count - 1) * run.sourceStepBytes;
+        int64_t previousHostOffset = run.hostStartBytesByLane.front() + (run.count - 1) * run.hostStepBytes;
+        int64_t sourceDelta = copy.sourceByteOffset - previousSourceOffset;
+        int64_t hostDelta = copy.hostByteOffset - previousHostOffset;
+        if (run.count == 1) {
+          run.sourceStepBytes = sourceDelta;
+          run.hostStepBytes = hostDelta;
+          ++run.count;
+          continue;
+        }
+        if (run.sourceStepBytes == sourceDelta && run.hostStepBytes == hostDelta) {
+          ++run.count;
+          continue;
+        }
+      }
+    }
+
+    LaneLocalCopyRun laneRun;
+    laneRun.run.source = copy.source;
+    laneRun.run.sourceType = copy.sourceType;
+    laneRun.run.hostTargetIndex = copy.hostTargetIndex;
+    laneRun.run.count = 1;
+    laneRun.run.byteSize = copy.byteSize;
+    laneRun.run.sourceStartBytesByLane.push_back(copy.sourceByteOffset);
+    laneRun.run.hostStartBytesByLane.push_back(copy.hostByteOffset);
+    laneRun.lane = copy.lane;
+    laneRuns.push_back(std::move(laneRun));
+  }
+
+  SmallVector<FragmentAssemblyCopyRun, 8> mergedRuns;
+  for (const LaneLocalCopyRun& laneRun : laneRuns) {
+    size_t laneIndex = static_cast<size_t>(laneRun.lane);
+    auto mergedIt = llvm::find_if(mergedRuns, [&](const FragmentAssemblyCopyRun& run) {
+      return run.source == laneRun.run.source && run.sourceType == laneRun.run.sourceType
+             && run.hostTargetIndex == laneRun.run.hostTargetIndex && run.count == laneRun.run.count
+             && run.byteSize == laneRun.run.byteSize && run.sourceStepBytes == laneRun.run.sourceStepBytes
+             && run.hostStepBytes == laneRun.run.hostStepBytes && laneIndex < run.sourceStartBytesByLane.size()
+             && run.sourceStartBytesByLane[laneIndex] == std::numeric_limits<int64_t>::min();
+    });
+
+    if (mergedIt == mergedRuns.end()) {
+      FragmentAssemblyCopyRun merged = laneRun.run;
+      merged.sourceStartBytesByLane.assign(laneCount, std::numeric_limits<int64_t>::min());
+      merged.hostStartBytesByLane.assign(laneCount, std::numeric_limits<int64_t>::min());
+      merged.sourceStartBytesByLane[laneIndex] = laneRun.run.sourceStartBytesByLane.front();
+      merged.hostStartBytesByLane[laneIndex] = laneRun.run.hostStartBytesByLane.front();
+      mergedRuns.push_back(std::move(merged));
+      continue;
+    }
+
+    mergedIt->sourceStartBytesByLane[laneIndex] = laneRun.run.sourceStartBytesByLane.front();
+    mergedIt->hostStartBytesByLane[laneIndex] = laneRun.run.hostStartBytesByLane.front();
+  }
+
+  for (const FragmentAssemblyCopyRun& run : mergedRuns) {
+    if (llvm::any_of(run.sourceStartBytesByLane,
+                     [](int64_t value) { return value == std::numeric_limits<int64_t>::min(); }))
+      return failure();
+    if (llvm::any_of(run.hostStartBytesByLane,
+                     [](int64_t value) { return value == std::numeric_limits<int64_t>::min(); }))
+      return failure();
+  }
+
+  return mergedRuns;
+}
+
+static FailureOr<mlir::Value> emitFragmentAssemblyCopyRun(OpBuilder& builder,
+                                                          Location loc,
+                                                          const FragmentAssemblyCopyRun& run,
+                                                          mlir::Value hostTarget,
+                                                          Operation* anchor,
+                                                          std::optional<mlir::Value> laneArg,
+                                                          mlir::Value baseHostOffset,
+                                                          mlir::Value sourceRunStartDelta = {},
+                                                          mlir::Value hostRunStartDelta = {}) {
+  auto sizeAttr = pim::getCheckedI32Attr(builder, anchor, run.byteSize, "fragment assembly host copy byte size");
+  if (failed(sizeAttr))
+    return failure();
+
+  mlir::Value hostStart;
+  mlir::Value sourceStart;
+  if (laneArg) {
+    hostStart = createIndexedOffset(builder, loc, *laneArg, run.hostStartBytesByLane);
+    sourceStart = createIndexedOffset(builder, loc, *laneArg, run.sourceStartBytesByLane);
+  } else {
+    hostStart = arith::ConstantIndexOp::create(builder, loc, run.hostStartBytesByLane.front());
+    sourceStart = arith::ConstantIndexOp::create(builder, loc, run.sourceStartBytesByLane.front());
+  }
+
+  if (hostRunStartDelta)
+    hostStart = arith::AddIOp::create(builder, loc, hostStart, hostRunStartDelta).getResult();
+  if (sourceRunStartDelta)
+    sourceStart = arith::AddIOp::create(builder, loc, sourceStart, sourceRunStartDelta).getResult();
+  if (baseHostOffset)
+    hostStart = arith::AddIOp::create(builder, loc, baseHostOffset, hostStart).getResult();
+
+  if (run.count == 1) {
+    return pim::PimMemCopyDevToHostOp::create(builder,
+                                              loc,
+                                              hostTarget.getType(),
+                                              hostStart,
+                                              sourceStart,
+                                              hostTarget,
+                                              run.source,
+                                              *sizeAttr)
+      .getOutput();
+  }
+
+  mlir::Value lowerBound = arith::ConstantIndexOp::create(builder, loc, 0);
+  mlir::Value upperBound = arith::ConstantIndexOp::create(builder, loc, run.count);
+  mlir::Value step = arith::ConstantIndexOp::create(builder, loc, 1);
+  FailureOr<NormalizedLoopResult> loop = buildNormalizedScfFor(
+    builder,
+    loc,
+    lowerBound,
+    upperBound,
+    step,
+    ValueRange {hostTarget},
+    [&](OpBuilder& loopBuilder,
+        Location bodyLoc,
+        mlir::Value flatIndex,
+        ValueRange iterArgs,
+        SmallVectorImpl<mlir::Value>& yielded) {
+      mlir::Value hostOffset = createSteppedOffset(loopBuilder, bodyLoc, hostStart, flatIndex, run.hostStepBytes);
+      mlir::Value sourceOffset =
+        createSteppedOffset(loopBuilder, bodyLoc, sourceStart, flatIndex, run.sourceStepBytes);
+      mlir::Value copied =
+        pim::PimMemCopyDevToHostOp::create(loopBuilder,
+                                           bodyLoc,
+                                           iterArgs.front().getType(),
+                                           hostOffset,
+                                           sourceOffset,
+                                           iterArgs.front(),
+                                           run.source,
+                                           *sizeAttr)
+          .getOutput();
+      yielded.push_back(copied);
+      return success();
+    });
+  if (failed(loop))
+    return failure();
+  return loop->results.front();
+}
+
+static FailureOr<mlir::Value> emitFragmentAssemblyCopyRunFamily(OpBuilder& builder,
+                                                                Location loc,
+                                                                const FragmentAssemblyCopyRunFamily& family,
+                                                                mlir::Value hostTarget,
+                                                                Operation* anchor,
+                                                                std::optional<mlir::Value> laneArg,
+                                                                mlir::Value baseHostOffset) {
+  if (family.sourceRunStartDeltas.size() == 1)
+    return emitFragmentAssemblyCopyRun(
+      builder, loc, family.prototype, hostTarget, anchor, laneArg, baseHostOffset);
+
+  mlir::Value lowerBound = arith::ConstantIndexOp::create(builder, loc, 0);
+  mlir::Value upperBound = arith::ConstantIndexOp::create(builder, loc, family.sourceRunStartDeltas.size());
+  mlir::Value step = arith::ConstantIndexOp::create(builder, loc, 1);
+  FailureOr<NormalizedLoopResult> outerLoop = buildNormalizedScfFor(
+    builder,
+    loc,
+    lowerBound,
+    upperBound,
+    step,
+    ValueRange {hostTarget},
+    [&](OpBuilder& loopBuilder,
+        Location bodyLoc,
+        mlir::Value runIndex,
+        ValueRange iterArgs,
+        SmallVectorImpl<mlir::Value>& yielded) {
+      mlir::Value sourceRunStartDelta =
+        createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.sourceRunStartDeltas);
+      mlir::Value hostRunStartDelta =
+        createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.hostRunStartDeltas);
+      FailureOr<mlir::Value> copied = emitFragmentAssemblyCopyRun(loopBuilder,
+                                                                  bodyLoc,
+                                                                  family.prototype,
+                                                                  iterArgs.front(),
+                                                                  anchor,
+                                                                  laneArg,
+                                                                  baseHostOffset,
+                                                                  sourceRunStartDelta,
+                                                                  hostRunStartDelta);
+      if (failed(copied))
+        return failure();
+      yielded.push_back(*copied);
+      return success();
+    });
+  if (failed(outerLoop))
+    return failure();
+  return outerLoop->results.front();
+}
+
+FailureOr<mlir::Value> emitFragmentAssemblyCopyRuns(IRRewriter& rewriter,
+                                                    Location loc,
+                                                    ArrayRef<FragmentAssemblyCopyRun> runs,
+                                                    mlir::Value hostTarget,
+                                                    Operation* anchor,
+                                                    std::optional<mlir::Value> laneArg,
+                                                    mlir::Value baseHostOffset) {
+  for (const FragmentAssemblyCopyRunFamily& family : groupFragmentAssemblyCopyRunFamilies(runs)) {
+    FailureOr<mlir::Value> updatedHostTarget =
+      emitFragmentAssemblyCopyRunFamily(rewriter, loc, family, hostTarget, anchor, laneArg, baseHostOffset);
+    if (failed(updatedHostTarget))
+      return failure();
+    hostTarget = *updatedHostTarget;
+  }
+
+  return hostTarget;
+}
+
 } // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/Common.hpp

@@ -1,10 +1,23 @@
 #pragma once
 
+#include <optional>
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/STLFunctionalExtras.h"
+
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/Value.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Support/LogicalResult.h"
 
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 
+namespace onnx_mlir::spatial {
+class SpatBlueprintOp;
+}
+
 namespace onnx_mlir {
 
 mlir::FailureOr<mlir::IntegerAttr>
@@ -29,6 +42,62 @@ mlir::SmallVector<mlir::Value> getOpOperandsSortedByUses(mlir::Operation* operat
 
 mlir::Value getBestOutputTensorFromOperandsOrAllocate(mlir::RewriterBase& rewriter, mlir::Operation* operation);
 
+mlir::LogicalResult validateFragmentAssemblyMetadata(onnx_mlir::spatial::SpatBlueprintOp blueprint,
+                                                     int64_t resultRank,
+                                                     size_t operandCount,
+                                                     llvm::ArrayRef<int64_t> operandIndices,
+                                                     llvm::ArrayRef<int64_t> sourceOffsets,
+                                                     llvm::ArrayRef<int64_t> flatOffsets,
+                                                     llvm::ArrayRef<int64_t> flatSizes,
+                                                     llvm::ArrayRef<int64_t> flatStrides);
+
+mlir::FailureOr<mlir::SmallVector<int64_t, 4>>
+getStaticSliceOffsetsForElementOffset(mlir::Operation* anchor,
+                                      mlir::ShapedType sourceType,
+                                      llvm::ArrayRef<int64_t> fragmentShape,
+                                      int64_t sourceElementOffset,
+                                      llvm::StringRef fieldName);
+
+mlir::LogicalResult
+forEachContiguousDestinationChunk(llvm::ArrayRef<int64_t> destShape,
+                                  llvm::ArrayRef<int64_t> baseOffsets,
+                                  llvm::ArrayRef<int64_t> sizes,
+                                  llvm::function_ref<mlir::LogicalResult(llvm::ArrayRef<int64_t>, int64_t, int64_t)>
+                                    callback);
+
+struct FragmentAssemblyCopy {
+  mlir::Value source;
+  mlir::RankedTensorType sourceType;
+  unsigned hostTargetIndex = 0;
+  int64_t lane = 0;
+  int64_t sourceByteOffset = 0;
+  int64_t hostByteOffset = 0;
+  int64_t byteSize = 0;
+};
+
+struct FragmentAssemblyCopyRun {
+  mlir::Value source;
+  mlir::RankedTensorType sourceType;
+  unsigned hostTargetIndex = 0;
+  int64_t count = 0;
+  int64_t sourceStepBytes = 0;
+  int64_t hostStepBytes = 0;
+  int64_t byteSize = 0;
+  mlir::SmallVector<int64_t, 8> sourceStartBytesByLane;
+  mlir::SmallVector<int64_t, 8> hostStartBytesByLane;
+};
+
+mlir::FailureOr<mlir::SmallVector<FragmentAssemblyCopyRun, 8>>
+groupFragmentAssemblyCopyRuns(llvm::ArrayRef<FragmentAssemblyCopy> copies, uint32_t laneCount = 1);
+
+mlir::FailureOr<mlir::Value> emitFragmentAssemblyCopyRuns(mlir::IRRewriter& rewriter,
+                                                          mlir::Location loc,
+                                                          llvm::ArrayRef<FragmentAssemblyCopyRun> runs,
+                                                          mlir::Value hostTarget,
+                                                          mlir::Operation* anchor,
+                                                          std::optional<mlir::Value> laneArg = std::nullopt,
+                                                          mlir::Value baseHostOffset = {});
+
 inline mlir::tensor::EmptyOp
 createEmptyTensorFromShaped(mlir::IRRewriter& rewriter, mlir::Location loc, mlir::ShapedType shapedType) {
   return mlir::tensor::EmptyOp::create(rewriter, loc, shapedType.getShape(), shapedType.getElementType());

src/PIM/Conversion/SpatialToPim/ComputeLikeRegionUtils.cpp

@@ -17,10 +17,10 @@ std::optional<unsigned> getDirectComputeLikeInputIndex(Operation* owner, unsigne
     return operandNumber - inputBegin;
   };
 
-  if (auto compute = dyn_cast<spatial::SpatCompute>(owner))
+  if (auto compute = dyn_cast<spatial::SpatScheduledCompute>(owner))
     return getInputIndex(owner, compute.getInputs().size());
 
-  if (auto computeBatch = dyn_cast<spatial::SpatComputeBatch>(owner))
+  if (auto computeBatch = dyn_cast<spatial::SpatScheduledComputeBatch>(owner))
     return getInputIndex(owner, computeBatch.getInputs().size());
 
   return std::nullopt;
@@ -32,13 +32,13 @@ void replaceAndEraseDirectComputeLikeInput(PatternRewriter& rewriter,
                                            Value replacement) {
   Block& body = owner->getRegion(0).front();
   BlockArgument bodyArgument;
-  if (auto compute = dyn_cast<spatial::SpatCompute>(owner)) {
+  if (auto compute = dyn_cast<spatial::SpatScheduledCompute>(owner)) {
     auto computeArg = compute.getInputArgument(inputIndex);
     assert(computeArg && "expected compute input block argument");
     bodyArgument = *computeArg;
   }
   else {
-    auto batchArg = cast<spatial::SpatComputeBatch>(owner).getInputArgument(inputIndex);
+    auto batchArg = cast<spatial::SpatScheduledComputeBatch>(owner).getInputArgument(inputIndex);
     assert(batchArg && "expected compute_batch input block argument");
     bodyArgument = *batchArg;
   }
@@ -46,10 +46,10 @@ void replaceAndEraseDirectComputeLikeInput(PatternRewriter& rewriter,
 
   rewriter.startOpModification(owner);
   bodyArgument.replaceAllUsesWith(replacement);
-  if (auto compute = dyn_cast<spatial::SpatCompute>(owner))
+  if (auto compute = dyn_cast<spatial::SpatScheduledCompute>(owner))
     compute.getInputsMutable().erase(inputIndex);
   else
-    cast<spatial::SpatComputeBatch>(owner).getInputsMutable().erase(inputIndex);
+    cast<spatial::SpatScheduledComputeBatch>(owner).getInputsMutable().erase(inputIndex);
   body.eraseArgument(bodyArgIndex);
   rewriter.finalizeOpModification(owner);
 }

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp

@@ -1,6 +1,7 @@
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/IR/IRMapping.h"
@@ -8,6 +9,8 @@
 
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Conversion/SpatialToPim/SpatialToPimPass.hpp"
+#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
@@ -30,6 +33,90 @@ static bool isChannelUseChainOp(Operation* op) {
              pim::PimTransposeOp>(op);
 }
 
+static FailureOr<Value> lowerFragmentAssemblyBlueprint(IRRewriter& rewriter,
+                                                       spatial::SpatBlueprintOp blueprint,
+                                                       IRMapping& mapping) {
+  auto resultType = dyn_cast<ShapedType>(blueprint.getOutput().getType());
+  if (!resultType || !resultType.hasStaticShape())
+    return blueprint.emitOpError("fragment assembly lowering requires a static ranked tensor result");
+
+  std::optional<StringRef> modeAttr = blueprint.getMode();
+  std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
+  std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
+  std::optional<ArrayRef<int64_t>> fragmentStridesAttr = blueprint.getFragmentStrides();
+  if (!modeAttr || *modeAttr != "fragment_assembly" || !operandIndicesAttr || !sourceOffsetsAttr
+      || !fragmentStridesAttr)
+    return blueprint.emitOpError("fragment assembly lowering requires explicit fragment metadata");
+
+  ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
+  ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
+  ArrayRef<int64_t> flatOffsets = blueprint.getFragmentOffsets();
+  ArrayRef<int64_t> flatSizes = blueprint.getFragmentSizes();
+  ArrayRef<int64_t> flatStrides = *fragmentStridesAttr;
+  int64_t rank = resultType.getRank();
+
+  SmallVector<Value> fragmentOperands {blueprint.getInput()};
+  llvm::append_range(fragmentOperands, blueprint.getFragments());
+  if (failed(validateFragmentAssemblyMetadata(blueprint,
+                                              rank,
+                                              fragmentOperands.size(),
+                                              operandIndices,
+                                              sourceOffsets,
+                                              flatOffsets,
+                                              flatSizes,
+                                              flatStrides)))
+    return failure();
+
+  SmallVector<int64_t> hostStrides = computeRowMajorStrides(resultType.getShape());
+  SmallVector<FragmentAssemblyCopy, 8> copies;
+  for (int64_t fragmentIndex = 0; fragmentIndex < static_cast<int64_t>(operandIndices.size()); ++fragmentIndex) {
+    int64_t operandIndex = operandIndices[fragmentIndex];
+
+    SmallVector<int64_t, 4> fragmentOffsets;
+    SmallVector<int64_t, 4> fragmentSizes;
+    for (int64_t dim = 0; dim < rank; ++dim) {
+      int64_t flatIndex = fragmentIndex * rank + dim;
+      if (flatStrides[flatIndex] != 1)
+        return blueprint.emitOpError("fragment assembly lowering only supports unit strides");
+      fragmentOffsets.push_back(flatOffsets[flatIndex]);
+      fragmentSizes.push_back(flatSizes[flatIndex]);
+    }
+
+    Value source = mapping.lookupOrDefault(fragmentOperands[operandIndex]);
+    auto sourceType = dyn_cast<RankedTensorType>(source.getType());
+    if (!sourceType || !sourceType.hasStaticShape())
+      return blueprint.emitOpError("fragment assembly lowering requires static ranked tensor operands");
+    size_t elementSize = getElementTypeSizeInBytes(sourceType.getElementType());
+    if (failed(forEachContiguousDestinationChunk(
+          resultType.getShape(),
+          fragmentOffsets,
+          fragmentSizes,
+          [&](ArrayRef<int64_t> chunkOffsets, int64_t relativeSourceOffset, int64_t chunkElements) -> LogicalResult {
+            int64_t hostElementOffset = 0;
+            for (auto [dim, offset] : llvm::enumerate(chunkOffsets))
+              hostElementOffset += offset * hostStrides[dim];
+
+            FragmentAssemblyCopy copy;
+            copy.source = source;
+            copy.sourceType = sourceType;
+            copy.sourceByteOffset =
+              (sourceOffsets[fragmentIndex] + relativeSourceOffset) * static_cast<int64_t>(elementSize);
+            copy.hostByteOffset = hostElementOffset * static_cast<int64_t>(elementSize);
+            copy.byteSize = chunkElements * static_cast<int64_t>(elementSize);
+            copies.push_back(copy);
+            return success();
+          })))
+      return failure();
+  }
+
+  Value currentOutput = createEmptyTensorFromShaped(rewriter, blueprint.getLoc(), resultType);
+  FailureOr<SmallVector<FragmentAssemblyCopyRun, 8>> runs = groupFragmentAssemblyCopyRuns(copies);
+  if (failed(runs))
+    return failure();
+  return emitFragmentAssemblyCopyRuns(
+    rewriter, blueprint.getLoc(), *runs, currentOutput, blueprint.getOperation());
+}
+
 static void
 cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewriter& rewriter, OperationFolder& constantFolder) {
   for (Value operand : op->getOperands()) {
@@ -55,18 +142,7 @@ cloneMappedHelperOperands(Operation* op, IRMapping& mapping, IRRewriter& rewrite
   }
 }
 
-static FailureOr<int32_t> getPimCoreIdForComputeOp(spatial::SpatCompute computeOp, size_t& fallbackCoreId) {
-  if (auto spatialCoreIdAttr = computeOp->getAttrOfType<IntegerAttr>(onnx_mlir::kCoreIdAttrName))
-    return pim::checkedI32(spatialCoreIdAttr.getInt(), computeOp, "spatial compute core id");
-  auto checkedCoreId =
-    pim::checkedI32(static_cast<uint64_t>(fallbackCoreId), computeOp, "fallback spatial compute core id");
-  if (failed(checkedCoreId))
-    return failure();
-  ++fallbackCoreId;
-  return *checkedCoreId;
-}
-
-static LogicalResult collectHelperComputeChain(spatial::SpatCompute computeOp,
+static LogicalResult collectHelperComputeChain(spatial::SpatScheduledCompute computeOp,
                                                SmallVectorImpl<Operation*>& helperChain,
                                                bool requireReturnUse = true) {
   if (computeOp.getInputs().size() != 1 || computeOp.getNumResults() != 1)
@@ -104,13 +180,13 @@ static LogicalResult collectHelperComputeChain(spatial::SpatCompute computeOp,
   return success();
 }
 
-static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatCompute computeOp,
+static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatScheduledCompute computeOp,
                                                            IRRewriter& rewriter,
                                                            OperationFolder& constantFolder) {
   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 isa<spatial::SpatScheduledCompute, spatial::SpatScheduledComputeBatch, pim::PimCoreOp, pim::PimCoreBatchOp>(user);
       }))
     return false;
 
@@ -131,6 +207,17 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatCompute
     mapping.map(*weightArg, weight);
   }
   for (Operation& op : block.without_terminator()) {
+    if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
+      std::optional<StringRef> modeAttr = blueprint.getMode();
+      if (modeAttr && *modeAttr == "fragment_assembly") {
+        auto lowered = lowerFragmentAssemblyBlueprint(rewriter, blueprint, mapping);
+        if (failed(lowered))
+          return false;
+        mapping.map(blueprint.getOutput(), *lowered);
+        continue;
+      }
+    }
+
     cloneMappedHelperOperands(&op, mapping, rewriter, constantFolder);
     Operation* clonedOp = rewriter.clone(op, mapping);
     for (auto [originalResult, newResult] : llvm::zip(op.getResults(), clonedOp->getResults()))
@@ -145,7 +232,7 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatCompute
 
 } // namespace
 
-LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatCompute computeOp,
+LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatScheduledCompute computeOp,
                                                        IRRewriter& rewriter,
                                                        OperationFolder& constantFolder) {
   Location loc = computeOp->getLoc();
@@ -214,7 +301,7 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeOp(spatial::SpatCompute comp
   if (!computeOp.getWeights().empty())
     computeWeights.append(computeOp.getWeights().begin(), computeOp.getWeights().end());
   rewriter.setInsertionPointAfter(computeOp);
-  auto checkedCoreId = getPimCoreIdForComputeOp(computeOp, coreId);
+  auto checkedCoreId = getRequiredScheduledCoreId(computeOp, "spatial compute core id");
   if (failed(checkedCoreId))
     return failure();
   auto coreIdAttr = pim::getCheckedI32Attr(rewriter, computeOp, static_cast<int64_t>(*checkedCoreId), "pim core id");

src/PIM/Conversion/SpatialToPim/Patterns.cpp

@@ -1,6 +1,10 @@
+#include "mlir/Transforms/DialectConversion.h"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
 
@@ -11,6 +15,92 @@ namespace raptor {
 
 } // namespace raptor
 
+struct LowerFragmentAssemblyBlueprintPattern
+  : OpConversionPattern<spatial::SpatBlueprintOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(spatial::SpatBlueprintOp op,
+                                OpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    std::optional<StringRef> modeAttr = op.getMode();
+    if (!modeAttr || *modeAttr != "fragment_assembly")
+      return failure();
+
+    auto resultType = dyn_cast<ShapedType>(op.getOutput().getType());
+    if (!resultType || !resultType.hasStaticShape())
+      return op.emitOpError("fragment assembly lowering requires a static ranked tensor result");
+
+    std::optional<ArrayRef<int64_t>> operandIndicesAttr = op.getFragmentOperandIndices();
+    std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = op.getFragmentSourceOffsets();
+    std::optional<ArrayRef<int64_t>> fragmentStridesAttr = op.getFragmentStrides();
+    if (!operandIndicesAttr || !sourceOffsetsAttr || !fragmentStridesAttr)
+      return op.emitOpError("fragment assembly lowering requires explicit fragment metadata");
+
+    ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
+    ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
+    ArrayRef<int64_t> flatOffsets = op.getFragmentOffsets();
+    ArrayRef<int64_t> flatSizes = op.getFragmentSizes();
+    ArrayRef<int64_t> flatStrides = *fragmentStridesAttr;
+    int64_t rank = resultType.getRank();
+
+    SmallVector<Value> fragmentOperands {adaptor.getInput()};
+    llvm::append_range(fragmentOperands, adaptor.getFragments());
+    if (failed(validateFragmentAssemblyMetadata(
+          op, rank, fragmentOperands.size(), operandIndices, sourceOffsets, flatOffsets, flatSizes, flatStrides)))
+      return failure();
+
+    Value currentOutput =
+      tensor::EmptyOp::create(rewriter, op.getLoc(), resultType.getShape(), resultType.getElementType()).getResult();
+    for (int64_t fragmentIndex = 0; fragmentIndex < static_cast<int64_t>(operandIndices.size()); ++fragmentIndex) {
+      int64_t operandIndex = operandIndices[fragmentIndex];
+
+      SmallVector<int64_t, 4> fragmentOffsets;
+      for (int64_t dim = 0; dim < rank; ++dim) {
+        int64_t flatIndex = fragmentIndex * rank + dim;
+        if (flatStrides[flatIndex] != 1)
+          return op.emitOpError("fragment assembly lowering only supports unit strides");
+        fragmentOffsets.push_back(flatOffsets[flatIndex]);
+      }
+
+      Value source = fragmentOperands[operandIndex];
+      auto sourceType = dyn_cast<RankedTensorType>(source.getType());
+      if (!sourceType || !sourceType.hasStaticShape())
+        return op.emitOpError("fragment assembly lowering requires static ranked tensor operands");
+
+      SmallVector<int64_t, 4> fragmentShape;
+      fragmentShape.reserve(rank);
+      for (int64_t dim = 0; dim < rank; ++dim)
+        fragmentShape.push_back(flatSizes[fragmentIndex * rank + dim]);
+
+      Value fragment = source;
+      if (llvm::to_vector(sourceType.getShape()) != fragmentShape || sourceOffsets[fragmentIndex] != 0) {
+        FailureOr<SmallVector<int64_t, 4>> extractOffsets = getStaticSliceOffsetsForElementOffset(
+          op, sourceType, fragmentShape, sourceOffsets[fragmentIndex], "fragment assembly source slice");
+        if (failed(extractOffsets))
+          return failure();
+        fragment = tensor::ExtractSliceOp::create(rewriter,
+                                                  op.getLoc(),
+                                                  source,
+                                                  getStaticIndexAttrs(rewriter, *extractOffsets),
+                                                  getStaticIndexAttrs(rewriter, fragmentShape),
+                                                  getUnitStrides(rewriter, rank));
+      }
+
+      currentOutput = tensor::InsertSliceOp::create(rewriter,
+                                                    op.getLoc(),
+                                                    fragment,
+                                                    currentOutput,
+                                                    getStaticIndexAttrs(rewriter, fragmentOffsets),
+                                                    getStaticIndexAttrs(rewriter, fragmentShape),
+                                                    getUnitStrides(rewriter, rank))
+                        .getResult();
+    }
+
+    rewriter.replaceOp(op, currentOutput);
+    return success();
+  }
+};
+
 void populateInitialPatterns(RewritePatternSet& patterns) {
   raptor::populateWithGenerated(patterns);
   populateTransposeLoweringPatterns(patterns);
@@ -19,6 +109,7 @@ void populateInitialPatterns(RewritePatternSet& patterns) {
 void populateCoreBodyPatterns(RewritePatternSet& patterns) {
   raptor::populateWithGenerated(patterns);
   populateTransposeLoweringPatterns(patterns);
+  patterns.add<LowerFragmentAssemblyBlueprintPattern>(patterns.getContext());
 }
 
 } // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/Patterns/ChannelLowering.cpp

@@ -10,6 +10,14 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace {
 
+static void copyRaptorDebugAttrs(Operation* source, Operation* target) {
+  for (NamedAttribute attr : source->getAttrs()) {
+    StringRef name = attr.getName().strref();
+    if (name.starts_with("raptor."))
+      target->setAttr(attr.getName(), attr.getValue());
+  }
+}
+
 struct ChannelSendLowering : OpRewritePattern<spatial::SpatChannelSendOp> {
   using OpRewritePattern::OpRewritePattern;
 
@@ -17,7 +25,8 @@ struct ChannelSendLowering : OpRewritePattern<spatial::SpatChannelSendOp> {
     auto sizeAttr = getTensorSizeInBytesAttr(rewriter, op.getOperation(), op.getInput());
     if (failed(sizeAttr))
       return failure();
-    pim::PimSendOp::create(rewriter, op.getLoc(), op.getInput(), *sizeAttr, op.getTargetCoreId());
+    auto send = pim::PimSendOp::create(rewriter, op.getLoc(), op.getInput(), *sizeAttr, op.getTargetCoreId());
+    copyRaptorDebugAttrs(op.getOperation(), send.getOperation());
     rewriter.eraseOp(op);
     return success();
   }
@@ -37,9 +46,10 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
     auto sizeAttr = getTensorSizeInBytesAttr(rewriter, op.getOperation(), op.getResult());
     if (failed(sizeAttr))
       return failure();
-    Value received = pim::PimReceiveOp::create(
-                       rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, *sizeAttr, op.getSourceCoreId())
-                       .getOutput();
+    auto receive = pim::PimReceiveOp::create(
+      rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, *sizeAttr, op.getSourceCoreId());
+    copyRaptorDebugAttrs(op.getOperation(), receive.getOperation());
+    Value received = receive.getOutput();
     rewriter.replaceOp(op, received);
     return success();
   }

src/PIM/Conversion/SpatialToPim/Patterns/GlobalTensorMaterialization.cpp

@@ -59,7 +59,7 @@ struct MoveExtractSliceIntoCompute final : OpRewritePattern<mlir::tensor::Extrac
       return failure();
 
     for (auto& uses : extractSliceOp->getUses()) {
-      if (isa<spatial::SpatCompute>(uses.getOwner())) {
+      if (isa<spatial::SpatScheduledCompute>(uses.getOwner())) {
         if (!getDirectComputeLikeInputIndex(uses.getOwner(), uses.getOperandNumber()))
           return failure();
       }
@@ -72,7 +72,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())) {
+      if (auto spatCompute = dyn_cast<spatial::SpatScheduledCompute>(uses.getOwner())) {
         auto inputIndex = getDirectComputeLikeInputIndex(spatCompute, uses.getOperandNumber());
         if (!inputIndex)
           return failure();
@@ -92,7 +92,7 @@ struct MoveExtractSliceIntoCompute final : OpRewritePattern<mlir::tensor::Extrac
         replaceAndEraseDirectComputeLikeInput(
           rewriter, spatCompute.getOperation(), *inputIndex, mapSpatToExtract[spatCompute.getOperation()]);
       }
-      else if (auto spatComputeBatch = dyn_cast<spatial::SpatComputeBatch>(uses.getOwner())) {
+      else if (auto spatComputeBatch = dyn_cast<spatial::SpatScheduledComputeBatch>(uses.getOwner())) {
         auto inputIndex = getDirectComputeLikeInputIndex(spatComputeBatch, uses.getOperandNumber());
         if (!inputIndex)
           return failure();
@@ -114,7 +114,7 @@ struct MoveExtractSliceIntoCompute final : OpRewritePattern<mlir::tensor::Extrac
       }
       else {
         {
-          if (auto spatCompute = uses.getOwner()->getParentOfType<spatial::SpatCompute>()) {
+          if (auto spatCompute = uses.getOwner()->getParentOfType<spatial::SpatScheduledCompute>()) {
             rewriter.setInsertionPoint(&spatCompute.getBody().front().front());
             if (!mapSpatToExtract.contains(spatCompute.getOperation())) {
               auto newExtractSlice = rewriter.clone(*extractSliceOp.getOperation());
@@ -125,7 +125,7 @@ struct MoveExtractSliceIntoCompute final : OpRewritePattern<mlir::tensor::Extrac
             uses.set(mapSpatToExtract[spatCompute.getOperation()]);
             rewriter.finalizeOpModification(spatCompute.getOperation());
           }
-          else if (auto spatComputeBatch = uses.getOwner()->getParentOfType<spatial::SpatComputeBatch>()) {
+          else if (auto spatComputeBatch = uses.getOwner()->getParentOfType<spatial::SpatScheduledComputeBatch>()) {
             rewriter.setInsertionPoint(&spatComputeBatch.getBody().front().front());
             if (!mapSpatToExtract.contains(spatComputeBatch.getOperation())) {
               auto newExtractSlice = rewriter.clone(*extractSliceOp.getOperation());
@@ -179,7 +179,7 @@ struct FuncOpArgToGlobalMemoryPattern final : OpRewritePattern<mlir::func::FuncO
 
       for (auto& argUses : llvm::make_early_inc_range(arg.getUses())) {
         auto argUser = argUses.getOwner();
-        if (auto spatCompute = dyn_cast<spatial::SpatCompute>(argUser)) {
+        if (auto spatCompute = dyn_cast<spatial::SpatScheduledCompute>(argUser)) {
           auto inputIndex = getDirectComputeLikeInputIndex(spatCompute, argUses.getOperandNumber());
           if (!inputIndex)
             return failure();
@@ -191,7 +191,7 @@ struct FuncOpArgToGlobalMemoryPattern final : OpRewritePattern<mlir::func::FuncO
 
           replaceAndEraseDirectComputeLikeInput(rewriter, spatCompute.getOperation(), BBArgIndex, toTensor);
         }
-        else if (auto spatComputeBatch = dyn_cast<spatial::SpatComputeBatch>(argUser)) {
+        else if (auto spatComputeBatch = dyn_cast<spatial::SpatScheduledComputeBatch>(argUser)) {
           auto inputIndex = getDirectComputeLikeInputIndex(spatComputeBatch, argUses.getOperandNumber());
           if (!inputIndex)
             return failure();

src/PIM/Conversion/SpatialToPim/ReturnPathNormalization.cpp

@@ -2,6 +2,7 @@
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/IR/BuiltinOps.h"
@@ -11,6 +12,7 @@
 
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Conversion/SpatialToPim/SpatialToPimPass.hpp"
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
@@ -86,7 +88,7 @@ getCheckedByteOffset(int64_t elementOffset, size_t elementSize, Operation* ancho
   return pim::checkedCast<int64_t>(*byteOffset, anchor, fieldName);
 }
 
-static LogicalResult collectHelperComputeChain(spatial::SpatCompute computeOp,
+static LogicalResult collectHelperComputeChain(spatial::SpatScheduledCompute computeOp,
                                                SmallVectorImpl<Operation*>& helperChain) {
   if (computeOp.getInputs().size() != 1 || computeOp.getNumResults() != 1)
     return failure();
@@ -149,6 +151,40 @@ static std::optional<ReturnUseInfo> analyzeReturnUse(Value value) {
   };
 }
 
+static FailureOr<SmallVector<std::pair<spatial::SpatBlueprintOp, size_t>, 4>>
+analyzeTopLevelFragmentAssemblyUses(Value value) {
+  SmallVector<std::pair<spatial::SpatBlueprintOp, size_t>, 4> uses;
+  for (OpOperand& use : value.getUses()) {
+    auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(use.getOwner());
+    if (!blueprint || blueprint->getParentOp() != blueprint->getParentOfType<func::FuncOp>())
+      return failure();
+    std::optional<StringRef> mode = blueprint.getMode();
+    if (!mode || *mode != "fragment_assembly")
+      return failure();
+    if (!blueprint.getOutput().hasOneUse() || !isa<func::ReturnOp>(*blueprint.getOutput().getUsers().begin()))
+      return failure();
+    std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
+    std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
+    std::optional<ArrayRef<int64_t>> stridesAttr = blueprint.getFragmentStrides();
+    auto resultType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
+    if (!operandIndicesAttr || !sourceOffsetsAttr || !stridesAttr || !resultType || !resultType.hasStaticShape())
+      return failure();
+    SmallVector<Value> fragmentOperands {blueprint.getInput()};
+    llvm::append_range(fragmentOperands, blueprint.getFragments());
+    if (failed(validateFragmentAssemblyMetadata(blueprint,
+                                                resultType.getRank(),
+                                                fragmentOperands.size(),
+                                                *operandIndicesAttr,
+                                                *sourceOffsetsAttr,
+                                                blueprint.getFragmentOffsets(),
+                                                blueprint.getFragmentSizes(),
+                                                *stridesAttr)))
+      return failure();
+    uses.emplace_back(blueprint, use.getOperandNumber());
+  }
+  return uses;
+}
+
 static std::optional<ConcatReturnUseInfo> analyzeConcatReturnUse(Value value) {
   auto getConcatResult = [](Operation* op) -> Value {
     if (auto tensorConcat = dyn_cast<tensor::ConcatOp>(op))
@@ -212,7 +248,7 @@ static std::optional<ConcatReturnUseInfo> analyzeConcatReturnUse(Value value) {
   }
 
   SmallVector<Operation*> helperChain;
-  if (auto helperCompute = dyn_cast<spatial::SpatCompute>(currentUser)) {
+  if (auto helperCompute = dyn_cast<spatial::SpatScheduledCompute>(currentUser)) {
     if (helperCompute.getInputs().size() != 1 || helperCompute.getInputs().front() != currentValue)
       return std::nullopt;
 
@@ -559,6 +595,115 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
     }
   }
 
+  FailureOr<SmallVector<std::pair<spatial::SpatBlueprintOp, size_t>, 4>> fragmentAssemblyUses =
+    analyzeTopLevelFragmentAssemblyUses(producedValue);
+  if (succeeded(fragmentAssemblyUses)) {
+    auto sourceType = dyn_cast<RankedTensorType>(storedValue.getType());
+    if (!sourceType || !sourceType.hasStaticShape()) {
+      producerOp->emitOpError("fragment assembly publication requires a static ranked tensor source");
+      return ReturnPathLoweringResult::Failure;
+    }
+
+    size_t elementSize = getElementTypeSizeInBytes(sourceType.getElementType());
+    for (auto [blueprint, operandNumber] : *fragmentAssemblyUses) {
+      rewriter.setInsertionPointAfterValue(storedValue);
+      std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
+      std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
+      std::optional<ArrayRef<int64_t>> stridesAttr = blueprint.getFragmentStrides();
+      if (!operandIndicesAttr || !sourceOffsetsAttr || !stridesAttr) {
+        blueprint.emitOpError(
+          "fragment assembly lowering requires explicit operand, source-offset, and stride metadata");
+        return ReturnPathLoweringResult::Failure;
+      }
+
+      size_t returnIndex = blueprint.getOutput().getUses().begin()->getOperandNumber();
+      Value outputTensor = outputTensors[returnIndex](rewriter, loc);
+      auto outputType = dyn_cast<RankedTensorType>(outputTensor.getType());
+      auto resultType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
+      if (!outputType || !resultType || !resultType.hasStaticShape()) {
+        blueprint.emitOpError("fragment assembly lowering requires static ranked host outputs");
+        return ReturnPathLoweringResult::Failure;
+      }
+
+      ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
+      ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
+      ArrayRef<int64_t> flatOffsets = blueprint.getFragmentOffsets();
+      ArrayRef<int64_t> flatSizes = blueprint.getFragmentSizes();
+      ArrayRef<int64_t> flatStrides = *stridesAttr;
+      int64_t rank = resultType.getRank();
+      if (failed(validateFragmentAssemblyMetadata(blueprint,
+                                                  rank,
+                                                  1 + blueprint.getFragments().size(),
+                                                  operandIndices,
+                                                  sourceOffsets,
+                                                  flatOffsets,
+                                                  flatSizes,
+                                                  flatStrides)))
+        return ReturnPathLoweringResult::Failure;
+      SmallVector<FragmentAssemblyCopy, 8> copies;
+      for (int64_t fragmentIndex = 0; fragmentIndex < static_cast<int64_t>(operandIndices.size()); ++fragmentIndex) {
+        if (operandIndices[fragmentIndex] != static_cast<int64_t>(operandNumber))
+          continue;
+
+        SmallVector<int64_t, 4> fragmentOffsets;
+        SmallVector<int64_t, 4> fragmentSizes;
+        for (int64_t dim = 0; dim < rank; ++dim) {
+          int64_t flatIndex = fragmentIndex * rank + dim;
+          if (flatStrides[flatIndex] != 1) {
+            blueprint.emitOpError("fragment assembly lowering only supports unit strides");
+            return ReturnPathLoweringResult::Failure;
+          }
+          fragmentOffsets.push_back(flatOffsets[flatIndex]);
+          fragmentSizes.push_back(flatSizes[flatIndex]);
+        }
+
+        bool failedChunk = false;
+        if (failed(forEachContiguousDestinationChunk(
+          outputType.getShape(),
+          fragmentOffsets,
+          fragmentSizes,
+          [&](ArrayRef<int64_t> chunkOffsets, int64_t relativeSourceOffset, int64_t chunkElements) -> LogicalResult {
+            auto hostOffset =
+              getCheckedByteOffset(computeFlatElementIndex(chunkOffsets, outputType.getShape()),
+                                   elementSize,
+                                   producerOp,
+                                   "fragment assembly host offset");
+            auto sourceOffset = getCheckedByteOffset(sourceOffsets[fragmentIndex] + relativeSourceOffset,
+                                                     elementSize,
+                                                     producerOp,
+                                                     "fragment assembly source offset");
+            auto fragmentBytes =
+              getCheckedByteOffset(chunkElements, elementSize, producerOp, "fragment assembly host copy byte size");
+            if (failed(hostOffset) || failed(sourceOffset) || failed(fragmentBytes)) {
+              failedChunk = true;
+              return failure();
+            }
+            FragmentAssemblyCopy copy;
+            copy.source = storedValue;
+            copy.sourceType = sourceType;
+            copy.hostByteOffset = *hostOffset;
+            copy.sourceByteOffset = *sourceOffset;
+            copy.byteSize = *fragmentBytes;
+            copies.push_back(copy);
+            return success();
+          })))
+          failedChunk = true;
+        if (failedChunk)
+          return ReturnPathLoweringResult::Failure;
+      }
+      FailureOr<SmallVector<FragmentAssemblyCopyRun, 8>> runs = groupFragmentAssemblyCopyRuns(copies);
+      if (failed(runs))
+        return ReturnPathLoweringResult::Failure;
+      FailureOr<Value> updatedOutput =
+        emitFragmentAssemblyCopyRuns(rewriter, blueprint.getLoc(), *runs, outputTensor, producerOp);
+      if (failed(updatedOutput))
+        return ReturnPathLoweringResult::Failure;
+      outputTensor = *updatedOutput;
+      markOpToRemove(blueprint.getOperation());
+    }
+    return ReturnPathLoweringResult::Handled;
+  }
+
   if (auto concatReturnUse = analyzeConcatReturnUse(producedValue)) {
     size_t elementSize = getElementTypeSizeInBytes(storedTensorType.getElementType());
     auto storedByteSize =
@@ -643,7 +788,7 @@ raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::low
 }
 
 raptor::SpatialToPimPass::ReturnPathLoweringResult raptor::SpatialToPimPass::lowerComputeResultReturnPath(
-  spatial::SpatCompute computeOp, OpResult result, Value yieldValue, IRRewriter& rewriter) {
+  spatial::SpatScheduledCompute computeOp, OpResult result, Value yieldValue, IRRewriter& rewriter) {
   return lowerProducedValueReturnPath(computeOp.getOperation(), result, yieldValue, rewriter);
 }
 
@@ -656,7 +801,7 @@ void raptor::SpatialToPimPass::replaceReturnWithOutputBuffers(func::ReturnOp ret
     if (!isExclusivelyOwnedByReturnChain && op->hasOneUse()) {
       Operation* onlyUser = *op->getUsers().begin();
       isExclusivelyOwnedByReturnChain =
-        isa<func::ReturnOp, tensor::ConcatOp, spatial::SpatConcatOp, pim::PimConcatOp, spatial::SpatCompute>(onlyUser)
+        isa<func::ReturnOp, tensor::ConcatOp, spatial::SpatConcatOp, pim::PimConcatOp, spatial::SpatScheduledCompute>(onlyUser)
         || isReturnHelperChainOp(onlyUser);
     }
     if (!isExclusivelyOwnedByReturnChain)
@@ -669,7 +814,17 @@ void raptor::SpatialToPimPass::replaceReturnWithOutputBuffers(func::ReturnOp ret
       return;
     }
 
-    if (auto computeOp = dyn_cast<spatial::SpatCompute>(op)) {
+    if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
+      std::optional<StringRef> mode = blueprint.getMode();
+      if (mode && *mode == "fragment_assembly") {
+        markOpToRemove(blueprint.getOperation());
+        for (Value operand : blueprint->getOperands())
+          markOwnedReturnChain(operand.getDefiningOp(), markOwnedReturnChain);
+        return;
+      }
+    }
+
+    if (auto computeOp = dyn_cast<spatial::SpatScheduledCompute>(op)) {
       markOpToRemove(computeOp);
       if (!computeOp.getInputs().empty())
         for (Value input : computeOp.getInputs())

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

@@ -25,9 +25,11 @@
 #include <cassert>
 #include <utility>
 
+#include "Common/IR/ShapeUtils.hpp"
 #include "Common/IR/ConstantUtils.hpp"
 #include "Common/PimCommon.hpp"
 #include "Common/Support/CheckedArithmetic.hpp"
+#include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
 #include "Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "Conversion/SpatialToPim/Common.hpp"
 #include "Conversion/SpatialToPim/Patterns.hpp"
@@ -42,63 +44,29 @@ using namespace pim;
 
 namespace onnx_mlir {
 
-static memref::GlobalOp getOrCreateZeroGlobal(IRRewriter& rewriter, Location loc, RankedTensorType tensorType) {
-  auto moduleOp = rewriter.getBlock()->getParentOp()->getParentOfType<ModuleOp>();
-  auto memRefType = MemRefType::get(tensorType.getShape(), tensorType.getElementType());
-  auto zeroAttr = DenseElementsAttr::get(tensorType, rewriter.getZeroAttr(tensorType.getElementType()));
-
-  for (auto globalOp : moduleOp.getOps<memref::GlobalOp>()) {
-    if (!globalOp.getConstant() || globalOp.getType() != memRefType || !globalOp.getInitialValue())
-      continue;
-    if (dyn_cast<DenseElementsAttr>(*globalOp.getInitialValue()) == zeroAttr)
-      return globalOp;
-  }
-
-  std::string nameStem;
-  llvm::raw_string_ostream nameStream(nameStem);
-  nameStream << "__pim_zero_" << tensorType.getRank() << "d_" << tensorType.getNumElements();
-  nameStream.flush();
-
-  std::string symbolName = nameStem;
-  unsigned suffix = 0;
-  while (SymbolTable::lookupSymbolIn(moduleOp, symbolName))
-    symbolName = (nameStem + "_" + Twine(suffix++)).str();
-
-  OpBuilder::InsertionGuard guard(rewriter);
-  rewriter.setInsertionPointToStart(moduleOp.getBody());
-  return memref::GlobalOp::create(rewriter,
-                                  loc,
-                                  rewriter.getStringAttr(symbolName),
-                                  rewriter.getStringAttr("private"),
-                                  TypeAttr::get(memRefType),
-                                  zeroAttr,
-                                  rewriter.getUnitAttr(),
-                                  IntegerAttr {});
-}
-
-static FailureOr<Value> createZeroedDeviceHVector(IRRewriter& rewriter,
-                                                  Location loc,
-                                                  RankedTensorType tensorType,
-                                                  OperationFolder& constantFolder) {
-  auto outputBuffer = createEmptyTensorFromShaped(rewriter, loc, tensorType);
-  auto zeroGlobal = getOrCreateZeroGlobal(rewriter, loc, tensorType);
-  auto zeroValue = memref::GetGlobalOp::create(rewriter, loc, zeroGlobal.getType(), zeroGlobal.getName());
-  auto zeroIndex = getOrCreateIndexConstant(constantFolder, outputBuffer.getOperation(), 0);
-  auto byteSize =
-    pim::getCheckedShapedTypeSizeInBytes(tensorType, outputBuffer.getOperation(), "host-to-device zero copy byte size");
-  if (failed(byteSize))
-    return failure();
-  auto sizeAttr =
-    pim::getCheckedI32Attr(rewriter, outputBuffer.getOperation(), *byteSize, "host-to-device zero copy byte size");
-  if (failed(sizeAttr))
-    return failure();
-  return PimMemCopyHostToDevOp::create(
-           rewriter, loc, tensorType, zeroIndex, zeroIndex, outputBuffer, zeroValue, *sizeAttr)
-    .getOutput();
-}
-
 static FailureOr<Value>
-padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector, OperationFolder& constantFolder) {
+createZeroPaddedTensor(IRRewriter& rewriter, Location loc, Value value, RankedTensorType resultType) {
+  auto sourceType = cast<RankedTensorType>(value.getType());
+  SmallVector<OpFoldResult> lowPads(sourceType.getRank(), rewriter.getIndexAttr(0));
+  SmallVector<OpFoldResult> highPads;
+  highPads.reserve(sourceType.getRank());
+  for (auto [sourceDim, resultDim] : llvm::zip(sourceType.getShape(), resultType.getShape()))
+    highPads.push_back(rewriter.getIndexAttr(resultDim - sourceDim));
+
+  auto padOp = tensor::PadOp::create(rewriter, loc, resultType, value, lowPads, highPads);
+  auto* padBlock = new Block();
+  for (int64_t i = 0; i < sourceType.getRank(); ++i)
+    padBlock->addArgument(rewriter.getIndexType(), loc);
+  padOp.getRegion().push_back(padBlock);
+  rewriter.setInsertionPointToStart(padBlock);
+  auto zero = getOrCreateConstant(
+    rewriter, padOp.getOperation(), rewriter.getZeroAttr(sourceType.getElementType()), sourceType.getElementType());
+  tensor::YieldOp::create(rewriter, loc, zero);
+  rewriter.setInsertionPointAfter(padOp);
+  return padOp.getResult();
+}
+
+static FailureOr<Value> padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector) {
   auto vectorType = cast<RankedTensorType>(vector.getType());
   ArrayRef<int64_t> shape = vectorType.getShape();
   assert(isHVectorShape(shape) && "expected a horizontal vector");
@@ -109,22 +77,10 @@ padHVectorInputToCrossbarSize(IRRewriter& rewriter, Location loc, Value vector,
 
   auto paddedType = RankedTensorType::get(
     {shape[0], static_cast<int64_t>(crossbarSize)}, vectorType.getElementType(), vectorType.getEncoding());
-  auto zeroed = createZeroedDeviceHVector(rewriter, loc, paddedType, constantFolder);
-  if (failed(zeroed))
-    return failure();
-  Value zeroIndex = getOrCreateIndexConstant(constantFolder, zeroed->getDefiningOp(), 0);
-  auto byteSize =
-    pim::getCheckedShapedTypeSizeInBytes(vectorType, zeroed->getDefiningOp(), "device padding copy byte size");
-  if (failed(byteSize))
-    return failure();
-  auto sizeAttr = pim::getCheckedI32Attr(rewriter, zeroed->getDefiningOp(), *byteSize, "device padding copy byte size");
-  if (failed(sizeAttr))
-    return failure();
-  return PimMemCopyOp::create(rewriter, loc, paddedType, zeroIndex, zeroIndex, *zeroed, vector, *sizeAttr).getOutput();
+  return createZeroPaddedTensor(rewriter, loc, vector, paddedType);
 }
 
 void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
-  coreId = 0;
   outputTensors.clear();
   operationsToRemove.clear();
   ModuleOp moduleOp = getOperation();
@@ -137,6 +93,12 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
     return;
   }
   func::FuncOp funcOp = *entryFunc;
+  if (failed(verifyScheduledSpatialInvariants(funcOp))) {
+    funcOp.emitOpError(
+      "scheduled Spatial verification failed at the start of SpatialToPim");
+    signalPassFailure();
+    return;
+  }
 
   IRRewriter rewriter(&getContext());
   OperationFolder constantFolder(&getContext());
@@ -176,19 +138,19 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
     return;
   }
 
-  for (auto computeOp : funcOp.getOps<spatial::SpatCompute>()) {
+  for (auto computeOp : funcOp.getOps<spatial::SpatScheduledCompute>()) {
     markOpToRemove(computeOp);
     if (failed(lowerComputeOp(computeOp, rewriter, constantFolder))) {
-      computeOp.emitOpError("failed to lower spat.compute to pim.core");
+      computeOp.emitOpError("failed to lower spat.scheduled_compute to pim.core");
       signalPassFailure();
       return;
     }
   }
 
-  for (auto computeBatchOp : funcOp.getOps<spatial::SpatComputeBatch>()) {
+  for (auto computeBatchOp : funcOp.getOps<spatial::SpatScheduledComputeBatch>()) {
     markOpToRemove(computeBatchOp);
     if (failed(lowerComputeBatchOp(computeBatchOp, rewriter))) {
-      computeBatchOp.emitOpError("failed to lower spat.compute_batch to pim.core_batch");
+      computeBatchOp.emitOpError("failed to lower spat.scheduled_compute_batch to pim.core_batch");
       signalPassFailure();
       return;
     }
@@ -292,7 +254,6 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
 }
 
 LogicalResult raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func::FuncOp funcOp, IRRewriter& rewriter) {
-  OperationFolder constantFolder(funcOp.getContext());
   bool hasFailure = false;
   funcOp.walk([&](PimVMMOp vmmOp) {
     auto outputType = cast<RankedTensorType>(vmmOp.getOutput().getType());
@@ -301,7 +262,7 @@ LogicalResult raptor::SpatialToPimPass::enlargeVMMOutTensorsToCrossbarSize(func:
     assert(outputShape[1] <= static_cast<int64_t>(crossbarSize) && "output width must fit in one crossbar");
 
     rewriter.setInsertionPoint(vmmOp);
-    auto paddedInput = padHVectorInputToCrossbarSize(rewriter, vmmOp.getLoc(), vmmOp.getInput(), constantFolder);
+    auto paddedInput = padHVectorInputToCrossbarSize(rewriter, vmmOp.getLoc(), vmmOp.getInput());
     if (failed(paddedInput)) {
       hasFailure = true;
       return WalkResult::interrupt();
@@ -374,7 +335,7 @@ LogicalResult raptor::SpatialToPimPass::allocateAndInitializeCoreLocalVariables(
   };
 
   for (auto& op : funcOp.getBody().getOps())
-    if (auto computeOp = dyn_cast<spatial::SpatCompute>(op)) {
+    if (auto computeOp = dyn_cast<spatial::SpatScheduledCompute>(op)) {
       if (!computeOp.getInputs().empty() || computeOp.getBody().front().getNumArguments() != 0)
         continue;
       for (auto getGlobal : computeOp.getOps<memref::GetGlobalOp>()) {

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.hpp

@@ -36,13 +36,15 @@ private:
   using OutputTensorFactory = std::function<mlir::Value(mlir::IRRewriter& rewriter, mlir::Location loc)>;
 
   llvm::SmallVector<OutputTensorFactory> outputTensors;
-  size_t coreId = 0;
   llvm::SmallVector<mlir::Operation*> operationsToRemove;
 
   mlir::LogicalResult allocateAndInitializeCoreLocalVariables(mlir::func::FuncOp funcOp, mlir::IRRewriter& rewriter);
   mlir::LogicalResult
-  lowerComputeOp(spatial::SpatCompute computeOp, mlir::IRRewriter& rewriter, mlir::OperationFolder& constantFolder);
-  mlir::LogicalResult lowerComputeBatchOp(spatial::SpatComputeBatch computeBatchOp, mlir::IRRewriter& rewriter);
+  lowerComputeOp(spatial::SpatScheduledCompute computeOp,
+                 mlir::IRRewriter& rewriter,
+                 mlir::OperationFolder& constantFolder);
+  mlir::LogicalResult lowerComputeBatchOp(spatial::SpatScheduledComputeBatch computeBatchOp,
+                                          mlir::IRRewriter& rewriter);
 
   enum class ReturnPathLoweringResult {
     Handled,
@@ -51,7 +53,7 @@ private:
   };
 
   void addReturnOutputBuffers(mlir::func::ReturnOp returnOp, mlir::IRRewriter& rewriter);
-  ReturnPathLoweringResult lowerComputeResultReturnPath(spatial::SpatCompute computeOp,
+  ReturnPathLoweringResult lowerComputeResultReturnPath(spatial::SpatScheduledCompute computeOp,
                                                         mlir::OpResult result,
                                                         mlir::Value yieldValue,
                                                         mlir::IRRewriter& rewriter);

src/PIM/Dialect/Pim/Transforms/Bufferization/BufferizationUtils.cpp

@@ -13,10 +13,13 @@ using namespace bufferization;
 
 namespace onnx_mlir::pim {
 
-FailureOr<Value> materializeContiguousInputMemRef(Value memrefValue, Location loc, RewriterBase& rewriter) {
+FailureOr<Value> materializeContiguousInputMemRef(Value memrefValue,
+                                                  Location loc,
+                                                  RewriterBase& rewriter,
+                                                  const StaticValueKnowledge& knowledge) {
   bool isContiguous =
-    succeeded(resolveContiguousAddress(memrefValue)) || succeeded(compileContiguousAddressExpr(memrefValue));
-  if (isContiguous && isDeviceLocalPimAddress(memrefValue))
+    succeeded(resolveContiguousAddress(memrefValue, knowledge)) || succeeded(compileContiguousAddressExpr(memrefValue));
+  if (isContiguous && isDeviceLocalPimAddress(memrefValue, knowledge))
     return memrefValue;
 
   auto shapedType = cast<ShapedType>(memrefValue.getType());
@@ -32,7 +35,7 @@ FailureOr<Value> materializeContiguousInputMemRef(Value memrefValue, Location lo
   if (failed(sizeAttr))
     return failure();
 
-  if (isHostBackedPimAddress(memrefValue)) {
+  if (isHostBackedPimAddress(memrefValue, knowledge)) {
     return PimMemCopyHostToDevOp::create(
              rewriter, loc, contiguousType, zeroOffset, zeroOffset, contiguousBuffer, memrefValue, *sizeAttr)
       .getOutput();

src/PIM/Dialect/Pim/Transforms/Bufferization/BufferizationUtils.hpp

@@ -3,10 +3,15 @@
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/IR/PatternMatch.h"
 
+#include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
+
 namespace onnx_mlir::pim {
 
 llvm::FailureOr<mlir::Value>
-materializeContiguousInputMemRef(mlir::Value memrefValue, mlir::Location loc, mlir::RewriterBase& rewriter);
+materializeContiguousInputMemRef(mlir::Value memrefValue,
+                                 mlir::Location loc,
+                                 mlir::RewriterBase& rewriter,
+                                 const onnx_mlir::StaticValueKnowledge& knowledge = {});
 mlir::Value
 allocateContiguousResultMemRefLike(mlir::Value memrefValue, mlir::Location loc, mlir::RewriterBase& rewriter);
 

src/PIM/Dialect/Pim/Transforms/Bufferization/OpBufferizationInterfaces.cpp

@@ -15,6 +15,26 @@ using namespace bufferization;
 namespace onnx_mlir {
 namespace pim {
 
+static StaticValueKnowledge getEnclosingBufferizationKnowledge(Operation* op) {
+  StaticValueKnowledge knowledge;
+
+  if (auto coreBatchOp = op->getParentOfType<PimCoreBatchOp>()) {
+    knowledge.indexValues[coreBatchOp.getLaneArgument()] = 0;
+    for (auto [index, weight] : llvm::enumerate(coreBatchOp.getWeights()))
+      knowledge.aliases[coreBatchOp.getWeightArgument(index)] = weight;
+    for (auto [index, input] : llvm::enumerate(coreBatchOp.getInputs()))
+      knowledge.aliases[coreBatchOp.getInputArgument(index)] = input;
+    return knowledge;
+  }
+
+  if (auto coreOp = op->getParentOfType<PimCoreOp>()) {
+    for (auto [index, weight] : llvm::enumerate(coreOp.getWeights()))
+      knowledge.aliases[coreOp.getWeightArgument(index)] = weight;
+  }
+
+  return knowledge;
+}
+
 struct MemCopyHostToDevOpInterface
 : DstBufferizableOpInterfaceExternalModel<MemCopyHostToDevOpInterface, PimMemCopyHostToDevOp> {
   LogicalResult bufferize(Operation* op,
@@ -148,7 +168,8 @@ struct ConcatOpInterface : DstBufferizableOpInterfaceExternalModel<ConcatOpInter
       auto inputOpt = getBufferOrValue(rewriter, input, options, state);
       if (failed(inputOpt))
         return failure();
-      auto contiguous = materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter);
+      auto contiguous =
+        materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter, getEnclosingBufferizationKnowledge(op));
       if (failed(contiguous))
         return failure();
       inputs.push_back(*contiguous);
@@ -182,7 +203,8 @@ struct SendOpInterface : BufferizableOpInterface::ExternalModel<SendOpInterface,
     auto inputOpt = getBufferOrValue(rewriter, sendOp.getInput(), options, state);
     if (failed(inputOpt))
       return failure();
-    auto contiguousInput = materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter);
+    auto contiguousInput =
+      materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter, getEnclosingBufferizationKnowledge(op));
     if (failed(contiguousInput))
       return failure();
 
@@ -410,7 +432,8 @@ struct TransposeOpInterface : DstBufferizableOpInterfaceExternalModel<TransposeO
     if (failed(outputBufferOpt))
       return failure();
 
-    auto contiguousInput = materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter);
+    auto contiguousInput =
+      materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter, getEnclosingBufferizationKnowledge(op));
     if (failed(contiguousInput))
       return failure();
     Value contiguousOutput = allocateContiguousResultMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
@@ -456,7 +479,8 @@ struct VMMOpInterface : DstBufferizableOpInterfaceExternalModel<VMMOpInterface,
     if (failed(outputBufferOpt))
       return failure();
 
-    auto contiguousInput = materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter);
+    auto contiguousInput =
+      materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter, getEnclosingBufferizationKnowledge(op));
     if (failed(contiguousInput))
       return failure();
     Value contiguousOutput = allocateContiguousResultMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
@@ -497,10 +521,12 @@ struct BinaryDstOpInterface : DstBufferizableOpInterfaceExternalModel<BinaryDstO
     if (failed(outputBufferOpt))
       return failure();
 
-    auto contiguousLhs = materializeContiguousInputMemRef(*lhsOpt, op->getLoc(), rewriter);
+    auto contiguousLhs =
+      materializeContiguousInputMemRef(*lhsOpt, op->getLoc(), rewriter, getEnclosingBufferizationKnowledge(op));
     if (failed(contiguousLhs))
       return failure();
-    auto contiguousRhs = materializeContiguousInputMemRef(*rhsOpt, op->getLoc(), rewriter);
+    auto contiguousRhs =
+      materializeContiguousInputMemRef(*rhsOpt, op->getLoc(), rewriter, getEnclosingBufferizationKnowledge(op));
     if (failed(contiguousRhs))
       return failure();
     Value contiguousOutput = allocateContiguousResultMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
@@ -534,10 +560,12 @@ struct VVDMulOpInterface : DstBufferizableOpInterfaceExternalModel<VVDMulOpInter
     if (failed(outputBufferOpt))
       return failure();
 
-    auto contiguousLhs = materializeContiguousInputMemRef(*lhsOpt, op->getLoc(), rewriter);
+    auto contiguousLhs =
+      materializeContiguousInputMemRef(*lhsOpt, op->getLoc(), rewriter, getEnclosingBufferizationKnowledge(op));
     if (failed(contiguousLhs))
       return failure();
-    auto contiguousRhs = materializeContiguousInputMemRef(*rhsOpt, op->getLoc(), rewriter);
+    auto contiguousRhs =
+      materializeContiguousInputMemRef(*rhsOpt, op->getLoc(), rewriter, getEnclosingBufferizationKnowledge(op));
     if (failed(contiguousRhs))
       return failure();
     Value contiguousOutput = allocateContiguousResultMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);
@@ -574,7 +602,8 @@ struct UnaryDstOpInterface : DstBufferizableOpInterfaceExternalModel<UnaryDstOpI
     if (failed(outputBufferOpt))
       return failure();
 
-    auto contiguousInput = materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter);
+    auto contiguousInput =
+      materializeContiguousInputMemRef(*inputOpt, op->getLoc(), rewriter, getEnclosingBufferizationKnowledge(op));
     if (failed(contiguousInput))
       return failure();
     Value contiguousOutput = allocateContiguousResultMemRefLike(*outputBufferOpt, op->getLoc(), rewriter);

src/PIM/Dialect/Pim/Transforms/Bufferization/PimBufferizationPass.cpp

@@ -116,6 +116,36 @@ lowerMemRefCopyToPimCopy(memref::CopyOp copyOp, PatternRewriter& rewriter, const
   return success();
 }
 
+static LogicalResult verifyLoweredPimCopy(pim::PimMemCopyHostToDevOp copyOp, const StaticValueKnowledge& knowledge) {
+  bool sourceIsHost = isHostBackedPimAddress(copyOp.getHostSource(), knowledge);
+  bool targetIsHost = isHostBackedPimAddress(copyOp.getDeviceTarget(), knowledge);
+  bool sourceIsDevice = isDeviceLocalPimAddress(copyOp.getHostSource(), knowledge);
+  bool targetIsDevice = isDeviceLocalPimAddress(copyOp.getDeviceTarget(), knowledge);
+  if (!sourceIsHost || !targetIsDevice || targetIsHost || sourceIsDevice)
+    return copyOp.emitOpError("pim.memcp_hd requires a host-backed source and a device-local target");
+  return success();
+}
+
+static LogicalResult verifyLoweredPimCopy(pim::PimMemCopyDevToHostOp copyOp, const StaticValueKnowledge& knowledge) {
+  bool sourceIsHost = isHostBackedPimAddress(copyOp.getDeviceSource(), knowledge);
+  bool targetIsHost = isHostBackedPimAddress(copyOp.getHostTarget(), knowledge);
+  bool sourceIsDevice = isDeviceLocalPimAddress(copyOp.getDeviceSource(), knowledge);
+  bool targetIsDevice = isDeviceLocalPimAddress(copyOp.getHostTarget(), knowledge);
+  if (!targetIsHost || !sourceIsDevice || sourceIsHost || targetIsDevice)
+    return copyOp.emitOpError("pim.memcp_dh requires a device-local source and a host-backed target");
+  return success();
+}
+
+static LogicalResult verifyLoweredPimCopy(pim::PimMemCopyOp copyOp, const StaticValueKnowledge& knowledge) {
+  bool sourceIsHost = isHostBackedPimAddress(copyOp.getSource(), knowledge);
+  bool targetIsHost = isHostBackedPimAddress(copyOp.getTarget(), knowledge);
+  bool sourceIsDevice = isDeviceLocalPimAddress(copyOp.getSource(), knowledge);
+  bool targetIsDevice = isDeviceLocalPimAddress(copyOp.getTarget(), knowledge);
+  if (!sourceIsDevice || !targetIsDevice || sourceIsHost || targetIsHost)
+    return copyOp.emitOpError("pim.memcp requires device-local source and target operands");
+  return success();
+}
+
 struct PimBufferizationPass : PassWrapper<PimBufferizationPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(PimBufferizationPass)
   StringRef getArgument() const override { return "bufferize-pim"; }
@@ -129,6 +159,7 @@ struct PimBufferizationPass : PassWrapper<PimBufferizationPass, OperationPass<Mo
 private:
   void annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncOp funcOp) const;
   LogicalResult verifyContiguousRuntimeOperands(ModuleOp moduleOp) const;
+  LogicalResult verifyPimCopyAddressSpaces(ModuleOp moduleOp) const;
 };
 
 static LogicalResult applyPatternsOnce(Operation* op, PatternApplicator& applicator, PatternRewriter& rewriter) {
@@ -240,6 +271,10 @@ void PimBufferizationPass::runOnOperation() {
     signalPassFailure();
     return;
   }
+  if (failed(verifyPimCopyAddressSpaces(moduleOp))) {
+    signalPassFailure();
+    return;
+  }
 
   annotateWeightsMemrefs(moduleOp, funcOp);
 
@@ -267,57 +302,60 @@ void PimBufferizationPass::annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncO
 
 LogicalResult PimBufferizationPass::verifyContiguousRuntimeOperands(ModuleOp moduleOp) const {
   bool hasFailure = false;
-  moduleOp.walk([&](Operation* op) {
+
+  auto verifyWithKnowledge = [&](auto coreLikeOp, const StaticValueKnowledge& initialKnowledge) {
+    (void) walkPimCoreBlockStructurally(
+      coreLikeOp.getBody().front(), initialKnowledge, [&](Operation& op, const StaticValueKnowledge& knowledge) {
         auto verifyOperand = [&](Value operand, unsigned operandIndex) {
           if (!isa<BaseMemRefType>(operand.getType()))
             return;
-      if (succeeded(resolveContiguousAddress(operand)) || succeeded(compileContiguousAddressExpr(operand)))
+          if (succeeded(resolveContiguousAddress(operand, knowledge)) || succeeded(compileContiguousAddressExpr(operand)))
             return;
-      op->emitOpError() << "operand #" << operandIndex
+          op.emitOpError() << "operand #" << operandIndex
                            << " is not backed by contiguous addressable storage after PIM bufferization";
           hasFailure = true;
         };
 
-    if (auto memCopyOp = dyn_cast<PimMemCopyOp>(op)) {
+        if (auto memCopyOp = dyn_cast<PimMemCopyOp>(&op)) {
           if (!pim::isNormalizedCopyOp(memCopyOp)) {
             memCopyOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
             hasFailure = true;
           }
           verifyOperand(memCopyOp.getTarget(), 0);
           verifyOperand(memCopyOp.getSource(), 1);
-      return;
+          return success();
         }
-    if (auto loadOp = dyn_cast<PimMemCopyHostToDevOp>(op)) {
+        if (auto loadOp = dyn_cast<PimMemCopyHostToDevOp>(&op)) {
           if (!pim::isNormalizedCopyOp(loadOp)) {
             loadOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
             hasFailure = true;
           }
           verifyOperand(loadOp.getDeviceTarget(), 2);
           verifyOperand(loadOp.getHostSource(), 3);
-      return;
+          return success();
         }
-    if (auto storeOp = dyn_cast<PimMemCopyDevToHostOp>(op)) {
+        if (auto storeOp = dyn_cast<PimMemCopyDevToHostOp>(&op)) {
           if (!pim::isNormalizedCopyOp(storeOp)) {
             storeOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
             hasFailure = true;
           }
           verifyOperand(storeOp.getHostTarget(), 2);
           verifyOperand(storeOp.getDeviceSource(), 3);
-      return;
+          return success();
         }
-    if (auto sendOp = dyn_cast<PimSendOp>(op)) {
+        if (auto sendOp = dyn_cast<PimSendOp>(&op)) {
           verifyOperand(sendOp.getInput(), 0);
-      return;
+          return success();
         }
-    if (auto receiveOp = dyn_cast<PimReceiveOp>(op)) {
+        if (auto receiveOp = dyn_cast<PimReceiveOp>(&op)) {
           verifyOperand(receiveOp.getOutputBuffer(), 0);
-      return;
+          return success();
         }
-    if (auto concatOp = dyn_cast<PimConcatOp>(op)) {
+        if (auto concatOp = dyn_cast<PimConcatOp>(&op)) {
           verifyOperand(concatOp.getOutputBuffer(), 0);
           for (auto inputAndIndex : llvm::enumerate(concatOp.getInputs()))
             verifyOperand(inputAndIndex.value(), inputAndIndex.index() + 1);
-      return;
+          return success();
         }
         if (isa<PimTransposeOp,
                 PimVMMOp,
@@ -330,13 +368,21 @@ LogicalResult PimBufferizationPass::verifyContiguousRuntimeOperands(ModuleOp mod
                 PimVReluOp,
                 PimVTanhOp,
                 PimVSigmOp,
-            PimVSoftmaxOp>(op)) {
-      for (auto operandAndIndex : llvm::enumerate(op->getOperands())) {
-        if (auto vmmOp = dyn_cast<PimVMMOp>(op); vmmOp && operandAndIndex.index() == 0)
+                PimVSoftmaxOp>(&op)) {
+          for (auto operandAndIndex : llvm::enumerate(op.getOperands())) {
+            if (auto vmmOp = dyn_cast<PimVMMOp>(&op); vmmOp && operandAndIndex.index() == 0)
               continue;
             verifyOperand(operandAndIndex.value(), operandAndIndex.index());
           }
         }
+        return success();
+      });
+  };
+
+  moduleOp.walk([&](pim::PimCoreOp coreOp) { verifyWithKnowledge(coreOp, seedCoreKnowledge(coreOp)); });
+  moduleOp.walk([&](pim::PimCoreBatchOp coreBatchOp) {
+    StaticValueKnowledge knowledge = seedCoreBatchKnowledge(coreBatchOp, 0);
+    verifyWithKnowledge(coreBatchOp, knowledge);
   });
 
   if (hasFailure) {
@@ -346,6 +392,31 @@ LogicalResult PimBufferizationPass::verifyContiguousRuntimeOperands(ModuleOp mod
   return success();
 }
 
+LogicalResult PimBufferizationPass::verifyPimCopyAddressSpaces(ModuleOp moduleOp) const {
+  bool hasFailure = false;
+  auto verifyWithKnowledge = [&](auto coreLikeOp, const StaticValueKnowledge& initialKnowledge) {
+    (void) walkPimCoreBlockStructurally(
+      coreLikeOp.getBody().front(), initialKnowledge, [&](Operation& op, const StaticValueKnowledge& knowledge) {
+        if (auto copyOp = dyn_cast<pim::PimMemCopyOp>(&op); copyOp && failed(verifyLoweredPimCopy(copyOp, knowledge)))
+          hasFailure = true;
+        if (auto copyOp = dyn_cast<pim::PimMemCopyHostToDevOp>(&op);
+            copyOp && failed(verifyLoweredPimCopy(copyOp, knowledge)))
+          hasFailure = true;
+        if (auto copyOp = dyn_cast<pim::PimMemCopyDevToHostOp>(&op);
+            copyOp && failed(verifyLoweredPimCopy(copyOp, knowledge)))
+          hasFailure = true;
+        return success();
+      });
+  };
+
+  moduleOp.walk([&](pim::PimCoreOp coreOp) { verifyWithKnowledge(coreOp, seedCoreKnowledge(coreOp)); });
+  moduleOp.walk([&](pim::PimCoreBatchOp coreBatchOp) {
+    StaticValueKnowledge knowledge = seedCoreBatchKnowledge(coreBatchOp, 0);
+    verifyWithKnowledge(coreBatchOp, knowledge);
+  });
+  return success(!hasFailure);
+}
+
 std::unique_ptr<Pass> createPimBufferizationPass() { return std::make_unique<PimBufferizationPass>(); }
 
 } // namespace onnx_mlir

src/PIM/Dialect/Pim/Transforms/HostConstantFolding/Patterns/Constant.cpp

@@ -96,8 +96,7 @@ struct FoldConstantCoreMapPattern final : OpRewritePattern<linalg::MapOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(linalg::MapOp mapOp, PatternRewriter& rewriter) const override {
-    auto coreOp = mapOp->getParentOfType<pim::PimCoreOp>();
-    if (!coreOp)
+    if (!mapOp->getParentOfType<pim::PimCoreOp>() && !mapOp->getParentOfType<pim::PimCoreBatchOp>())
       return failure();
 
     auto initType = dyn_cast<MemRefType>(mapOp.getInit().getType());

src/PIM/Dialect/Pim/Transforms/Verification/CMakeLists.txt

@@ -5,6 +5,7 @@ add_pim_library(OMPimVerification
 
   LINK_LIBS PUBLIC
   OMPimCommon
+  OMPimCompilerOptions
   OMPimBufferization
   PimOps
   SpatialOps

src/PIM/Dialect/Pim/Transforms/Verification/VerificationPass.cpp

@@ -5,12 +5,17 @@
 #include "mlir/Pass/Pass.h"
 
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <string>
 
 #include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/CoreBlockUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/SubviewUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/ContiguityPatterns.hpp"
@@ -143,6 +148,479 @@ static bool isHostAddressableValue(Value value, const StaticValueKnowledge& know
   return isa_and_nonnull<memref::GetGlobalOp>(base.getDefiningOp());
 }
 
+
+enum class CommunicationEventKind { Send, Receive };
+
+struct CommunicationEvent {
+  CommunicationEventKind kind = CommunicationEventKind::Send;
+  int64_t coreId = 0;
+  int64_t peerCoreId = 0;
+  int64_t size = 0;
+  uint64_t ordinal = 0;
+  std::optional<int64_t> minChannelId;
+  std::string materializer;
+  std::optional<int64_t> traceId;
+  std::optional<int64_t> commOrder;
+  std::optional<int64_t> traceClassId;
+  std::optional<int64_t> traceBlockOrdinal;
+  std::string traceKind;
+  std::string tracePhase;
+  std::string traceClassKind;
+  std::string tracePayload;
+  std::string traceMessages;
+  std::string tracePrevOp;
+  std::string traceNextOp;
+  Operation* op = nullptr;
+};
+
+using CommunicationEventVector = SmallVector<CommunicationEvent, 0>;
+
+static StringRef getCommunicationEventKindName(CommunicationEventKind kind) {
+  return kind == CommunicationEventKind::Send ? "send" : "receive";
+}
+
+constexpr StringLiteral kRaptorMinChannelIdAttr = "raptor.min_channel_id";
+constexpr StringLiteral kRaptorMaterializerAttr = "raptor.materializer";
+constexpr StringLiteral kRaptorCommOrderAttr = "raptor.comm_order";
+constexpr StringLiteral kRaptorCommTraceIdAttr = "raptor.comm_trace_id";
+constexpr StringLiteral kRaptorCommTraceKindAttr = "raptor.comm_trace_kind";
+constexpr StringLiteral kRaptorCommTracePhaseAttr = "raptor.comm_trace_phase";
+constexpr StringLiteral kRaptorCommTraceClassIdAttr = "raptor.comm_trace_class_id";
+constexpr StringLiteral kRaptorCommTraceClassKindAttr = "raptor.comm_trace_class_kind";
+constexpr StringLiteral kRaptorCommTraceBlockOrdinalAttr = "raptor.comm_trace_block_ordinal";
+constexpr StringLiteral kRaptorCommTracePayloadAttr = "raptor.comm_trace_payload";
+constexpr StringLiteral kRaptorCommTraceMessagesAttr = "raptor.comm_trace_messages";
+constexpr StringLiteral kRaptorCommTracePrevOpAttr = "raptor.comm_trace_prev_op";
+constexpr StringLiteral kRaptorCommTraceNextOpAttr = "raptor.comm_trace_next_op";
+
+static std::optional<int64_t> getNearestIntegerAttr(Operation* op, StringRef name) {
+  for (Operation* current = op; current; current = current->getParentOp())
+    if (auto attr = current->getAttrOfType<IntegerAttr>(name))
+      return attr.getInt();
+  return std::nullopt;
+}
+
+static std::string getNearestStringAttr(Operation* op, StringRef name) {
+  for (Operation* current = op; current; current = current->getParentOp())
+    if (auto attr = current->getAttrOfType<StringAttr>(name))
+      return attr.getValue().str();
+  return "";
+}
+
+static std::string formatLocation(Location loc) {
+  std::string text;
+  llvm::raw_string_ostream os(text);
+  loc.print(os);
+  return os.str();
+}
+
+static std::string formatOperationSummary(Operation* op) {
+  std::string text;
+  llvm::raw_string_ostream os(text);
+  OpPrintingFlags flags;
+  flags.skipRegions();
+  flags.elideLargeElementsAttrs();
+  op->print(os, flags);
+  return os.str();
+}
+
+static std::string formatCommunicationEvent(const CommunicationEvent& event) {
+  std::string text;
+  llvm::raw_string_ostream os(text);
+  os << "core " << event.coreId << " " << getCommunicationEventKindName(event.kind) << " "
+     << (event.kind == CommunicationEventKind::Send ? "to" : "from") << " " << event.peerCoreId
+     << " size " << event.size << "B ordinal " << event.ordinal;
+  if (event.minChannelId)
+    os << " min_channel " << *event.minChannelId;
+  if (event.commOrder)
+    os << " comm_order " << *event.commOrder;
+  if (!event.materializer.empty())
+    os << " materializer " << event.materializer;
+  if (event.traceId)
+    os << " trace#" << *event.traceId;
+  if (!event.tracePhase.empty())
+    os << " phase " << event.tracePhase;
+  if (event.traceClassId)
+    os << " class " << event.traceClassKind << "#" << *event.traceClassId;
+  if (event.traceBlockOrdinal)
+    os << " block_ordinal " << *event.traceBlockOrdinal;
+  if (!event.tracePayload.empty())
+    os << " payload " << event.tracePayload;
+  if (!event.traceMessages.empty())
+    os << " messages {" << event.traceMessages << "}";
+  if (!event.tracePrevOp.empty() || !event.traceNextOp.empty())
+    os << " inserted_between [" << event.tracePrevOp << " | " << event.traceNextOp << "]";
+  return os.str();
+}
+
+static bool areMatchedCommunicationEvents(const CommunicationEvent& lhs, const CommunicationEvent& rhs) {
+  if (lhs.coreId != rhs.peerCoreId || lhs.peerCoreId != rhs.coreId || lhs.size != rhs.size)
+    return false;
+
+  return (lhs.kind == CommunicationEventKind::Send && rhs.kind == CommunicationEventKind::Receive)
+      || (lhs.kind == CommunicationEventKind::Receive && rhs.kind == CommunicationEventKind::Send);
+}
+
+
+static std::optional<size_t> findMatchingCounterpartIndex(const CommunicationEventVector& events,
+                                                           const CommunicationEvent& event,
+                                                           size_t begin) {
+  for (size_t index = begin; index < events.size(); ++index)
+    if (areMatchedCommunicationEvents(event, events[index]))
+      return index;
+  return std::nullopt;
+}
+
+static void printCounterpartProbe(llvm::raw_ostream& os,
+                                  const DenseMap<int64_t, CommunicationEventVector>& coreEvents,
+                                  const DenseMap<int64_t, size_t>& programCounters,
+                                  const CommunicationEvent& blockedEvent) {
+  auto peerEventsIt = coreEvents.find(blockedEvent.peerCoreId);
+  if (peerEventsIt == coreEvents.end()) {
+    os << "  no local stream was collected for peer core " << blockedEvent.peerCoreId << "\n";
+    return;
+  }
+
+  const CommunicationEventVector& peerEvents = peerEventsIt->second;
+  size_t peerPc = 0;
+  auto peerPcIt = programCounters.find(blockedEvent.peerCoreId);
+  if (peerPcIt != programCounters.end())
+    peerPc = peerPcIt->second;
+
+  os << "  counterpart probe for " << formatCommunicationEvent(blockedEvent) << "\n";
+  os << "    peer core " << blockedEvent.peerCoreId << " current pc " << peerPc << " of " << peerEvents.size()
+     << "\n";
+
+  std::optional<size_t> nextMatch = findMatchingCounterpartIndex(peerEvents, blockedEvent, peerPc);
+  std::optional<size_t> anyMatch = findMatchingCounterpartIndex(peerEvents, blockedEvent, 0);
+
+  if (!nextMatch && !anyMatch) {
+    os << "    no matching counterpart exists anywhere in the peer stream\n";
+    return;
+  }
+
+  if (!nextMatch && anyMatch) {
+    os << "    matching counterpart exists only before the peer pc at ordinal " << *anyMatch
+       << "; this usually means the static stream expansion or ordering metadata is inconsistent\n";
+    os << "      " << formatCommunicationEvent(peerEvents[*anyMatch]) << "\n";
+    os << "      op: " << formatOperationSummary(peerEvents[*anyMatch].op) << "\n";
+    return;
+  }
+
+  const CommunicationEvent& match = peerEvents[*nextMatch];
+  os << "    next matching counterpart is at peer ordinal " << *nextMatch << " (distance +"
+     << (*nextMatch >= peerPc ? *nextMatch - peerPc : 0) << ")\n";
+  os << "      " << formatCommunicationEvent(match) << "\n";
+  os << "      op: " << formatOperationSummary(match.op) << "\n";
+
+  if (*nextMatch == peerPc)
+    return;
+
+  os << "    peer operations blocking before that counterpart:\n";
+  size_t end = std::min(peerEvents.size(), std::min(*nextMatch + static_cast<size_t>(1), peerPc + static_cast<size_t>(12)));
+  for (size_t index = peerPc; index < end; ++index) {
+    os << (index == peerPc ? "      pc => " : "           ") << "#" << index << " "
+       << formatCommunicationEvent(peerEvents[index]) << "\n";
+    os << "             op: " << formatOperationSummary(peerEvents[index].op) << "\n";
+  }
+  if (end <= *nextMatch)
+    os << "      ... " << (*nextMatch - end + 1) << " more peer communication event(s) before the counterpart\n";
+}
+
+static CommunicationEvent makeCommunicationEvent(CommunicationEventKind kind,
+                                                 int64_t coreId,
+                                                 int64_t peerCoreId,
+                                                 int64_t size,
+                                                 uint64_t ordinal,
+                                                 Operation* op) {
+  return CommunicationEvent {kind,
+                             coreId,
+                             peerCoreId,
+                             size,
+                             ordinal,
+                             getNearestIntegerAttr(op, kRaptorMinChannelIdAttr),
+                             getNearestStringAttr(op, kRaptorMaterializerAttr),
+                             getNearestIntegerAttr(op, kRaptorCommTraceIdAttr),
+                             getNearestIntegerAttr(op, kRaptorCommOrderAttr),
+                             getNearestIntegerAttr(op, kRaptorCommTraceClassIdAttr),
+                             getNearestIntegerAttr(op, kRaptorCommTraceBlockOrdinalAttr),
+                             getNearestStringAttr(op, kRaptorCommTraceKindAttr),
+                             getNearestStringAttr(op, kRaptorCommTracePhaseAttr),
+                             getNearestStringAttr(op, kRaptorCommTraceClassKindAttr),
+                             getNearestStringAttr(op, kRaptorCommTracePayloadAttr),
+                             getNearestStringAttr(op, kRaptorCommTraceMessagesAttr),
+                             getNearestStringAttr(op, kRaptorCommTracePrevOpAttr),
+                             getNearestStringAttr(op, kRaptorCommTraceNextOpAttr),
+                             op};
+}
+
+static LogicalResult appendCoreCommunicationEvents(Block& block,
+                                                   int64_t coreId,
+                                                   const StaticValueKnowledge& initialKnowledge,
+                                                   SmallVectorImpl<CommunicationEvent>& events,
+                                                   pim::CappedDiagnosticReporter& diagnostics) {
+  return walkPimCoreBlock(block, initialKnowledge, [&](Operation& op, const StaticValueKnowledge& knowledge) {
+    if (auto sendOp = dyn_cast<pim::PimSendOp>(&op)) {
+      auto targetCoreId = resolveIndexValue(sendOp.getTargetCoreId(), knowledge);
+      if (failed(targetCoreId)) {
+        diagnostics.report(&op, [](Operation* illegalOp) {
+          illegalOp->emitOpError("cannot statically resolve send target core for PIM communication deadlock check");
+        });
+        return failure();
+      }
+
+      events.push_back(makeCommunicationEvent(CommunicationEventKind::Send,
+                                              coreId,
+                                              *targetCoreId,
+                                              sendOp.getSize(),
+                                              static_cast<uint64_t>(events.size()),
+                                              &op));
+      return success();
+    }
+
+    if (auto receiveOp = dyn_cast<pim::PimReceiveOp>(&op)) {
+      auto sourceCoreId = resolveIndexValue(receiveOp.getSourceCoreId(), knowledge);
+      if (failed(sourceCoreId)) {
+        diagnostics.report(&op, [](Operation* illegalOp) {
+          illegalOp->emitOpError("cannot statically resolve receive source core for PIM communication deadlock check");
+        });
+        return failure();
+      }
+
+      events.push_back(makeCommunicationEvent(CommunicationEventKind::Receive,
+                                              coreId,
+                                              *sourceCoreId,
+                                              receiveOp.getSize(),
+                                              static_cast<uint64_t>(events.size()),
+                                              &op));
+      return success();
+    }
+
+    return success();
+  });
+}
+
+static void printCommunicationWindow(llvm::raw_ostream& os,
+                                     const DenseMap<int64_t, CommunicationEventVector>& coreEvents,
+                                     int64_t coreId,
+                                     size_t pc,
+                                     unsigned radius = 4) {
+  auto eventsIt = coreEvents.find(coreId);
+  if (eventsIt == coreEvents.end())
+    return;
+
+  const CommunicationEventVector& events = eventsIt->second;
+  size_t begin = pc > radius ? pc - radius : 0;
+  size_t end = std::min(events.size(), pc + static_cast<size_t>(radius) + 1);
+  os << "  local stream for core " << coreId << " around pc " << pc << " of " << events.size() << ":\n";
+  for (size_t index = begin; index < end; ++index) {
+    os << (index == pc ? "    => " : "       ") << formatCommunicationEvent(events[index]) << "\n";
+    os << "          op: " << formatOperationSummary(events[index].op) << "\n";
+  }
+}
+
+static void printCommunicationDeadlockReport(const DenseMap<int64_t, CommunicationEventVector>& coreEvents,
+                                             const DenseMap<int64_t, size_t>& programCounters,
+                                             ArrayRef<int64_t> cycle) {
+  llvm::errs() << "\n=== PIM static communication deadlock report ===\n";
+  llvm::errs() << "wait cycle:";
+  for (int64_t coreId : cycle)
+    llvm::errs() << " " << coreId;
+  if (!cycle.empty())
+    llvm::errs() << " -> " << cycle.front();
+  llvm::errs() << "\n\nblocked heads:\n";
+
+  for (int64_t coreId : cycle) {
+    auto eventsIt = coreEvents.find(coreId);
+    auto pcIt = programCounters.find(coreId);
+    if (eventsIt == coreEvents.end() || pcIt == programCounters.end() || pcIt->second >= eventsIt->second.size())
+      continue;
+    const CommunicationEvent& event = eventsIt->second[pcIt->second];
+    llvm::errs() << "  " << formatCommunicationEvent(event) << "\n";
+    llvm::errs() << "    loc: " << formatLocation(event.op->getLoc()) << "\n";
+    llvm::errs() << "    op : " << formatOperationSummary(event.op) << "\n";
+  }
+
+  llvm::errs() << "\npeer counterpart probes:\n";
+  for (int64_t coreId : cycle) {
+    auto eventsIt = coreEvents.find(coreId);
+    auto pcIt = programCounters.find(coreId);
+    if (eventsIt == coreEvents.end() || pcIt == programCounters.end() || pcIt->second >= eventsIt->second.size())
+      continue;
+    printCounterpartProbe(llvm::errs(), coreEvents, programCounters, eventsIt->second[pcIt->second]);
+  }
+
+  llvm::errs() << "\nlocal communication streams:\n";
+  for (int64_t coreId : cycle) {
+    auto pcIt = programCounters.find(coreId);
+    if (pcIt == programCounters.end())
+      continue;
+    printCommunicationWindow(llvm::errs(), coreEvents, coreId, pcIt->second);
+  }
+  llvm::errs() << "=== end PIM static communication deadlock report ===\n\n";
+}
+
+static void emitCommunicationDeadlockCycle(ModuleOp moduleOp,
+                                           const DenseMap<int64_t, CommunicationEventVector>& coreEvents,
+                                           const DenseMap<int64_t, size_t>& programCounters,
+                                           ArrayRef<int64_t> cycle) {
+  printCommunicationDeadlockReport(coreEvents, programCounters, cycle);
+
+  auto diagnostic = moduleOp.emitError()
+      << "PIM communication deadlock check found a blocking send/receive cycle while statically simulating the "
+         "expanded per-core communication streams; see the PIM static communication deadlock report above";
+
+  for (int64_t coreId : cycle) {
+    auto eventsIt = coreEvents.find(coreId);
+    auto pcIt = programCounters.find(coreId);
+    if (eventsIt == coreEvents.end() || pcIt == programCounters.end() || pcIt->second >= eventsIt->second.size())
+      continue;
+
+    const CommunicationEvent& event = eventsIt->second[pcIt->second];
+    Diagnostic& note = diagnostic.attachNote(event.op->getLoc());
+    note << formatCommunicationEvent(event);
+    if (!event.materializer.empty())
+      note << " emitted by " << event.materializer;
+  }
+}
+
+static FailureOr<SmallVector<int64_t>> findCommunicationWaitCycle(
+  const DenseMap<int64_t, CommunicationEventVector>& coreEvents,
+  const DenseMap<int64_t, size_t>& programCounters) {
+  for (const auto& [startCoreId, events] : coreEvents) {
+    auto startPcIt = programCounters.find(startCoreId);
+    if (startPcIt == programCounters.end() || startPcIt->second >= events.size())
+      continue;
+
+    DenseMap<int64_t, size_t> positionInPath;
+    SmallVector<int64_t, 8> path;
+    int64_t currentCoreId = startCoreId;
+    while (true) {
+      auto eventsIt = coreEvents.find(currentCoreId);
+      auto pcIt = programCounters.find(currentCoreId);
+      if (eventsIt == coreEvents.end() || pcIt == programCounters.end() || pcIt->second >= eventsIt->second.size())
+        break;
+
+      auto positionIt = positionInPath.find(currentCoreId);
+      if (positionIt != positionInPath.end()) {
+        SmallVector<int64_t> cycle;
+        for (size_t index = positionIt->second; index < path.size(); ++index)
+          cycle.push_back(path[index]);
+        return cycle;
+      }
+
+      positionInPath[currentCoreId] = path.size();
+      path.push_back(currentCoreId);
+      currentCoreId = eventsIt->second[pcIt->second].peerCoreId;
+    }
+  }
+
+  return failure();
+}
+
+static LogicalResult verifyNoStaticCommunicationDeadlock(ModuleOp moduleOp,
+                                                          pim::CappedDiagnosticReporter& diagnostics) {
+  DenseMap<int64_t, CommunicationEventVector> coreEvents;
+  bool hasFailure = false;
+
+  for (func::FuncOp funcOp : moduleOp.getOps<func::FuncOp>()) {
+    if (funcOp.isExternal())
+      continue;
+
+    for (Operation& op : funcOp.getBody().front().getOperations()) {
+      if (auto coreOp = dyn_cast<pim::PimCoreOp>(&op)) {
+        int64_t coreId = coreOp.getCoreId();
+        if (failed(appendCoreCommunicationEvents(
+              coreOp.getBody().front(), coreId, StaticValueKnowledge {}, coreEvents[coreId], diagnostics)))
+          hasFailure = true;
+        continue;
+      }
+
+      if (auto coreBatchOp = dyn_cast<pim::PimCoreBatchOp>(&op)) {
+        SmallVector<int32_t> coreIds = getBatchCoreIds(coreBatchOp);
+        size_t laneCount = static_cast<size_t>(coreBatchOp.getLaneCount());
+        for (size_t lane = 0; lane < laneCount; ++lane) {
+          StaticValueKnowledge laneKnowledge;
+          laneKnowledge.indexValues[coreBatchOp.getLaneArgument()] = static_cast<int64_t>(lane);
+          for (unsigned inputIndex = 0; inputIndex < coreBatchOp.getInputs().size(); ++inputIndex)
+            laneKnowledge.aliases[coreBatchOp.getInputArgument(inputIndex)] = coreBatchOp.getInputs()[inputIndex];
+
+          SmallVector<int32_t> laneCoreIds = getLaneChunkCoreIds(coreIds, laneCount, static_cast<unsigned>(lane));
+          for (int32_t coreId : laneCoreIds) {
+            if (failed(appendCoreCommunicationEvents(
+                  coreBatchOp.getBody().front(), coreId, laneKnowledge, coreEvents[coreId], diagnostics)))
+              hasFailure = true;
+          }
+        }
+      }
+    }
+  }
+
+  if (hasFailure)
+    return failure();
+
+  DenseMap<int64_t, size_t> programCounters;
+  for (const auto& [coreId, events] : coreEvents)
+    programCounters[coreId] = 0;
+
+  while (true) {
+    bool madeProgress = false;
+    for (const auto& [coreId, events] : coreEvents) {
+      size_t pc = programCounters[coreId];
+      if (pc >= events.size())
+        continue;
+
+      const CommunicationEvent& event = events[pc];
+      auto peerEventsIt = coreEvents.find(event.peerCoreId);
+      if (peerEventsIt == coreEvents.end())
+        continue;
+
+      size_t peerPc = programCounters[event.peerCoreId];
+      if (peerPc >= peerEventsIt->second.size())
+        continue;
+
+      const CommunicationEvent& peerEvent = peerEventsIt->second[peerPc];
+      if (!areMatchedCommunicationEvents(event, peerEvent))
+        continue;
+
+      ++programCounters[coreId];
+      ++programCounters[event.peerCoreId];
+      madeProgress = true;
+    }
+
+    if (madeProgress)
+      continue;
+
+    bool allDone = true;
+    for (const auto& [coreId, events] : coreEvents) {
+      if (programCounters[coreId] < events.size()) {
+        allDone = false;
+        break;
+      }
+    }
+    if (allDone)
+      return success();
+
+    auto cycle = findCommunicationWaitCycle(coreEvents, programCounters);
+    if (succeeded(cycle)) {
+      emitCommunicationDeadlockCycle(moduleOp, coreEvents, programCounters, *cycle);
+      return failure();
+    }
+
+    auto diagnostic = moduleOp.emitError()
+        << "PIM communication deadlock check stalled without finding a closed wait cycle; this usually means a "
+           "send/receive peer is missing or ordered after a finished core";
+    for (const auto& [coreId, events] : coreEvents) {
+      size_t pc = programCounters[coreId];
+      if (pc >= events.size())
+        continue;
+      const CommunicationEvent& event = events[pc];
+      diagnostic.attachNote(event.op->getLoc()) << formatCommunicationEvent(event);
+    }
+    return failure();
+  }
+}
+
 struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(VerificationPass)
 
@@ -212,11 +690,18 @@ struct VerificationPass : PassWrapper<VerificationPass, OperationPass<ModuleOp>>
       }
     }
 
+    bool hasFailure = false;
+    if (pimDetectCommunicationDeadlock && failed(verifyNoStaticCommunicationDeadlock(moduleOp, diagnostics)))
+      hasFailure = true;
+
     if (diagnostics.hasFailure()) {
       diagnostics.emitSuppressedSummary(moduleOp, "verification failures");
       moduleOp.emitError("PIM codegen verification failed; see diagnostics above");
-      signalPassFailure();
+      hasFailure = true;
     }
+
+    if (hasFailure)
+      signalPassFailure();
   }
 
 private:

src/PIM/Dialect/Spatial/CMakeLists.txt

@@ -7,10 +7,13 @@ add_pim_library(SpatialOps
   SpatialOpsVerify.cpp
   SpatialOpsCanonicalization.cpp
   ${PIM_SRC_ROOT}/Conversion/ONNXToSpatial/CompileTime.cpp
-  Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp
-  Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp
   Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp
   Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp
+  Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp
+  Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp
+  Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp
+  Transforms/MergeComputeNodes/ScheduledComputeReport.cpp
+  Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp
   Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp
   Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp
 

src/PIM/Dialect/Spatial/Spatial.td

@@ -6,6 +6,7 @@ include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/BuiltinTypes.td"
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/RegionKindInterface.td"
+include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/ParallelCombiningOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
@@ -26,8 +27,8 @@ def SpatTensor :
 // Execution
 //===----------------------------------------------------------------------===//
 
-def SpatCompute : SpatOp<"compute",
-    [SingleBlock, AttrSizedOperandSegments,
+class SpatComputeLikeBase<string mnemonic> : SpatOp<mnemonic,
+    [AttrSizedOperandSegments,
      DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmBlockArgumentNames"]>]> {
   let summary = "Compute region with attached constant weights";
 
@@ -40,8 +41,14 @@ def SpatCompute : SpatOp<"compute",
     Variadic<SpatTensor>:$outputs
   );
 
-  let regions = (region SizedRegion<1>:$body);
+  let regions = (region MinSizedRegion<1>:$body);
 
+  let hasVerifier = 1;
+  let hasFolder = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def SpatGraphCompute : SpatComputeLikeBase<"graph_compute"> {
   let extraClassDeclaration = [{
     std::optional<::mlir::BlockArgument> getWeightArgument(unsigned idx);
     std::optional<::mlir::BlockArgument> getInputArgument(unsigned idx);
@@ -50,17 +57,27 @@ def SpatCompute : SpatOp<"compute",
     std::optional<std::tuple<::mlir::Value, ::mlir::BlockArgument>>
       insertInput(unsigned idx, ::mlir::Value input, ::mlir::Location loc);
     ::llvm::SetVector<::mlir::Value, ::llvm::SmallVector<::mlir::Value, 4>, ::llvm::SmallDenseSet<::mlir::Value, 4>> getCrossbarWeights();
-    ::mlir::FailureOr<std::tuple<::mlir::OpResult, SpatCompute>>
+    ::mlir::FailureOr<std::tuple<::mlir::OpResult, SpatGraphCompute>>
       insertOutput(::mlir::RewriterBase &rewriter, unsigned idx, ::mlir::Type type, ::mlir::Location loc);
   }];
-
-  let hasVerifier = 1;
-  let hasFolder = 1;
-  let hasCustomAssemblyFormat = 1;
 }
 
-def SpatComputeBatch : SpatOp<"compute_batch",
-    [SingleBlock, AttrSizedOperandSegments,
+def SpatScheduledCompute : SpatComputeLikeBase<"scheduled_compute"> {
+  let extraClassDeclaration = [{
+    std::optional<::mlir::BlockArgument> getWeightArgument(unsigned idx);
+    std::optional<::mlir::BlockArgument> getInputArgument(unsigned idx);
+    std::optional<std::tuple<::mlir::Value, ::mlir::BlockArgument>>
+      insertWeight(unsigned idx, ::mlir::Value weight, ::mlir::Location loc);
+    std::optional<std::tuple<::mlir::Value, ::mlir::BlockArgument>>
+      insertInput(unsigned idx, ::mlir::Value input, ::mlir::Location loc);
+    ::llvm::SetVector<::mlir::Value, ::llvm::SmallVector<::mlir::Value, 4>, ::llvm::SmallDenseSet<::mlir::Value, 4>> getCrossbarWeights();
+    ::mlir::FailureOr<std::tuple<::mlir::OpResult, SpatScheduledCompute>>
+      insertOutput(::mlir::RewriterBase &rewriter, unsigned idx, ::mlir::Type type, ::mlir::Location loc);
+  }];
+}
+
+class SpatComputeBatchLikeBase<string mnemonic> : SpatOp<mnemonic,
+    [AttrSizedOperandSegments,
      DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmBlockArgumentNames"]>]> {
   let summary = "Tensor-native batch of equivalent compute lanes with shared weights and packed inputs";
 
@@ -74,8 +91,14 @@ def SpatComputeBatch : SpatOp<"compute_batch",
     Variadic<SpatTensor>:$outputs
   );
 
-  let regions = (region SizedRegion<1>:$body);
+  let regions = (region MinSizedRegion<1>:$body);
 
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def SpatGraphComputeBatch : SpatComputeBatchLikeBase<"graph_compute_batch"> {
+  let hasCanonicalizer = 1;
   let extraClassDeclaration = [{
     std::optional<::mlir::BlockArgument> getLaneArgument();
     std::optional<::mlir::BlockArgument> getWeightArgument(unsigned idx);
@@ -86,21 +109,34 @@ def SpatComputeBatch : SpatOp<"compute_batch",
     std::optional<std::tuple<::mlir::Value, ::mlir::BlockArgument>>
       insertInput(unsigned idx, ::mlir::Value input, ::mlir::Location loc);
     ::llvm::SetVector<::mlir::Value, ::llvm::SmallVector<::mlir::Value, 4>, ::llvm::SmallDenseSet<::mlir::Value, 4>> getCrossbarWeights();
-    ::mlir::FailureOr<std::tuple<::mlir::OpResult, ::mlir::BlockArgument, SpatComputeBatch>>
+    ::mlir::FailureOr<std::tuple<::mlir::OpResult, ::mlir::BlockArgument, SpatGraphComputeBatch>>
       insertOutput(::mlir::RewriterBase &rewriter, unsigned idx, ::mlir::Type type, ::mlir::Location loc);
   }];
+}
 
-  let hasVerifier = 1;
-  let hasCustomAssemblyFormat = 1;
+def SpatScheduledComputeBatch : SpatComputeBatchLikeBase<"scheduled_compute_batch"> {
+  let hasCanonicalizer = 1;
+  let extraClassDeclaration = [{
+    std::optional<::mlir::BlockArgument> getLaneArgument();
+    std::optional<::mlir::BlockArgument> getWeightArgument(unsigned idx);
+    std::optional<::mlir::BlockArgument> getInputArgument(unsigned idx);
+    std::optional<::mlir::BlockArgument> getOutputArgument(unsigned idx);
+    std::optional<std::tuple<::mlir::Value, ::mlir::BlockArgument>>
+      insertWeight(unsigned idx, ::mlir::Value weight, ::mlir::Location loc);
+    std::optional<std::tuple<::mlir::Value, ::mlir::BlockArgument>>
+      insertInput(unsigned idx, ::mlir::Value input, ::mlir::Location loc);
+    ::llvm::SetVector<::mlir::Value, ::llvm::SmallVector<::mlir::Value, 4>, ::llvm::SmallDenseSet<::mlir::Value, 4>> getCrossbarWeights();
+    ::mlir::FailureOr<std::tuple<::mlir::OpResult, ::mlir::BlockArgument, SpatScheduledComputeBatch>>
+      insertOutput(::mlir::RewriterBase &rewriter, unsigned idx, ::mlir::Type type, ::mlir::Location loc);
+  }];
 }
 
 def SpatInParallelOp : SpatOp<"in_parallel", [
        Pure,
        Terminator,
        DeclareOpInterfaceMethods<InParallelOpInterface>,
-       HasParent<"SpatComputeBatch">,
       ] # GraphRegionNoTerminator.traits> {
-  let summary = "Parallel combining terminator for resultful spat.compute_batch";
+  let summary = "Parallel combining terminator for resultful Spatial compute batches";
 
   let regions = (region SizedRegion<1>:$region);
 
@@ -128,6 +164,54 @@ def SpatYieldOp : SpatOp<"yield", [Terminator]> {
   let hasCustomAssemblyFormat = 1;
 }
 
+def SpatBlockYieldOp : SpatOp<"block_yield", [
+    Terminator,
+    DeclareOpInterfaceMethods<BranchOpInterface, ["getSuccessorForOperands"]>
+  ]> {
+  let summary = "Terminate a scheduled structural compute block";
+
+  let arguments = (ins
+    Variadic<AnyType>:$outputs
+  );
+
+  let successors = (successor
+    VariadicSuccessor<AnySuccessor>:$next
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def SpatDeferredCommunicationOp : SpatOp<"deferred_communication", [SingleBlock]> {
+  let summary = "Temporary scheduled communication placeholder";
+
+  let arguments = (ins
+    Variadic<SpatTensor>:$sources,
+    StrAttr:$exchangeId,
+    StrAttr:$logicalProducer,
+    StrAttr:$logicalConsumer,
+    I64Attr:$sourceClass,
+    I64Attr:$targetClass,
+    I64Attr:$sourceCore,
+    I64Attr:$targetCore,
+    I64Attr:$sourceLane,
+    I64Attr:$targetLane,
+    StrAttr:$transferKind,
+    I64Attr:$resultIndex,
+    DenseI64ArrayAttr:$projectedTransfer,
+    I64Attr:$hostOutputOwner
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let regions = (region SizedRegion<1>:$body);
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
 def SpatExtractRowsOp : SpatOp<"extract_rows", []> {
   let summary = "Extract every row of a rank-2 tensor as separate rank-2 row tensors";
 
@@ -159,6 +243,106 @@ def SpatConcatOp : SpatOp<"concat", []> {
   let hasCustomAssemblyFormat = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// Planning
+//===----------------------------------------------------------------------===//
+
+def SpatConv2DPlanOp : SpatOp<"conv2d_plan", []> {
+  let summary = "Structured Conv2D planning op that preserves logical ONNX geometry";
+
+  let arguments = (ins
+    SpatTensor:$input,
+    SpatTensor:$weight,
+    Optional<SpatTensor>:$bias,
+    DenseI64ArrayAttr:$pads,
+    DenseI64ArrayAttr:$strides,
+    DenseI64ArrayAttr:$dilations,
+    I64Attr:$group,
+    StrAttr:$logicalLayout
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+}
+
+def SpatReluPlanOp : SpatOp<"relu_plan", []> {
+  let summary = "Layout-aware ReLU planning op";
+
+  let arguments = (ins
+    SpatTensor:$input,
+    StrAttr:$logicalLayout
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+}
+
+def SpatBiasAddPlanOp : SpatOp<"bias_add_plan", []> {
+  let summary = "Layout-aware Conv-style bias add planning op";
+
+  let arguments = (ins
+    SpatTensor:$input,
+    SpatTensor:$bias,
+    StrAttr:$logicalLayout
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+}
+
+def SpatBlueprintOp : SpatOp<"blueprint", []> {
+  let summary = "Blueprint for assembling logical tensors from published fragments";
+
+  let arguments = (ins
+    SpatTensor:$input,
+    Variadic<SpatTensor>:$fragments,
+    StrAttr:$logicalLayout,
+    StrAttr:$physicalLayout,
+    DenseI64ArrayAttr:$fragmentOffsets,
+    DenseI64ArrayAttr:$fragmentSizes,
+    StrAttr:$indexMap,
+    OptionalAttr<StrAttr>:$mode,
+    OptionalAttr<DenseI64ArrayAttr>:$fragmentOperandIndices,
+    OptionalAttr<DenseI64ArrayAttr>:$fragmentSourceOffsets,
+    OptionalAttr<DenseI64ArrayAttr>:$fragmentStrides,
+    OptionalAttr<StrAttr>:$conflictPolicy,
+    OptionalAttr<StrAttr>:$coveragePolicy
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def SpatMaterializeLayoutOp : SpatOp<"materialize_layout", []> {
+  let summary = "Explicit layout conversion or materialization barrier";
+
+  let arguments = (ins
+    SpatTensor:$input,
+    StrAttr:$logicalLayout,
+    StrAttr:$sourcePhysicalLayout,
+    StrAttr:$targetPhysicalLayout
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // Communication
 //===----------------------------------------------------------------------===//

src/PIM/Dialect/Spatial/SpatialOps.cpp

@@ -29,11 +29,19 @@ std::optional<BlockArgument> insertBlockArgument(Region& body, unsigned argIdx,
 }
 
 void setComputeOperandSegmentSizes(Operation* op, int32_t weightCount, int32_t inputCount) {
-  if (auto compute = dyn_cast<SpatCompute>(op)) {
+  if (auto compute = dyn_cast<SpatGraphCompute>(op)) {
     compute.getProperties().setOperandSegmentSizes({weightCount, inputCount});
     return;
   }
-  cast<SpatComputeBatch>(op).getProperties().setOperandSegmentSizes({weightCount, inputCount});
+  if (auto compute = dyn_cast<SpatScheduledCompute>(op)) {
+    compute.getProperties().setOperandSegmentSizes({weightCount, inputCount});
+    return;
+  }
+  if (auto batch = dyn_cast<SpatGraphComputeBatch>(op)) {
+    batch.getProperties().setOperandSegmentSizes({weightCount, inputCount});
+    return;
+  }
+  cast<SpatScheduledComputeBatch>(op).getProperties().setOperandSegmentSizes({weightCount, inputCount});
 }
 
 using CrossbarWeightSet = llvm::SetVector<Value, llvm::SmallVector<Value, 4>, llvm::SmallDenseSet<Value, 4>>;
@@ -47,116 +55,214 @@ CrossbarWeightSet collectCrossbarWeights(Region& body) {
   return weights;
 }
 
-} // namespace
-
-std::optional<BlockArgument> SpatCompute::getWeightArgument(unsigned idx) { return getBlockArgument(getBody(), idx); }
-
-std::optional<BlockArgument> SpatCompute::getInputArgument(unsigned idx) {
-  return getBlockArgument(getBody(), getWeights().size() + idx);
+template <typename ComputeOpTy>
+std::optional<BlockArgument> getComputeWeightArgument(ComputeOpTy compute, unsigned idx) {
+  return getBlockArgument(compute.getBody(), idx);
 }
 
-std::optional<std::tuple<Value, BlockArgument>> SpatCompute::insertWeight(unsigned idx, Value weight, Location loc) {
-  if (auto existing = llvm::find(getWeights(), weight); existing != getWeights().end()) {
-    auto index = std::distance(getWeights().begin(), existing);
-    return {
-      {*existing, *getWeightArgument(index)}
-    };
+template <typename ComputeOpTy>
+std::optional<BlockArgument> getComputeInputArgument(ComputeOpTy compute, unsigned idx) {
+  return getBlockArgument(compute.getBody(), compute.getWeights().size() + idx);
 }
 
-  unsigned weightCount = getWeights().size();
-  unsigned inputCount = getInputs().size();
-  getOperation()->insertOperands(idx, ValueRange {weight});
+template <typename ComputeOpTy>
+std::optional<std::tuple<Value, BlockArgument>>
+insertComputeWeight(ComputeOpTy compute, unsigned idx, Value weight, Location loc) {
+  if (auto existing = llvm::find(compute.getWeights(), weight); existing != compute.getWeights().end()) {
+    auto index = std::distance(compute.getWeights().begin(), existing);
+    return {{*existing, *getComputeWeightArgument(compute, index)}};
+  }
+
+  unsigned weightCount = compute.getWeights().size();
+  unsigned inputCount = compute.getInputs().size();
+  compute.getOperation()->insertOperands(idx, ValueRange {weight});
   setComputeOperandSegmentSizes(
-    getOperation(), static_cast<int32_t>(weightCount + 1), static_cast<int32_t>(inputCount));
-  auto blockArg = insertBlockArgument(getBody(), idx, weight.getType(), loc);
+    compute.getOperation(), static_cast<int32_t>(weightCount + 1), static_cast<int32_t>(inputCount));
+  auto blockArg = insertBlockArgument(compute.getBody(), idx, weight.getType(), loc);
   if (!blockArg)
     return std::nullopt;
-  return std::make_tuple(getOperation()->getOperand(idx), *blockArg);
+  return std::make_tuple(compute.getOperation()->getOperand(idx), *blockArg);
 }
 
-std::optional<std::tuple<Value, BlockArgument>> SpatCompute::insertInput(unsigned idx, Value input, Location loc) {
-  unsigned weightCount = getWeights().size();
-  unsigned inputCount = getInputs().size();
-  getOperation()->insertOperands(weightCount + idx, ValueRange {input});
+template <typename ComputeBatchOpTy>
+std::optional<std::tuple<Value, BlockArgument>>
+insertComputeBatchWeight(ComputeBatchOpTy batch, unsigned idx, Value weight, Location loc) {
+  if (auto existing = llvm::find(batch.getWeights(), weight); existing != batch.getWeights().end()) {
+    auto index = std::distance(batch.getWeights().begin(), existing);
+    return {{*existing, *batch.getWeightArgument(index)}};
+  }
+
+  unsigned weightCount = batch.getWeights().size();
+  unsigned inputCount = batch.getInputs().size();
+  batch.getOperation()->insertOperands(idx, ValueRange {weight});
   setComputeOperandSegmentSizes(
-    getOperation(), static_cast<int32_t>(weightCount), static_cast<int32_t>(inputCount + 1));
-  auto blockArg = insertBlockArgument(getBody(), weightCount + idx, input.getType(), loc);
+    batch.getOperation(), static_cast<int32_t>(weightCount + 1), static_cast<int32_t>(inputCount));
+
+  auto blockArg = insertBlockArgument(batch.getBody(), 1 + idx, weight.getType(), loc);
   if (!blockArg)
     return std::nullopt;
-  return std::make_tuple(getOperation()->getOperand(weightCount + idx), *blockArg);
+  return std::make_tuple(batch.getOperation()->getOperand(idx), *blockArg);
 }
 
-CrossbarWeightSet SpatCompute::getCrossbarWeights() { return collectCrossbarWeights(getBody()); }
-
-FailureOr<std::tuple<OpResult, SpatCompute>>
-SpatCompute::insertOutput(RewriterBase& rewriter, unsigned idx, Type type, Location loc) {
-  if (idx > getNumResults())
-    return failure();
-
-  rewriter.setInsertionPoint(getOperation());
-  SmallVector<Type> resultTypes(getResultTypes().begin(), getResultTypes().end());
-  resultTypes.insert(resultTypes.begin() + idx, type);
-  auto newCompute = SpatCompute::create(rewriter, getLoc(), TypeRange(resultTypes), getWeights(), getInputs());
-  newCompute->setAttrs((*this)->getAttrs());
-  setComputeOperandSegmentSizes(newCompute.getOperation(),
-                                static_cast<int32_t>(newCompute.getWeights().size()),
-                                static_cast<int32_t>(newCompute.getInputs().size()));
-  rewriter.inlineRegionBefore(getBody(), newCompute.getBody(), newCompute.getBody().end());
-  for (unsigned oldResultIdx = 0; oldResultIdx < getNumResults(); ++oldResultIdx)
-    getResult(oldResultIdx)
-      .replaceAllUsesWith(newCompute.getResult(oldResultIdx < idx ? oldResultIdx : oldResultIdx + 1));
-  rewriter.eraseOp(getOperation());
-  return std::make_tuple(cast<OpResult>(newCompute.getResult(idx)), newCompute);
+template <typename ComputeOpTy>
+std::optional<std::tuple<Value, BlockArgument>>
+insertComputeInput(ComputeOpTy compute, unsigned idx, Value input, Location loc) {
+  unsigned weightCount = compute.getWeights().size();
+  unsigned inputCount = compute.getInputs().size();
+  compute.getOperation()->insertOperands(weightCount + idx, ValueRange {input});
+  setComputeOperandSegmentSizes(
+    compute.getOperation(), static_cast<int32_t>(weightCount), static_cast<int32_t>(inputCount + 1));
+  auto blockArg = insertBlockArgument(compute.getBody(), weightCount + idx, input.getType(), loc);
+  if (!blockArg)
+    return std::nullopt;
+  return std::make_tuple(compute.getOperation()->getOperand(weightCount + idx), *blockArg);
 }
 
-void SpatCompute::getAsmBlockArgumentNames(Region& region, OpAsmSetValueNameFn setNameFn) {
+template <typename ComputeOpTy>
+void setComputeAsmBlockArgumentNames(ComputeOpTy compute, Region& region, OpAsmSetValueNameFn setNameFn) {
   if (region.empty())
     return;
 
-  for (unsigned index = 0; index < getWeights().size(); ++index)
-    if (auto weightArg = getWeightArgument(index))
+  for (unsigned index = 0; index < compute.getWeights().size(); ++index)
+    if (auto weightArg = compute.getWeightArgument(index))
       setNameFn(*weightArg, ("w" + std::to_string(index)).c_str());
 
-  for (unsigned index = 0; index < getInputs().size(); ++index)
-    if (auto inputArg = getInputArgument(index))
+  for (unsigned index = 0; index < compute.getInputs().size(); ++index)
+    if (auto inputArg = compute.getInputArgument(index))
       setNameFn(*inputArg, ("in" + std::to_string(index)).c_str());
 }
 
-std::optional<BlockArgument> SpatComputeBatch::getLaneArgument() { return getBlockArgument(getBody(), 0); }
+template <typename ComputeOpTy>
+FailureOr<std::tuple<OpResult, ComputeOpTy>>
+insertComputeOutput(ComputeOpTy compute, RewriterBase& rewriter, unsigned idx, Type type, Location loc) {
+  if (idx > compute.getNumResults())
+    return failure();
 
-std::optional<BlockArgument> SpatComputeBatch::getWeightArgument(unsigned idx) {
+  rewriter.setInsertionPoint(compute.getOperation());
+  SmallVector<Type> resultTypes(compute.getResultTypes().begin(), compute.getResultTypes().end());
+  resultTypes.insert(resultTypes.begin() + idx, type);
+  auto newCompute =
+    ComputeOpTy::create(rewriter, compute.getLoc(), TypeRange(resultTypes), compute.getWeights(), compute.getInputs());
+  newCompute->setAttrs(compute->getAttrs());
+  setComputeOperandSegmentSizes(newCompute.getOperation(),
+                                static_cast<int32_t>(newCompute.getWeights().size()),
+                                static_cast<int32_t>(newCompute.getInputs().size()));
+  rewriter.inlineRegionBefore(compute.getBody(), newCompute.getBody(), newCompute.getBody().end());
+  for (unsigned oldResultIdx = 0; oldResultIdx < compute.getNumResults(); ++oldResultIdx)
+    compute.getResult(oldResultIdx)
+      .replaceAllUsesWith(newCompute.getResult(oldResultIdx < idx ? oldResultIdx : oldResultIdx + 1));
+  rewriter.eraseOp(compute.getOperation());
+  return std::make_tuple(cast<OpResult>(newCompute.getResult(idx)), newCompute);
+}
+
+template <typename ComputeBatchOpTy>
+FailureOr<std::tuple<OpResult, BlockArgument, ComputeBatchOpTy>>
+insertComputeBatchOutput(ComputeBatchOpTy batch, RewriterBase& rewriter, unsigned idx, Type type, Location loc) {
+  if (idx > batch.getNumResults())
+    return failure();
+
+  rewriter.setInsertionPoint(batch.getOperation());
+  SmallVector<Type> resultTypes(batch.getResultTypes().begin(), batch.getResultTypes().end());
+  resultTypes.insert(resultTypes.begin() + idx, type);
+  auto newBatch =
+    ComputeBatchOpTy::create(rewriter, batch.getLoc(), TypeRange(resultTypes), batch.getLaneCountAttr(), batch.getWeights(), batch.getInputs());
+  newBatch->setAttrs(batch->getAttrs());
+  setComputeOperandSegmentSizes(newBatch.getOperation(),
+                                static_cast<int32_t>(newBatch.getWeights().size()),
+                                static_cast<int32_t>(newBatch.getInputs().size()));
+  rewriter.inlineRegionBefore(batch.getBody(), newBatch.getBody(), newBatch.getBody().end());
+  if (newBatch.getBody().empty()) {
+    rewriter.eraseOp(newBatch);
+    return failure();
+  }
+  auto blockArg = newBatch.getBody().front().insertArgument(
+    1 + newBatch.getWeights().size() + newBatch.getInputs().size() + idx, type, loc);
+  for (unsigned oldResultIdx = 0; oldResultIdx < batch.getNumResults(); ++oldResultIdx)
+    batch.getResult(oldResultIdx)
+      .replaceAllUsesWith(newBatch.getResult(oldResultIdx < idx ? oldResultIdx : oldResultIdx + 1));
+  rewriter.eraseOp(batch.getOperation());
+  return std::make_tuple(cast<OpResult>(newBatch.getResult(idx)), blockArg, newBatch);
+}
+
+} // namespace
+
+bool isGraphComputeLike(Operation* op) { return isa<SpatGraphCompute, SpatGraphComputeBatch>(op); }
+
+bool isGraphBatchComputeLike(Operation* op) { return isa<SpatGraphComputeBatch>(op); }
+
+bool isScheduledComputeLike(Operation* op) { return isa<SpatScheduledCompute, SpatScheduledComputeBatch>(op); }
+
+bool isScheduledBatchComputeLike(Operation* op) { return isa<SpatScheduledComputeBatch>(op); }
+
+bool isAnySpatialComputeLike(Operation* op) {
+  return isa<SpatGraphCompute, SpatGraphComputeBatch, SpatScheduledCompute, SpatScheduledComputeBatch>(op);
+}
+
+bool isAnySpatialComputeBatchLike(Operation* op) { return isa<SpatGraphComputeBatch, SpatScheduledComputeBatch>(op); }
+
+std::optional<BlockArgument> SpatGraphCompute::getWeightArgument(unsigned idx) { return getComputeWeightArgument(*this, idx); }
+std::optional<BlockArgument> SpatGraphCompute::getInputArgument(unsigned idx) { return getComputeInputArgument(*this, idx); }
+std::optional<std::tuple<Value, BlockArgument>> SpatGraphCompute::insertWeight(unsigned idx, Value weight, Location loc) {
+  return insertComputeWeight(*this, idx, weight, loc);
+}
+std::optional<std::tuple<Value, BlockArgument>> SpatGraphCompute::insertInput(unsigned idx, Value input, Location loc) {
+  return insertComputeInput(*this, idx, input, loc);
+}
+CrossbarWeightSet SpatGraphCompute::getCrossbarWeights() { return collectCrossbarWeights(getBody()); }
+FailureOr<std::tuple<OpResult, SpatGraphCompute>>
+SpatGraphCompute::insertOutput(RewriterBase& rewriter, unsigned idx, Type type, Location loc) {
+  return insertComputeOutput(*this, rewriter, idx, type, loc);
+}
+void SpatGraphCompute::getAsmBlockArgumentNames(Region& region, OpAsmSetValueNameFn setNameFn) {
+  setComputeAsmBlockArgumentNames(*this, region, setNameFn);
+}
+
+std::optional<BlockArgument> SpatScheduledCompute::getWeightArgument(unsigned idx) {
+  return getComputeWeightArgument(*this, idx);
+}
+std::optional<BlockArgument> SpatScheduledCompute::getInputArgument(unsigned idx) { return getComputeInputArgument(*this, idx); }
+std::optional<std::tuple<Value, BlockArgument>>
+SpatScheduledCompute::insertWeight(unsigned idx, Value weight, Location loc) {
+  return insertComputeWeight(*this, idx, weight, loc);
+}
+std::optional<std::tuple<Value, BlockArgument>>
+SpatScheduledCompute::insertInput(unsigned idx, Value input, Location loc) {
+  return insertComputeInput(*this, idx, input, loc);
+}
+CrossbarWeightSet SpatScheduledCompute::getCrossbarWeights() { return collectCrossbarWeights(getBody()); }
+FailureOr<std::tuple<OpResult, SpatScheduledCompute>>
+SpatScheduledCompute::insertOutput(RewriterBase& rewriter, unsigned idx, Type type, Location loc) {
+  return insertComputeOutput(*this, rewriter, idx, type, loc);
+}
+void SpatScheduledCompute::getAsmBlockArgumentNames(Region& region, OpAsmSetValueNameFn setNameFn) {
+  setComputeAsmBlockArgumentNames(*this, region, setNameFn);
+}
+
+SuccessorOperands SpatBlockYieldOp::getSuccessorOperands(unsigned index) {
+  assert(index == 0 && "invalid successor index");
+  return SuccessorOperands(getOutputsMutable());
+}
+
+Block* SpatBlockYieldOp::getSuccessorForOperands(ArrayRef<Attribute>) {
+  return getOperation()->getNumSuccessors() == 0 ? nullptr : getOperation()->getSuccessor(0);
+}
+
+std::optional<BlockArgument> SpatGraphComputeBatch::getLaneArgument() { return getBlockArgument(getBody(), 0); }
+std::optional<BlockArgument> SpatGraphComputeBatch::getWeightArgument(unsigned idx) {
   return getBlockArgument(getBody(), 1 + idx);
 }
-
-std::optional<BlockArgument> SpatComputeBatch::getInputArgument(unsigned idx) {
+std::optional<BlockArgument> SpatGraphComputeBatch::getInputArgument(unsigned idx) {
   return getBlockArgument(getBody(), 1 + getWeights().size() + idx);
 }
-
-std::optional<BlockArgument> SpatComputeBatch::getOutputArgument(unsigned idx) {
+std::optional<BlockArgument> SpatGraphComputeBatch::getOutputArgument(unsigned idx) {
   return getBlockArgument(getBody(), 1 + getWeights().size() + getInputs().size() + idx);
 }
-
 std::optional<std::tuple<Value, BlockArgument>>
-SpatComputeBatch::insertWeight(unsigned idx, Value weight, Location loc) {
-  if (auto existing = llvm::find(getWeights(), weight); existing != getWeights().end()) {
-    auto index = std::distance(getWeights().begin(), existing);
-    return {
-      {*existing, *getWeightArgument(index)}
-    };
+SpatGraphComputeBatch::insertWeight(unsigned idx, Value weight, Location loc) {
+  return insertComputeBatchWeight(*this, idx, weight, loc);
 }
-
-  unsigned weightCount = getWeights().size();
-  unsigned inputCount = getInputs().size();
-  getOperation()->insertOperands(idx, ValueRange {weight});
-  setComputeOperandSegmentSizes(
-    getOperation(), static_cast<int32_t>(weightCount + 1), static_cast<int32_t>(inputCount));
-  auto blockArg = insertBlockArgument(getBody(), 1 + idx, weight.getType(), loc);
-  if (!blockArg)
-    return std::nullopt;
-  return std::make_tuple(getOperation()->getOperand(idx), *blockArg);
-}
-
-std::optional<std::tuple<Value, BlockArgument>> SpatComputeBatch::insertInput(unsigned idx, Value input, Location loc) {
+std::optional<std::tuple<Value, BlockArgument>>
+SpatGraphComputeBatch::insertInput(unsigned idx, Value input, Location loc) {
   unsigned weightCount = getWeights().size();
   unsigned inputCount = getInputs().size();
   getOperation()->insertOperands(weightCount + idx, ValueRange {input});
@@ -167,52 +273,68 @@ std::optional<std::tuple<Value, BlockArgument>> SpatComputeBatch::insertInput(un
     return std::nullopt;
   return std::make_tuple(getOperation()->getOperand(weightCount + idx), *blockArg);
 }
-
-CrossbarWeightSet SpatComputeBatch::getCrossbarWeights() { return collectCrossbarWeights(getBody()); }
-
-FailureOr<std::tuple<OpResult, BlockArgument, SpatComputeBatch>>
-SpatComputeBatch::insertOutput(RewriterBase& rewriter, unsigned idx, Type type, Location loc) {
-  if (idx > getNumResults())
-    return failure();
-
-  rewriter.setInsertionPoint(getOperation());
-  SmallVector<Type> resultTypes(getResultTypes().begin(), getResultTypes().end());
-  resultTypes.insert(resultTypes.begin() + idx, type);
-  auto newBatch =
-    SpatComputeBatch::create(rewriter, getLoc(), TypeRange(resultTypes), getLaneCountAttr(), getWeights(), getInputs());
-  newBatch->setAttrs((*this)->getAttrs());
-  setComputeOperandSegmentSizes(newBatch.getOperation(),
-                                static_cast<int32_t>(newBatch.getWeights().size()),
-                                static_cast<int32_t>(newBatch.getInputs().size()));
-  rewriter.inlineRegionBefore(getBody(), newBatch.getBody(), newBatch.getBody().end());
-  if (newBatch.getBody().empty()) {
-    rewriter.eraseOp(newBatch);
-    return failure();
+CrossbarWeightSet SpatGraphComputeBatch::getCrossbarWeights() { return collectCrossbarWeights(getBody()); }
+FailureOr<std::tuple<OpResult, BlockArgument, SpatGraphComputeBatch>>
+SpatGraphComputeBatch::insertOutput(RewriterBase& rewriter, unsigned idx, Type type, Location loc) {
+  return insertComputeBatchOutput(*this, rewriter, idx, type, loc);
 }
-  auto blockArg = newBatch.getBody().front().insertArgument(
-    1 + newBatch.getWeights().size() + newBatch.getInputs().size() + idx, type, loc);
-  for (unsigned oldResultIdx = 0; oldResultIdx < getNumResults(); ++oldResultIdx)
-    getResult(oldResultIdx)
-      .replaceAllUsesWith(newBatch.getResult(oldResultIdx < idx ? oldResultIdx : oldResultIdx + 1));
-  rewriter.eraseOp(getOperation());
-  return std::make_tuple(cast<OpResult>(newBatch.getResult(idx)), blockArg, newBatch);
-}
-
-void SpatComputeBatch::getAsmBlockArgumentNames(Region& region, OpAsmSetValueNameFn setNameFn) {
+void SpatGraphComputeBatch::getAsmBlockArgumentNames(Region& region, OpAsmSetValueNameFn setNameFn) {
   if (region.empty())
     return;
 
   if (auto laneArg = getLaneArgument())
     setNameFn(*laneArg, "lane");
+  setComputeAsmBlockArgumentNames(*this, region, setNameFn);
+  for (unsigned index = 0; index < getNumResults(); ++index) {
+    auto outputArg = getOutputArgument(index);
+    if (!outputArg)
+      continue;
+    if (index == 0) {
+      setNameFn(*outputArg, "out");
+      continue;
+    }
+    setNameFn(*outputArg, ("out" + std::to_string(index)).c_str());
+  }
+}
 
-  for (unsigned index = 0; index < getWeights().size(); ++index)
-    if (auto weightArg = getWeightArgument(index))
-      setNameFn(*weightArg, ("w" + std::to_string(index)).c_str());
-
-  for (unsigned index = 0; index < getInputs().size(); ++index)
-    if (auto inputArg = getInputArgument(index))
-      setNameFn(*inputArg, ("in" + std::to_string(index)).c_str());
+std::optional<BlockArgument> SpatScheduledComputeBatch::getLaneArgument() { return getBlockArgument(getBody(), 0); }
+std::optional<BlockArgument> SpatScheduledComputeBatch::getWeightArgument(unsigned idx) {
+  return getBlockArgument(getBody(), 1 + idx);
+}
+std::optional<BlockArgument> SpatScheduledComputeBatch::getInputArgument(unsigned idx) {
+  return getBlockArgument(getBody(), 1 + getWeights().size() + idx);
+}
+std::optional<BlockArgument> SpatScheduledComputeBatch::getOutputArgument(unsigned idx) {
+  return getBlockArgument(getBody(), 1 + getWeights().size() + getInputs().size() + idx);
+}
+std::optional<std::tuple<Value, BlockArgument>>
+SpatScheduledComputeBatch::insertWeight(unsigned idx, Value weight, Location loc) {
+  return insertComputeBatchWeight(*this, idx, weight, loc);
+}
+std::optional<std::tuple<Value, BlockArgument>>
+SpatScheduledComputeBatch::insertInput(unsigned idx, Value input, Location loc) {
+  unsigned weightCount = getWeights().size();
+  unsigned inputCount = getInputs().size();
+  getOperation()->insertOperands(weightCount + idx, ValueRange {input});
+  setComputeOperandSegmentSizes(
+    getOperation(), static_cast<int32_t>(weightCount), static_cast<int32_t>(inputCount + 1));
+  auto blockArg = insertBlockArgument(getBody(), 1 + weightCount + idx, input.getType(), loc);
+  if (!blockArg)
+    return std::nullopt;
+  return std::make_tuple(getOperation()->getOperand(weightCount + idx), *blockArg);
+}
+CrossbarWeightSet SpatScheduledComputeBatch::getCrossbarWeights() { return collectCrossbarWeights(getBody()); }
+FailureOr<std::tuple<OpResult, BlockArgument, SpatScheduledComputeBatch>>
+SpatScheduledComputeBatch::insertOutput(RewriterBase& rewriter, unsigned idx, Type type, Location loc) {
+  return insertComputeBatchOutput(*this, rewriter, idx, type, loc);
+}
+void SpatScheduledComputeBatch::getAsmBlockArgumentNames(Region& region, OpAsmSetValueNameFn setNameFn) {
+  if (region.empty())
+    return;
 
+  if (auto laneArg = getLaneArgument())
+    setNameFn(*laneArg, "lane");
+  setComputeAsmBlockArgumentNames(*this, region, setNameFn);
   for (unsigned index = 0; index < getNumResults(); ++index) {
     auto outputArg = getOutputArgument(index);
     if (!outputArg)
@@ -231,7 +353,11 @@ void SpatInParallelOp::build(OpBuilder& builder, OperationState& result) {
   builder.createBlock(bodyRegion);
 }
 
-OpResult SpatInParallelOp::getParentResult(int64_t idx) { return getOperation()->getParentOp()->getResult(idx); }
+OpResult SpatInParallelOp::getParentResult(int64_t idx) {
+  Operation* parent = getOperation()->getParentOp();
+  assert(isAnySpatialComputeBatchLike(parent) && "expected Spatial compute batch parent");
+  return parent->getResult(idx);
+}
 
 llvm::iterator_range<Block::iterator> SpatInParallelOp::getYieldingOps() { return getRegion().front().getOperations(); }
 

src/PIM/Dialect/Spatial/SpatialOps.hpp

@@ -26,3 +26,19 @@
 
 #define GET_OP_CLASSES
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp.inc"
+
+namespace onnx_mlir {
+namespace spatial {
+
+bool isGraphComputeLike(mlir::Operation* op);
+bool isGraphBatchComputeLike(mlir::Operation* op);
+bool isScheduledComputeLike(mlir::Operation* op);
+bool isScheduledBatchComputeLike(mlir::Operation* op);
+bool isAnySpatialComputeLike(mlir::Operation* op);
+bool isAnySpatialComputeBatchLike(mlir::Operation* op);
+
+using SpatCompute = SpatGraphCompute;
+using SpatComputeBatch = SpatGraphComputeBatch;
+
+} // namespace spatial
+} // namespace onnx_mlir

src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp

@@ -32,6 +32,14 @@ static IntegerAttr getI32Attr(OpAsmParser& parser, int32_t value) {
   return parser.getBuilder().getI32IntegerAttr(value);
 }
 
+static ParseResult parseBareStringAttr(OpAsmParser& parser, StringAttr& attr) {
+  StringRef value;
+  if (parser.parseKeyword(&value))
+    return failure();
+  attr = parser.getBuilder().getStringAttr(value);
+  return success();
+}
+
 static void printBlockArgumentList(OpAsmPrinter& printer, ArrayRef<BlockArgument> arguments) {
   printer << "(";
   for (auto [index, argument] : llvm::enumerate(arguments)) {
@@ -115,6 +123,254 @@ static ParseResult parseBoundValueList(OpAsmParser& parser,
   return success();
 }
 
+template <typename ComputeOpTy>
+void printComputeLikeOp(ComputeOpTy op, OpAsmPrinter& printer) {
+  SmallVector<Value> weightArgs;
+  weightArgs.reserve(op.getWeights().size());
+  for (unsigned index = 0; index < op.getWeights().size(); ++index) {
+    auto weightArg = op.getWeightArgument(index);
+    if (!weightArg)
+      return printer.printGenericOp(op.getOperation(), /*printOpName=*/false);
+    weightArgs.push_back(*weightArg);
+  }
+  SmallVector<Value> inputArgs;
+  inputArgs.reserve(op.getInputs().size());
+  for (unsigned index = 0; index < op.getInputs().size(); ++index) {
+    auto inputArg = op.getInputArgument(index);
+    if (!inputArg)
+      return printer.printGenericOp(op.getOperation(), /*printOpName=*/false);
+    inputArgs.push_back(*inputArg);
+  }
+
+  printer << " ";
+  printBoundValueList(printer, weightArgs, op.getWeights(), ListDelimiter::Square);
+  printer << " ";
+  printBoundValueList(printer, inputArgs, op.getInputs(), ListDelimiter::Paren);
+
+  if (auto coreIdAttr = op->template getAttrOfType<IntegerAttr>(onnx_mlir::kCoreIdAttrName))
+    printer << " coreId " << coreIdAttr.getInt();
+  printer << " crossbarWeights " << collectDistinctCrossbarWeights(op.getOperation()).size();
+
+  printer.printOptionalAttrDict(op->getAttrs(), {op.getOperandSegmentSizesAttrName().getValue(), onnx_mlir::kCoreIdAttrName});
+
+  printer << " : ";
+  printCompressedTypeList(printer, TypeRange(op.getWeights()), ListDelimiter::Square);
+  printer << " ";
+  printCompressedTypeList(printer, TypeRange(op.getInputs()), ListDelimiter::Paren);
+  printer << " -> ";
+  printCompressedTypeSequence(printer, op.getResultTypes());
+  printer << " ";
+  printer.printRegion(op.getBody(), /*printEntryBlockArgs=*/!op.getBody().hasOneBlock());
+}
+
+template <typename ComputeOpTy>
+ParseResult parseComputeLikeOp(OpAsmParser& parser, OperationState& result) {
+  SmallVector<OpAsmParser::Argument> weightArgs;
+  SmallVector<OpAsmParser::Argument> regionArgs;
+  SmallVector<OpAsmParser::UnresolvedOperand> weights;
+  SmallVector<OpAsmParser::UnresolvedOperand> inputs;
+  SmallVector<Type> weightTypes;
+  SmallVector<Type> inputTypes;
+  SmallVector<Type> outputTypes;
+  int32_t crossbarWeightCount = 0;
+  int32_t coreId = 0;
+
+  if (parseBoundValueList(parser, ListDelimiter::Square, weightArgs, weights))
+    return failure();
+
+  SmallVector<OpAsmParser::Argument> inputArgs;
+  if (parseBoundValueList(parser, ListDelimiter::Paren, inputArgs, inputs))
+    return failure();
+
+  bool hasCoreId = parseOptionalKeywordAlias(parser, "coreId", "core_id");
+  if (hasCoreId && parser.parseInteger(coreId))
+    return failure();
+
+  bool hasCrossbarWeightCount = parseOptionalKeywordAlias(parser, "crossbarWeights", "crossbar_weights");
+  if (hasCrossbarWeightCount && parser.parseInteger(crossbarWeightCount))
+    return failure();
+  (void) crossbarWeightCount;
+
+  if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()
+      || parseCompressedRepeatedList(
+        parser, ListDelimiter::Square, weightTypes, [&](Type& type) { return parser.parseType(type); })
+      || parseCompressedRepeatedList(
+        parser, ListDelimiter::Paren, inputTypes, [&](Type& type) { return parser.parseType(type); })
+      || parser.parseArrow() || parseCompressedTypeSequence(parser, outputTypes, /*allowEmpty=*/true))
+    return failure();
+
+  if (weights.size() != weightTypes.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of weights and weight types must match");
+  if (weightArgs.size() != weights.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of weight bindings and weight operands must match");
+  if (inputs.size() != inputTypes.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of inputs and input types must match");
+  if (inputArgs.size() != inputs.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of argument bindings and input operands must match");
+  if (hasCoreId && result.attributes.get(onnx_mlir::kCoreIdAttrName))
+    return parser.emitError(parser.getCurrentLocation(),
+                            "coreId cannot be specified both positionally and in attr-dict");
+
+  auto& builder = parser.getBuilder();
+  result.addAttribute(
+    "operandSegmentSizes",
+    builder.getDenseI32ArrayAttr({static_cast<int32_t>(weights.size()), static_cast<int32_t>(inputs.size())}));
+  if (hasCoreId)
+    result.addAttribute(onnx_mlir::kCoreIdAttrName, getI32Attr(parser, coreId));
+
+  if (parser.resolveOperands(weights, weightTypes, parser.getCurrentLocation(), result.operands)
+      || parser.resolveOperands(inputs, inputTypes, parser.getCurrentLocation(), result.operands))
+    return failure();
+  result.addTypes(outputTypes);
+
+  Region* body = result.addRegion();
+  applyArgumentTypes(weightTypes, weightArgs);
+  applyArgumentTypes(inputTypes, inputArgs);
+  llvm::append_range(regionArgs, weightArgs);
+  llvm::append_range(regionArgs, inputArgs);
+  return parser.parseRegion(*body, regionArgs);
+}
+
+template <typename ComputeBatchOpTy>
+void printComputeBatchLikeOp(ComputeBatchOpTy op, OpAsmPrinter& printer) {
+  auto laneArg = op.getLaneArgument();
+  SmallVector<Value> weightArgs;
+  weightArgs.reserve(op.getWeights().size());
+  for (unsigned index = 0; index < op.getWeights().size(); ++index) {
+    auto weightArg = op.getWeightArgument(index);
+    if (!weightArg)
+      return printer.printGenericOp(op.getOperation(), /*printOpName=*/false);
+    weightArgs.push_back(*weightArg);
+  }
+  SmallVector<Value> inputArgs;
+  inputArgs.reserve(op.getInputs().size());
+  for (unsigned index = 0; index < op.getInputs().size(); ++index) {
+    auto inputArg = op.getInputArgument(index);
+    if (!inputArg)
+      return printer.printGenericOp(op.getOperation(), /*printOpName=*/false);
+    inputArgs.push_back(*inputArg);
+  }
+
+  SmallVector<BlockArgument> outputArgs;
+  if (!laneArg)
+    return printer.printGenericOp(op.getOperation(), /*printOpName=*/false);
+  if (op.getNumResults() != 0) {
+    outputArgs.reserve(op.getNumResults());
+    for (unsigned index = 0; index < op.getNumResults(); ++index) {
+      auto outputArg = op.getOutputArgument(index);
+      if (!outputArg)
+        return printer.printGenericOp(op.getOperation(), /*printOpName=*/false);
+      outputArgs.push_back(*outputArg);
+    }
+  }
+
+  printer << " ";
+  printer.printOperand(*laneArg);
+  printer << " = 0 to " << op.getLaneCount();
+  printer << " ";
+  printBoundValueList(printer, weightArgs, op.getWeights(), ListDelimiter::Square);
+  printer << " ";
+  printBoundValueList(printer, inputArgs, op.getInputs(), ListDelimiter::Paren);
+  if (op.getNumResults() != 0) {
+    printer << " shared_outs";
+    printBlockArgumentList(printer, outputArgs);
+  }
+  printer << " crossbarWeights " << getComputeInstanceCrossbarUsage({op.getOperation(), 0, op.getLaneCount()}).size();
+  if (auto coreIdsAttr = op->template getAttrOfType<DenseI32ArrayAttr>(onnx_mlir::kCoreIdsAttrName)) {
+    printer << " coreIds ";
+    printCompressedIntegerList(printer, coreIdsAttr.asArrayRef());
+  }
+  printer.printOptionalAttrDict(
+    op->getAttrs(),
+    {op.getLaneCountAttrName().getValue(), op.getOperandSegmentSizesAttrName().getValue(), onnx_mlir::kCoreIdsAttrName});
+  printer << " : ";
+  printCompressedTypeList(printer, TypeRange(op.getWeights()), ListDelimiter::Square);
+  printer << " ";
+  printCompressedTypeList(printer, TypeRange(op.getInputs()), ListDelimiter::Paren);
+  printer << " -> ";
+  printCompressedTypeSequence(printer, op.getResultTypes());
+  printer << " ";
+  printer.printRegion(op.getBody(), /*printEntryBlockArgs=*/!op.getBody().hasOneBlock());
+}
+
+template <typename ComputeBatchOpTy>
+ParseResult parseComputeBatchLikeOp(OpAsmParser& parser, OperationState& result) {
+  int64_t lowerBound = 0;
+  int32_t laneCount = 0;
+  OpAsmParser::Argument laneArg;
+  SmallVector<OpAsmParser::Argument> weightArgs;
+  SmallVector<OpAsmParser::Argument> inputArgs;
+  SmallVector<OpAsmParser::Argument> outputArgs;
+  SmallVector<OpAsmParser::Argument> regionArgs;
+  SmallVector<OpAsmParser::UnresolvedOperand> weights;
+  SmallVector<OpAsmParser::UnresolvedOperand> inputs;
+  SmallVector<Type> weightTypes;
+  SmallVector<Type> inputTypes;
+  SmallVector<Type> outputTypes;
+  int32_t crossbarWeightCount = 0;
+  SmallVector<int32_t> coreIds;
+
+  if (parser.parseArgument(laneArg) || parser.parseEqual() || parser.parseInteger(lowerBound)
+      || parser.parseKeyword("to") || parser.parseInteger(laneCount))
+    return failure();
+  if (lowerBound != 0)
+    return parser.emitError(parser.getCurrentLocation(), "compute_batch currently requires a zero lower bound");
+  if (parseBoundValueList(parser, ListDelimiter::Square, weightArgs, weights))
+    return failure();
+  if (parseBoundValueList(parser, ListDelimiter::Paren, inputArgs, inputs))
+    return failure();
+  if (succeeded(parser.parseOptionalKeyword("shared_outs")))
+    if (parseBlockArgumentList(parser, outputArgs))
+      return failure();
+  bool hasCoreIds = parseOptionalKeywordAlias(parser, "coreIds", "core_ids");
+  if (hasCoreIds && parseCompressedIntegerList(parser, coreIds))
+    return failure();
+  bool hasCrossbarWeightCount = parseOptionalKeywordAlias(parser, "crossbarWeights", "crossbar_weights");
+  if (hasCrossbarWeightCount && parser.parseInteger(crossbarWeightCount))
+    return failure();
+  (void) crossbarWeightCount;
+
+  if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()
+      || parseCompressedOrTupleTypeList(parser, ListDelimiter::Square, weightTypes)
+      || parseCompressedRepeatedList(
+        parser, ListDelimiter::Paren, inputTypes, [&](Type& type) { return parser.parseType(type); })
+      || parser.parseArrow() || parseCompressedTypeSequence(parser, outputTypes, /*allowEmpty=*/true))
+    return failure();
+
+  if (weights.size() != weightTypes.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of weights and weight types must match");
+  if (weightArgs.size() != weights.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of weight bindings and weight operands must match");
+  if (inputs.size() != inputTypes.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of inputs and input types must match");
+  if (inputArgs.size() != inputs.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of argument bindings and input operands must match");
+  if (outputArgs.size() != outputTypes.size())
+    return parser.emitError(parser.getCurrentLocation(),
+                            "number of shared output bindings and result types must match");
+  if (hasCoreIds && result.attributes.get(onnx_mlir::kCoreIdsAttrName))
+    return parser.emitError(parser.getCurrentLocation(),
+                            "coreIds cannot be specified both positionally and in attr-dict");
+
+  auto& builder = parser.getBuilder();
+  result.addAttribute("laneCount", builder.getI32IntegerAttr(laneCount));
+  result.addAttribute(
+    "operandSegmentSizes",
+    builder.getDenseI32ArrayAttr({static_cast<int32_t>(weights.size()), static_cast<int32_t>(inputs.size())}));
+  if (hasCoreIds)
+    result.addAttribute(onnx_mlir::kCoreIdsAttrName, getDenseI32ArrayAttr(parser, coreIds));
+
+  if (parser.resolveOperands(weights, weightTypes, parser.getCurrentLocation(), result.operands)
+      || parser.resolveOperands(inputs, inputTypes, parser.getCurrentLocation(), result.operands))
+    return failure();
+  result.addTypes(outputTypes);
+
+  Region* body = result.addRegion();
+  applyBatchRegionArgumentTypes(
+    inputTypes, weightTypes, outputTypes, laneArg, weightArgs, inputArgs, outputArgs, regionArgs, parser.getBuilder());
+  return parser.parseRegion(*body, regionArgs);
+}
+
 } // namespace
 
 void SpatYieldOp::print(OpAsmPrinter& printer) {
@@ -151,6 +407,89 @@ ParseResult SpatYieldOp::parse(OpAsmParser& parser, OperationState& result) {
   return parser.resolveOperands(outputs, outputTypes, parser.getCurrentLocation(), result.operands);
 }
 
+void SpatBlockYieldOp::print(OpAsmPrinter& printer) {
+  printer << " ";
+  printCompressedValueSequence(printer, getOutputs());
+  if (getOperation()->getNumSuccessors() != 0) {
+    printer << " next ";
+    printer.printSuccessor(getOperation()->getSuccessor(0));
+  }
+  printer.printOptionalAttrDict((*this)->getAttrs());
+  printer << " : ";
+  printCompressedTypeSequence(printer, getOutputs().getTypes());
+}
+
+ParseResult SpatBlockYieldOp::parse(OpAsmParser& parser, OperationState& result) {
+  SmallVector<OpAsmParser::UnresolvedOperand> outputs;
+  SmallVector<Type> outputTypes;
+  Block* successor = nullptr;
+
+  OpAsmParser::UnresolvedOperand firstOutput;
+  OptionalParseResult firstOutputResult = parser.parseOptionalOperand(firstOutput);
+  if (firstOutputResult.has_value()) {
+    if (failed(*firstOutputResult))
+      return failure();
+    if (parseCompressedOperandEntryWithFirst(parser, firstOutput, outputs))
+      return failure();
+    while (succeeded(parser.parseOptionalComma()))
+      if (parseOneCompressedOperandEntry(parser, outputs))
+        return failure();
+  }
+
+  if (succeeded(parser.parseOptionalKeyword("next")) && parser.parseSuccessor(successor))
+    return failure();
+
+  if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()
+      || parseCompressedTypeSequence(parser, outputTypes, /*allowEmpty=*/true))
+    return failure();
+
+  if (outputs.size() != outputTypes.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of outputs and output types must match");
+  if (parser.resolveOperands(outputs, outputTypes, parser.getCurrentLocation(), result.operands))
+    return failure();
+  if (successor)
+    result.addSuccessors(successor);
+  return success();
+}
+
+void SpatDeferredCommunicationOp::print(OpAsmPrinter& printer) {
+  printer << " ";
+  printCompressedValueSequence(printer, getSources());
+  printer.printOptionalAttrDict((*this)->getAttrs());
+  printer << " : ";
+  printer.printFunctionalType(getSources().getTypes(), getOperation()->getResultTypes());
+  printer << " ";
+  printer.printRegion(getBody(), /*printEntryBlockArgs=*/false);
+}
+
+ParseResult SpatDeferredCommunicationOp::parse(OpAsmParser& parser, OperationState& result) {
+  SmallVector<OpAsmParser::UnresolvedOperand> sources;
+  Type functionTypeStorage;
+
+  if (parseCompressedOperandSequence(parser, sources) || parser.parseOptionalAttrDict(result.attributes)
+      || parser.parseColon() || parser.parseType(functionTypeStorage))
+    return failure();
+
+  auto functionType = dyn_cast<FunctionType>(functionTypeStorage);
+  if (!functionType)
+    return parser.emitError(parser.getCurrentLocation(), "expected deferred communication function type");
+  if (sources.size() != functionType.getNumInputs())
+    return parser.emitError(parser.getCurrentLocation(), "number of sources and source types must match");
+
+  if (parser.resolveOperands(sources, functionType.getInputs(), parser.getCurrentLocation(), result.operands))
+    return failure();
+  result.addTypes(functionType.getResults());
+
+  Region* body = result.addRegion();
+  SmallVector<OpAsmParser::Argument> bodyArgs;
+  for (Type type : functionType.getInputs()) {
+    OpAsmParser::Argument argument;
+    argument.type = type;
+    bodyArgs.push_back(argument);
+  }
+  return parser.parseRegion(*body, bodyArgs);
+}
+
 void SpatExtractRowsOp::print(OpAsmPrinter& printer) {
   printer << " ";
   printer.printOperand(getInput());
@@ -218,260 +557,146 @@ ParseResult SpatConcatOp::parse(OpAsmParser& parser, OperationState& result) {
   return success();
 }
 
-void SpatCompute::print(OpAsmPrinter& printer) {
-  SmallVector<Value> weightArgs;
-  weightArgs.reserve(getWeights().size());
-  for (unsigned index = 0; index < getWeights().size(); ++index) {
-    auto weightArg = getWeightArgument(index);
-    if (!weightArg)
-      return printer.printGenericOp(getOperation(), /*printOpName=*/false);
-    weightArgs.push_back(*weightArg);
-  }
-  SmallVector<Value> inputArgs;
-  inputArgs.reserve(getInputs().size());
-  for (unsigned index = 0; index < getInputs().size(); ++index) {
-    auto inputArg = getInputArgument(index);
-    if (!inputArg)
-      return printer.printGenericOp(getOperation(), /*printOpName=*/false);
-    inputArgs.push_back(*inputArg);
-  }
+void SpatBlueprintOp::print(OpAsmPrinter& printer) {
+  SmallVector<Value> operands {getInput()};
+  llvm::append_range(operands, getFragments());
 
-  printer << " ";
-  printBoundValueList(printer, weightArgs, getWeights(), ListDelimiter::Square);
-  printer << " ";
-  printBoundValueList(printer, inputArgs, getInputs(), ListDelimiter::Paren);
-
-  if (auto coreIdAttr = (*this)->getAttrOfType<IntegerAttr>(onnx_mlir::kCoreIdAttrName))
-    printer << " coreId " << coreIdAttr.getInt();
-  printer << " crossbarWeights " << collectDistinctCrossbarWeights(getOperation()).size();
+  printer << " fragments";
+  printCompressedValueList(printer, operands, ListDelimiter::Paren);
+  printer << " layout " << getLogicalLayout();
+  printer << " physical " << getPhysicalLayout();
+  printer << " offsets ";
+  printCompressedIntegerList(printer, getFragmentOffsets());
+  printer << " sizes ";
+  printCompressedIntegerList(printer, getFragmentSizes());
+  printer << " map " << getIndexMap();
+  if (std::optional<StringRef> mode = getMode())
+    printer << " mode " << *mode;
+  if (std::optional<ArrayRef<int64_t>> operandIndices = getFragmentOperandIndices()) {
+    printer << " operandIndices ";
+    printCompressedIntegerList(printer, *operandIndices);
+  }
+  if (std::optional<ArrayRef<int64_t>> sourceOffsets = getFragmentSourceOffsets()) {
+    printer << " sourceOffsets ";
+    printCompressedIntegerList(printer, *sourceOffsets);
+  }
+  if (std::optional<ArrayRef<int64_t>> strides = getFragmentStrides()) {
+    printer << " strides ";
+    printCompressedIntegerList(printer, *strides);
+  }
+  if (std::optional<StringRef> conflictPolicy = getConflictPolicy())
+    printer << " conflict " << *conflictPolicy;
+  if (std::optional<StringRef> coveragePolicy = getCoveragePolicy())
+    printer << " coverage " << *coveragePolicy;
 
   printer.printOptionalAttrDict((*this)->getAttrs(),
-                                {getOperandSegmentSizesAttrName().getValue(), onnx_mlir::kCoreIdAttrName});
-
+                                {getLogicalLayoutAttrName().getValue(),
+                                 getPhysicalLayoutAttrName().getValue(),
+                                 getFragmentOffsetsAttrName().getValue(),
+                                 getFragmentSizesAttrName().getValue(),
+                                 getIndexMapAttrName().getValue(),
+                                 getModeAttrName().getValue(),
+                                 getFragmentOperandIndicesAttrName().getValue(),
+                                 getFragmentSourceOffsetsAttrName().getValue(),
+                                 getFragmentStridesAttrName().getValue(),
+                                 getConflictPolicyAttrName().getValue(),
+                                 getCoveragePolicyAttrName().getValue()});
   printer << " : ";
-  printCompressedTypeList(printer, TypeRange(getWeights()), ListDelimiter::Square);
-  printer << " ";
-  printCompressedTypeList(printer, TypeRange(getInputs()), ListDelimiter::Paren);
+  printCompressedTypeList(printer, TypeRange(operands), ListDelimiter::Paren);
   printer << " -> ";
-  printCompressedTypeSequence(printer, getResultTypes());
-  printer << " ";
-  printer.printRegion(getBody(), /*printEntryBlockArgs=*/false);
+  printer.printType(getOutput().getType());
 }
 
-ParseResult SpatCompute::parse(OpAsmParser& parser, OperationState& result) {
-  SmallVector<OpAsmParser::Argument> weightArgs;
-  SmallVector<OpAsmParser::Argument> regionArgs;
-  SmallVector<OpAsmParser::UnresolvedOperand> weights;
-  SmallVector<OpAsmParser::UnresolvedOperand> inputs;
-  SmallVector<Type> weightTypes;
-  SmallVector<Type> inputTypes;
-  SmallVector<Type> outputTypes;
-  int32_t crossbarWeightCount = 0;
-  int32_t coreId = 0;
+ParseResult SpatBlueprintOp::parse(OpAsmParser& parser, OperationState& result) {
+  SmallVector<OpAsmParser::UnresolvedOperand> operands;
+  SmallVector<Type> operandTypes;
+  Type outputType;
+  StringAttr logicalLayout;
+  StringAttr physicalLayout;
+  StringAttr indexMap;
+  StringAttr mode;
+  StringAttr conflictPolicy;
+  StringAttr coveragePolicy;
+  SmallVector<int64_t> fragmentOffsets;
+  SmallVector<int64_t> fragmentSizes;
+  SmallVector<int64_t> fragmentOperandIndices;
+  SmallVector<int64_t> fragmentSourceOffsets;
+  SmallVector<int64_t> fragmentStrides;
 
-  if (parseBoundValueList(parser, ListDelimiter::Square, weightArgs, weights))
+  if (parser.parseKeyword("fragments")
+      || parseCompressedOperandList(parser, ListDelimiter::Paren, operands)
+      || parser.parseKeyword("layout") || parseBareStringAttr(parser, logicalLayout)
+      || parser.parseKeyword("physical") || parseBareStringAttr(parser, physicalLayout)
+      || parser.parseKeyword("offsets") || parseCompressedIntegerList(parser, fragmentOffsets)
+      || parser.parseKeyword("sizes") || parseCompressedIntegerList(parser, fragmentSizes)
+      || parser.parseKeyword("map") || parseBareStringAttr(parser, indexMap))
     return failure();
 
-  SmallVector<OpAsmParser::Argument> inputArgs;
-  if (parseBoundValueList(parser, ListDelimiter::Paren, inputArgs, inputs))
+  if (succeeded(parser.parseOptionalKeyword("mode")) && parseBareStringAttr(parser, mode))
     return failure();
-
-  bool hasCoreId = parseOptionalKeywordAlias(parser, "coreId", "core_id");
-  if (hasCoreId && parser.parseInteger(coreId))
+  if (succeeded(parser.parseOptionalKeyword("operandIndices"))
+      && parseCompressedIntegerList(parser, fragmentOperandIndices))
     return failure();
-
-  bool hasCrossbarWeightCount = parseOptionalKeywordAlias(parser, "crossbarWeights", "crossbar_weights");
-  if (hasCrossbarWeightCount && parser.parseInteger(crossbarWeightCount))
+  if (succeeded(parser.parseOptionalKeyword("sourceOffsets"))
+      && parseCompressedIntegerList(parser, fragmentSourceOffsets))
+    return failure();
+  if (succeeded(parser.parseOptionalKeyword("strides")) && parseCompressedIntegerList(parser, fragmentStrides))
+    return failure();
+  if (succeeded(parser.parseOptionalKeyword("conflict")) && parseBareStringAttr(parser, conflictPolicy))
+    return failure();
+  if (succeeded(parser.parseOptionalKeyword("coverage")) && parseBareStringAttr(parser, coveragePolicy))
     return failure();
-  (void) crossbarWeightCount;
 
   if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()
       || parseCompressedRepeatedList(
-        parser, ListDelimiter::Square, weightTypes, [&](Type& type) { return parser.parseType(type); })
-      || parseCompressedRepeatedList(
-        parser, ListDelimiter::Paren, inputTypes, [&](Type& type) { return parser.parseType(type); })
-      || parser.parseArrow() || parseCompressedTypeSequence(parser, outputTypes, /*allowEmpty=*/true))
+        parser, ListDelimiter::Paren, operandTypes, [&](Type& type) { return parser.parseType(type); })
+      || parser.parseArrow() || parser.parseType(outputType))
     return failure();
-
-  if (weights.size() != weightTypes.size())
-    return parser.emitError(parser.getCurrentLocation(), "number of weights and weight types must match");
-  if (weightArgs.size() != weights.size())
-    return parser.emitError(parser.getCurrentLocation(), "number of weight bindings and weight operands must match");
-  if (inputs.size() != inputTypes.size())
-    return parser.emitError(parser.getCurrentLocation(), "number of inputs and input types must match");
-  if (inputArgs.size() != inputs.size())
-    return parser.emitError(parser.getCurrentLocation(), "number of argument bindings and input operands must match");
-  if (hasCoreId && result.attributes.get(onnx_mlir::kCoreIdAttrName))
-    return parser.emitError(parser.getCurrentLocation(),
-                            "coreId cannot be specified both positionally and in attr-dict");
+  if (operands.empty())
+    return parser.emitError(parser.getCurrentLocation(), "spat.blueprint requires at least one fragment operand");
+  if (operands.size() != operandTypes.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of fragment operands and types must match");
 
   auto& builder = parser.getBuilder();
-  result.addAttribute(
-    "operandSegmentSizes",
-    builder.getDenseI32ArrayAttr({static_cast<int32_t>(weights.size()), static_cast<int32_t>(inputs.size())}));
-  if (hasCoreId)
-    result.addAttribute(onnx_mlir::kCoreIdAttrName, getI32Attr(parser, coreId));
+  result.addAttribute("logicalLayout", logicalLayout);
+  result.addAttribute("physicalLayout", physicalLayout);
+  result.addAttribute("fragmentOffsets", builder.getDenseI64ArrayAttr(fragmentOffsets));
+  result.addAttribute("fragmentSizes", builder.getDenseI64ArrayAttr(fragmentSizes));
+  result.addAttribute("indexMap", indexMap);
+  if (mode)
+    result.addAttribute("mode", mode);
+  if (!fragmentOperandIndices.empty())
+    result.addAttribute("fragmentOperandIndices", builder.getDenseI64ArrayAttr(fragmentOperandIndices));
+  if (!fragmentSourceOffsets.empty())
+    result.addAttribute("fragmentSourceOffsets", builder.getDenseI64ArrayAttr(fragmentSourceOffsets));
+  if (!fragmentStrides.empty())
+    result.addAttribute("fragmentStrides", builder.getDenseI64ArrayAttr(fragmentStrides));
+  if (conflictPolicy)
+    result.addAttribute("conflictPolicy", conflictPolicy);
+  if (coveragePolicy)
+    result.addAttribute("coveragePolicy", coveragePolicy);
 
-  if (parser.resolveOperands(weights, weightTypes, parser.getCurrentLocation(), result.operands)
-      || parser.resolveOperands(inputs, inputTypes, parser.getCurrentLocation(), result.operands))
+  if (parser.resolveOperands(operands, operandTypes, parser.getCurrentLocation(), result.operands))
     return failure();
-  result.addTypes(outputTypes);
-
-  Region* body = result.addRegion();
-  applyArgumentTypes(weightTypes, weightArgs);
-  applyArgumentTypes(inputTypes, inputArgs);
-  llvm::append_range(regionArgs, weightArgs);
-  llvm::append_range(regionArgs, inputArgs);
-  return parser.parseRegion(*body, regionArgs);
+  result.addTypes(outputType);
+  return success();
 }
 
-void SpatComputeBatch::print(OpAsmPrinter& printer) {
-  auto laneArg = getLaneArgument();
-  SmallVector<Value> weightArgs;
-  weightArgs.reserve(getWeights().size());
-  for (unsigned index = 0; index < getWeights().size(); ++index) {
-    auto weightArg = getWeightArgument(index);
-    if (!weightArg)
-      return printer.printGenericOp(getOperation(), /*printOpName=*/false);
-    weightArgs.push_back(*weightArg);
+void SpatGraphCompute::print(OpAsmPrinter& printer) { printComputeLikeOp(*this, printer); }
+ParseResult SpatGraphCompute::parse(OpAsmParser& parser, OperationState& result) {
+  return parseComputeLikeOp<SpatGraphCompute>(parser, result);
 }
-  SmallVector<Value> inputArgs;
-  inputArgs.reserve(getInputs().size());
-  for (unsigned index = 0; index < getInputs().size(); ++index) {
-    auto inputArg = getInputArgument(index);
-    if (!inputArg)
-      return printer.printGenericOp(getOperation(), /*printOpName=*/false);
-    inputArgs.push_back(*inputArg);
+void SpatScheduledCompute::print(OpAsmPrinter& printer) { printComputeLikeOp(*this, printer); }
+ParseResult SpatScheduledCompute::parse(OpAsmParser& parser, OperationState& result) {
+  return parseComputeLikeOp<SpatScheduledCompute>(parser, result);
 }
-
-  SmallVector<BlockArgument> outputArgs;
-  if (!laneArg)
-    return printer.printGenericOp(getOperation(), /*printOpName=*/false);
-  if (getNumResults() != 0) {
-    outputArgs.reserve(getNumResults());
-    for (unsigned index = 0; index < getNumResults(); ++index) {
-      auto outputArg = getOutputArgument(index);
-      if (!outputArg)
-        return printer.printGenericOp(getOperation(), /*printOpName=*/false);
-      outputArgs.push_back(*outputArg);
+void SpatGraphComputeBatch::print(OpAsmPrinter& printer) { printComputeBatchLikeOp(*this, printer); }
+ParseResult SpatGraphComputeBatch::parse(OpAsmParser& parser, OperationState& result) {
+  return parseComputeBatchLikeOp<SpatGraphComputeBatch>(parser, result);
 }
-  }
-
-  printer << " ";
-  printer.printOperand(*laneArg);
-  printer << " = 0 to " << getLaneCount();
-
-  printer << " ";
-  printBoundValueList(printer, weightArgs, getWeights(), ListDelimiter::Square);
-  printer << " ";
-  printBoundValueList(printer, inputArgs, getInputs(), ListDelimiter::Paren);
-
-  if (getNumResults() != 0) {
-    printer << " shared_outs";
-    printBlockArgumentList(printer, outputArgs);
-  }
-
-  printer << " crossbarWeights " << getComputeInstanceCrossbarUsage({getOperation(), 0, getLaneCount()}).size();
-
-  if (auto coreIdsAttr = (*this)->getAttrOfType<DenseI32ArrayAttr>(onnx_mlir::kCoreIdsAttrName)) {
-    printer << " coreIds ";
-    printCompressedIntegerList(printer, coreIdsAttr.asArrayRef());
-  }
-
-  printer.printOptionalAttrDict(
-    (*this)->getAttrs(),
-    {getLaneCountAttrName().getValue(), getOperandSegmentSizesAttrName().getValue(), onnx_mlir::kCoreIdsAttrName});
-
-  printer << " : ";
-  printCompressedTypeList(printer, TypeRange(getWeights()), ListDelimiter::Square);
-  printer << " ";
-  printCompressedTypeList(printer, TypeRange(getInputs()), ListDelimiter::Paren);
-  printer << " -> ";
-  printCompressedTypeSequence(printer, getResultTypes());
-  printer << " ";
-  printer.printRegion(getBody(), /*printEntryBlockArgs=*/false);
-}
-
-ParseResult SpatComputeBatch::parse(OpAsmParser& parser, OperationState& result) {
-  int64_t lowerBound = 0;
-  int32_t laneCount = 0;
-  OpAsmParser::Argument laneArg;
-  SmallVector<OpAsmParser::Argument> weightArgs;
-  SmallVector<OpAsmParser::Argument> inputArgs;
-  SmallVector<OpAsmParser::Argument> outputArgs;
-  SmallVector<OpAsmParser::Argument> regionArgs;
-  SmallVector<OpAsmParser::UnresolvedOperand> weights;
-  SmallVector<OpAsmParser::UnresolvedOperand> inputs;
-  SmallVector<Type> weightTypes;
-  SmallVector<Type> inputTypes;
-  SmallVector<Type> outputTypes;
-  int32_t crossbarWeightCount = 0;
-  SmallVector<int32_t> coreIds;
-
-  if (parser.parseArgument(laneArg) || parser.parseEqual() || parser.parseInteger(lowerBound)
-      || parser.parseKeyword("to") || parser.parseInteger(laneCount))
-    return failure();
-  if (lowerBound != 0)
-    return parser.emitError(parser.getCurrentLocation(), "compute_batch currently requires a zero lower bound");
-
-  if (parseBoundValueList(parser, ListDelimiter::Square, weightArgs, weights))
-    return failure();
-
-  if (parseBoundValueList(parser, ListDelimiter::Paren, inputArgs, inputs))
-    return failure();
-
-  if (succeeded(parser.parseOptionalKeyword("shared_outs")))
-    if (parseBlockArgumentList(parser, outputArgs))
-      return failure();
-
-  bool hasCoreIds = parseOptionalKeywordAlias(parser, "coreIds", "core_ids");
-  if (hasCoreIds && parseCompressedIntegerList(parser, coreIds))
-    return failure();
-
-  bool hasCrossbarWeightCount = parseOptionalKeywordAlias(parser, "crossbarWeights", "crossbar_weights");
-  if (hasCrossbarWeightCount && parser.parseInteger(crossbarWeightCount))
-    return failure();
-  (void) crossbarWeightCount;
-
-  if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()
-      || parseCompressedOrTupleTypeList(parser, ListDelimiter::Square, weightTypes)
-      || parseCompressedRepeatedList(
-        parser, ListDelimiter::Paren, inputTypes, [&](Type& type) { return parser.parseType(type); })
-      || parser.parseArrow() || parseCompressedTypeSequence(parser, outputTypes, /*allowEmpty=*/true))
-    return failure();
-
-  if (weights.size() != weightTypes.size())
-    return parser.emitError(parser.getCurrentLocation(), "number of weights and weight types must match");
-  if (weightArgs.size() != weights.size())
-    return parser.emitError(parser.getCurrentLocation(), "number of weight bindings and weight operands must match");
-  if (inputs.size() != inputTypes.size())
-    return parser.emitError(parser.getCurrentLocation(), "number of inputs and input types must match");
-  if (inputArgs.size() != inputs.size())
-    return parser.emitError(parser.getCurrentLocation(), "number of argument bindings and input operands must match");
-  if (outputArgs.size() != outputTypes.size())
-    return parser.emitError(parser.getCurrentLocation(),
-                            "number of shared output bindings and result types must match");
-  if (hasCoreIds && result.attributes.get(onnx_mlir::kCoreIdsAttrName))
-    return parser.emitError(parser.getCurrentLocation(),
-                            "coreIds cannot be specified both positionally and in attr-dict");
-
-  auto& builder = parser.getBuilder();
-  result.addAttribute("laneCount", builder.getI32IntegerAttr(laneCount));
-  result.addAttribute(
-    "operandSegmentSizes",
-    builder.getDenseI32ArrayAttr({static_cast<int32_t>(weights.size()), static_cast<int32_t>(inputs.size())}));
-  if (hasCoreIds)
-    result.addAttribute(onnx_mlir::kCoreIdsAttrName, getDenseI32ArrayAttr(parser, coreIds));
-
-  if (parser.resolveOperands(weights, weightTypes, parser.getCurrentLocation(), result.operands)
-      || parser.resolveOperands(inputs, inputTypes, parser.getCurrentLocation(), result.operands))
-    return failure();
-  result.addTypes(outputTypes);
-
-  Region* body = result.addRegion();
-  applyBatchRegionArgumentTypes(
-    inputTypes, weightTypes, outputTypes, laneArg, weightArgs, inputArgs, outputArgs, regionArgs, parser.getBuilder());
-  return parser.parseRegion(*body, regionArgs);
+void SpatScheduledComputeBatch::print(OpAsmPrinter& printer) { printComputeBatchLikeOp(*this, printer); }
+ParseResult SpatScheduledComputeBatch::parse(OpAsmParser& parser, OperationState& result) {
+  return parseComputeBatchLikeOp<SpatScheduledComputeBatch>(parser, result);
 }
 
 void SpatInParallelOp::print(OpAsmPrinter& printer) {

src/PIM/Dialect/Spatial/SpatialOpsCanonicalization.cpp

@@ -1,8 +1,13 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Block.h"
+#include "mlir/IR/IRMapping.h"
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/LogicalResult.h"
 
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
@@ -10,8 +15,9 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace spatial {
 
-LogicalResult SpatCompute::fold(FoldAdaptor adaptor, ::llvm::SmallVectorImpl<::mlir::OpFoldResult>& results) {
-  Block& block = getBody().front();
+template <typename ComputeOpTy>
+LogicalResult foldComputeLike(ComputeOpTy compute, ::llvm::SmallVectorImpl<::mlir::OpFoldResult>& results) {
+  Block& block = compute.getBody().front();
   if (!llvm::hasSingleElement(block))
     return failure();
 
@@ -22,7 +28,7 @@ LogicalResult SpatCompute::fold(FoldAdaptor adaptor, ::llvm::SmallVectorImpl<::m
   for (Value yieldedValue : yieldOp.getOperands()) {
     if (auto blockArg = dyn_cast<BlockArgument>(yieldedValue)) {
       if (blockArg.getOwner() == &block) {
-        results.push_back(getOperand(blockArg.getArgNumber()));
+        results.push_back(compute.getOperand(blockArg.getArgNumber()));
         continue;
       }
     }
@@ -31,5 +37,185 @@ LogicalResult SpatCompute::fold(FoldAdaptor adaptor, ::llvm::SmallVectorImpl<::m
   return success();
 }
 
+LogicalResult SpatGraphCompute::fold(FoldAdaptor adaptor, ::llvm::SmallVectorImpl<::mlir::OpFoldResult>& results) {
+  return foldComputeLike(*this, results);
+}
+
+LogicalResult SpatScheduledCompute::fold(FoldAdaptor adaptor, ::llvm::SmallVectorImpl<::mlir::OpFoldResult>& results) {
+  return foldComputeLike(*this, results);
+}
+
+template <typename ScalarComputeOpTy>
+static ScalarComputeOpTy createEmptyScalarCompute(PatternRewriter& rewriter,
+                                                  Location loc,
+                                                  TypeRange resultTypes,
+                                                  ValueRange weights,
+                                                  ValueRange inputs) {
+  auto computeOp = ScalarComputeOpTy::create(rewriter, loc, resultTypes, weights, inputs);
+  SmallVector<Type> blockArgTypes;
+  SmallVector<Location> blockArgLocs;
+  blockArgTypes.reserve(weights.size() + inputs.size());
+  blockArgLocs.reserve(weights.size() + inputs.size());
+  for (Value weight : weights) {
+    blockArgTypes.push_back(weight.getType());
+    blockArgLocs.push_back(weight.getLoc());
+  }
+  for (Value input : inputs) {
+    blockArgTypes.push_back(input.getType());
+    blockArgLocs.push_back(input.getLoc());
+  }
+  rewriter.createBlock(&computeOp.getBody(), computeOp.getBody().end(), blockArgTypes, blockArgLocs);
+  rewriter.setInsertionPointToStart(&computeOp.getBody().front());
+  return computeOp;
+}
+
+static SmallVector<OpFoldResult> remapMixedOffsets(ArrayRef<OpFoldResult> mixedOffsets, IRMapping& mapper) {
+  SmallVector<OpFoldResult> remapped;
+  remapped.reserve(mixedOffsets.size());
+  for (OpFoldResult ofr : mixedOffsets) {
+    if (auto value = dyn_cast<Value>(ofr))
+      remapped.push_back(cast<Value>(mapper.lookupOrDefault(value)));
+    else
+      remapped.push_back(cast<Attribute>(ofr));
+  }
+  return remapped;
+}
+
+static SmallVector<Value> createEmptyResults(PatternRewriter& rewriter, Location loc, TypeRange resultTypes) {
+  SmallVector<Value> resultValues;
+  resultValues.reserve(resultTypes.size());
+  for (Type resultType : resultTypes) {
+    auto tensorType = dyn_cast<RankedTensorType>(resultType);
+    if (!tensorType || !tensorType.hasStaticShape())
+      return {};
+    resultValues.push_back(tensor::EmptyOp::create(rewriter, loc, tensorType.getShape(), tensorType.getElementType()));
+  }
+  return resultValues;
+}
+
+template <typename ScalarComputeOpTy, typename ComputeBatchOpTy>
+static void copyCanonicalizedBatchAttrs(ScalarComputeOpTy compute, ComputeBatchOpTy batch, PatternRewriter& rewriter) {
+  for (NamedAttribute attr : batch->getAttrs()) {
+    if (attr.getName() == batch.getOperandSegmentSizesAttrName() || attr.getName() == batch.getLaneCountAttrName()
+        || attr.getName() == onnx_mlir::kCoreIdsAttrName)
+      continue;
+    compute->setAttr(attr.getName(), attr.getValue());
+  }
+  if constexpr (std::is_same_v<ComputeBatchOpTy, SpatScheduledComputeBatch>) {
+    if (auto coreIds = batch->template getAttrOfType<DenseI32ArrayAttr>(onnx_mlir::kCoreIdsAttrName)) {
+      assert(coreIds.size() == 1 && "single-lane scheduled compute_batch canonicalization expects exactly one core id");
+      compute->setAttr(onnx_mlir::kCoreIdAttrName, rewriter.getI32IntegerAttr(coreIds.asArrayRef().front()));
+    }
+  }
+}
+
+template <typename ComputeBatchOpTy, typename ScalarComputeOpTy>
+struct CanonicalizeSingleLaneComputeBatchPattern : OpRewritePattern<ComputeBatchOpTy> {
+  using OpRewritePattern<ComputeBatchOpTy>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ComputeBatchOpTy compute, PatternRewriter& rewriter) const override {
+    if (compute.getLaneCount() != 1)
+      return rewriter.notifyMatchFailure(compute, "lane count is not 1");
+
+    Block& oldBlock = compute.getBody().front();
+    auto oldLaneArg = compute.getLaneArgument();
+    if (!oldLaneArg)
+      return rewriter.notifyMatchFailure(compute, "missing compute_batch lane block argument");
+
+    rewriter.setInsertionPointAfter(compute);
+    auto newCompute =
+      createEmptyScalarCompute<ScalarComputeOpTy>(rewriter, compute.getLoc(), compute.getResultTypes(), compute.getWeights(), compute.getInputs());
+    copyCanonicalizedBatchAttrs(newCompute, compute, rewriter);
+    auto* newBlock = &newCompute.getBody().front();
+    rewriter.setInsertionPointToStart(newBlock);
+
+    IRMapping mapper;
+    Value zero = arith::ConstantIndexOp::create(rewriter, compute.getLoc(), 0);
+    mapper.map(*oldLaneArg, zero);
+    for (auto [index, weight] : llvm::enumerate(compute.getWeights())) {
+      auto oldArg = compute.getWeightArgument(index);
+      auto newArg = newCompute.getWeightArgument(index);
+      if (!oldArg || !newArg)
+        return rewriter.notifyMatchFailure(compute, "missing rewritten compute weight block argument");
+      mapper.map(*oldArg, *newArg);
+    }
+    for (auto [index, input] : llvm::enumerate(compute.getInputs())) {
+      auto oldArg = compute.getInputArgument(index);
+      auto newArg = newCompute.getInputArgument(index);
+      if (!oldArg || !newArg)
+        return rewriter.notifyMatchFailure(compute, "missing rewritten compute input block argument");
+      mapper.map(*oldArg, *newArg);
+    }
+
+    SmallVector<Value> resultValues = createEmptyResults(rewriter, compute.getLoc(), compute.getResultTypes());
+    if (resultValues.size() != compute.getNumResults())
+      return rewriter.notifyMatchFailure(compute, "single-lane compute_batch canonicalization requires static ranked results");
+    for (auto [index, resultValue] : llvm::enumerate(resultValues)) {
+      auto oldOutputArg = compute.getOutputArgument(index);
+      if (!oldOutputArg)
+        return rewriter.notifyMatchFailure(compute, "missing compute_batch output block argument");
+      mapper.map(*oldOutputArg, resultValue);
+    }
+
+    auto oldInParallel = dyn_cast<SpatInParallelOp>(oldBlock.getTerminator());
+    auto oldYield = dyn_cast<SpatYieldOp>(oldBlock.getTerminator());
+    for (Operation& op : oldBlock.without_terminator())
+      rewriter.clone(op, mapper);
+
+    if (oldYield) {
+      SpatYieldOp::create(rewriter, oldYield.getLoc(), ValueRange {});
+      rewriter.replaceOp(compute, newCompute.getResults());
+      return success();
+    }
+    if (!oldInParallel)
+      return rewriter.notifyMatchFailure(compute, "expected spat.in_parallel or empty spat.yield terminator");
+
+    DenseMap<BlockArgument, size_t> outputIndexByArg;
+    for (size_t index = 0; index < compute.getNumResults(); ++index) {
+      auto oldOutputArg = compute.getOutputArgument(index);
+      if (!oldOutputArg)
+        return rewriter.notifyMatchFailure(compute, "missing compute_batch output block argument");
+      outputIndexByArg[*oldOutputArg] = index;
+    }
+
+    for (Operation& op : oldInParallel.getRegion().front()) {
+      auto insertSlice = dyn_cast<tensor::ParallelInsertSliceOp>(&op);
+      if (!insertSlice)
+        return rewriter.notifyMatchFailure(compute, "expected only tensor.parallel_insert_slice in spat.in_parallel");
+      auto oldDest = dyn_cast<BlockArgument>(insertSlice.getDest());
+      if (!oldDest)
+        return rewriter.notifyMatchFailure(compute, "expected tensor.parallel_insert_slice destination to be a block argument");
+      auto resultIndexIt = outputIndexByArg.find(oldDest);
+      if (resultIndexIt == outputIndexByArg.end())
+        return rewriter.notifyMatchFailure(compute, "unexpected tensor.parallel_insert_slice destination");
+      size_t resultIndex = resultIndexIt->second;
+      Value remappedSource = mapper.lookupOrDefault(insertSlice.getSource());
+      auto remappedOffsets = remapMixedOffsets(insertSlice.getMixedOffsets(), mapper);
+      auto remappedSizes = remapMixedOffsets(insertSlice.getMixedSizes(), mapper);
+      auto remappedStrides = remapMixedOffsets(insertSlice.getMixedStrides(), mapper);
+      resultValues[resultIndex] = tensor::InsertSliceOp::create(rewriter,
+                                                                insertSlice.getLoc(),
+                                                                remappedSource,
+                                                                resultValues[resultIndex],
+                                                                remappedOffsets,
+                                                                remappedSizes,
+                                                                remappedStrides)
+                                    .getResult();
+    }
+
+    SpatYieldOp::create(rewriter, oldInParallel.getLoc(), resultValues);
+    rewriter.replaceOp(compute, newCompute.getResults());
+    return success();
+  }
+};
+
+void SpatGraphComputeBatch::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) {
+  results.add<CanonicalizeSingleLaneComputeBatchPattern<SpatGraphComputeBatch, SpatGraphCompute>>(context);
+}
+
+void SpatScheduledComputeBatch::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) {
+  results.add<CanonicalizeSingleLaneComputeBatchPattern<SpatScheduledComputeBatch, SpatScheduledCompute>>(context);
+}
+
 } // namespace spatial
 } // namespace onnx_mlir

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -1,5 +1,6 @@
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Block.h"
@@ -35,7 +36,8 @@ static FailureOr<ArrayRef<int64_t>> getWeightShapeForWeightedOp(Value weight) {
   return shapedType.getShape();
 }
 
-static bool isBatchOutputArgument(SpatComputeBatch batchOp, Value value) {
+template <typename ComputeBatchOpTy>
+static bool isBatchOutputArgument(ComputeBatchOpTy batchOp, Value value) {
   if (batchOp.getNumResults() == 0)
     return false;
   auto blockArg = dyn_cast<BlockArgument>(value);
@@ -58,8 +60,43 @@ static LogicalResult verifyStaticWeights(ComputeOpTy computeOp, StringRef kind)
   return success();
 }
 
+static bool isStaticScfForInductionVar(Value value) {
+  auto blockArg = dyn_cast<BlockArgument>(value);
+  if (!blockArg)
+    return false;
+
+  auto loop = dyn_cast_or_null<scf::ForOp>(blockArg.getOwner()->getParentOp());
+  if (!loop || loop.getInductionVar() != value)
+    return false;
+
+  std::optional<int64_t> lowerBound = matchConstantIndexValue(loop.getLowerBound());
+  std::optional<int64_t> upperBound = matchConstantIndexValue(loop.getUpperBound());
+  std::optional<int64_t> step = matchConstantIndexValue(loop.getStep());
+  return lowerBound && upperBound && step && *step > 0 && *upperBound >= *lowerBound;
+}
+
+static bool isStaticIndexExpr(Value value) {
+  if (matchConstantIndexValue(value))
+    return true;
+
+  auto affineApply = value.getDefiningOp<affine::AffineApplyOp>();
+  if (affineApply) {
+    if (!isSingleResultSymbolFreeAffineMap(affineApply.getAffineMap()))
+      return false;
+    return llvm::all_of(affineApply.getMapOperands(), isStaticIndexExpr);
+  }
+
+  if (auto addOp = value.getDefiningOp<arith::AddIOp>())
+    return isStaticIndexExpr(addOp.getLhs()) && isStaticIndexExpr(addOp.getRhs());
+
+  if (auto mulOp = value.getDefiningOp<arith::MulIOp>())
+    return isStaticIndexExpr(mulOp.getLhs()) && isStaticIndexExpr(mulOp.getRhs());
+
+  return false;
+}
+
 static bool isSupportedLaneOffsetExpr(Value value, BlockArgument laneArg) {
-  if (value == laneArg || matchConstantIndexValue(value))
+  if (value == laneArg || isStaticIndexExpr(value) || isStaticScfForInductionVar(value))
     return true;
 
   auto affineApply = value.getDefiningOp<affine::AffineApplyOp>();
@@ -83,10 +120,15 @@ static bool isSupportedLaneOffsetExpr(Value value, BlockArgument laneArg) {
   }
 
   auto addOp = value.getDefiningOp<arith::AddIOp>();
-  if (!addOp)
+  if (addOp)
+    return (isSupportedLaneOffsetExpr(addOp.getLhs(), laneArg) && isStaticIndexExpr(addOp.getRhs()))
+        || (isSupportedLaneOffsetExpr(addOp.getRhs(), laneArg) && isStaticIndexExpr(addOp.getLhs()));
+
+  auto mulOp = value.getDefiningOp<arith::MulIOp>();
+  if (!mulOp)
     return false;
-  return (addOp.getLhs() == laneArg && matchConstantIndexValue(addOp.getRhs()))
-      || (addOp.getRhs() == laneArg && matchConstantIndexValue(addOp.getLhs()));
+  return (isSupportedLaneOffsetExpr(mulOp.getLhs(), laneArg) && isStaticIndexExpr(mulOp.getRhs()))
+      || (isSupportedLaneOffsetExpr(mulOp.getRhs(), laneArg) && isStaticIndexExpr(mulOp.getLhs()));
 }
 
 static LogicalResult
@@ -150,25 +192,68 @@ static bool isConstantExternalValue(Value value) {
   return definingOp && definingOp->hasTrait<OpTrait::ConstantLike>();
 }
 
+static bool isRecordedDeferredCommunicationSource(Operation* op, Value value) {
+  auto transfer = dyn_cast<SpatDeferredCommunicationOp>(op);
+  return transfer && llvm::is_contained(transfer.getSources(), value);
+}
+
 static LogicalResult verifyOnlyConstantExternalValues(Operation* ownerOp, Region& region, StringRef kind) {
   bool hasFailure = false;
   region.walk([&](Operation* op) {
     for (OpOperand& operand : op->getOpOperands()) {
       Value value = operand.get();
-      if (isDefinedInsideRegion(value, region) || isConstantExternalValue(value))
+      if (isDefinedInsideRegion(value, region) || isConstantExternalValue(value)
+          || isRecordedDeferredCommunicationSource(op, value))
         continue;
 
-      InFlightDiagnostic diagnostic = ownerOp->emitOpError()
-                                   << kind << " body may only directly reference external constants";
+      InFlightDiagnostic diagnostic =
+        ownerOp->emitOpError() << kind << " body may not capture external values";
       diagnostic.attachNote(op->getLoc())
-        << "non-constant external operand #" << operand.getOperandNumber() << " is used by " << op->getName();
+        << "owner='" << ownerOp->getName() << "' nestedOp='" << op->getName() << "' operand#"
+        << operand.getOperandNumber() << " type=" << value.getType()
+        << " category=" << (isa<TensorType>(value.getType()) ? "tensor" : (value.getType().isIndex() ? "index"
+                                                                                                        : "scalar"));
+      if (Operation* definingOp = value.getDefiningOp())
+        diagnostic.attachNote(definingOp->getLoc()) << "defining op is '" << definingOp->getName() << "'";
+      else if (auto blockArg = dyn_cast<BlockArgument>(value))
+        diagnostic.attachNote(blockArg.getOwner()->getParentOp()->getLoc())
+          << "value is block argument #" << blockArg.getArgNumber() << " of '"
+          << blockArg.getOwner()->getParentOp()->getName() << "'";
       hasFailure = true;
     }
   });
   return success(!hasFailure);
 }
 
-static LogicalResult verifyBatchBody(SpatComputeBatch batchOp, Block& block) {
+static LogicalResult verifyYieldTypes(Operation* op, Region& region, TypeRange resultTypes, StringRef kind) {
+  if (region.empty())
+    return op->emitOpError() << kind << " requires a body block";
+  Block& block = region.front();
+  auto yield = dyn_cast_or_null<SpatYieldOp>(block.getTerminator());
+  if (!yield)
+    return op->emitOpError() << kind << " body must terminate with spat.yield";
+  if (yield.getOutputs().size() != resultTypes.size())
+    return op->emitOpError() << kind << " yield operand count must match result count";
+  for (auto [yieldType, resultType] : llvm::zip(yield.getOutputs().getTypes(), resultTypes))
+    if (yieldType != resultType)
+      return op->emitOpError() << kind << " yield operand types must match result types";
+  return success();
+}
+
+static LogicalResult verifyRegionArguments(Operation* op, Region& region, ValueRange operands, StringRef kind) {
+  if (region.empty())
+    return op->emitOpError() << kind << " requires a body block";
+  Block& block = region.front();
+  if (block.getNumArguments() != operands.size())
+    return op->emitOpError() << kind << " body argument count must match operand count";
+  for (auto [arg, operand] : llvm::zip(block.getArguments(), operands))
+    if (arg.getType() != operand.getType())
+      return op->emitOpError() << kind << " body argument types must match operand types";
+  return success();
+}
+
+template <typename ComputeBatchOpTy>
+static LogicalResult verifyBatchBody(ComputeBatchOpTy batchOp, Block& block, bool verifyLaneSliceOffsets = true) {
   if (batchOp.getNumResults() == 0) {
     auto yieldOp = dyn_cast_or_null<SpatYieldOp>(block.getTerminator());
     if (!yieldOp)
@@ -183,6 +268,7 @@ static LogicalResult verifyBatchBody(SpatComputeBatch batchOp, Block& block) {
   auto laneArg = batchOp.getLaneArgument();
   if (!laneArg)
     return batchOp.emitError("compute_batch body must have a lane block argument");
+  if (verifyLaneSliceOffsets)
     for (auto& bodyOp : block) {
       if (auto extractSlice = dyn_cast<tensor::ExtractSliceOp>(&bodyOp))
         if (failed(verifyStaticUnitStrideExtractSliceOp(extractSlice, *laneArg, "tensor.extract_slice")))
@@ -344,144 +430,479 @@ LogicalResult SpatConcatOp::verify() {
   return success();
 }
 
-LogicalResult verifyComputeResultsUses(Operation* op) {
-  if (!isa<SpatCompute, SpatComputeBatch>(op))
-    return op->emitError("verifyComputeResultUses: Op is not a SpatCompute/SpatComputeBatch operation");
-  if (!llvm::all_of(op->getResults(), [](Value result) {
-        return llvm::all_of(result.getUsers(), [](Operation* op) {
-          return !(op->getParentOfType<SpatCompute>() || op->getParentOfType<SpatComputeBatch>());
-        });
-      })) {
-    return op->emitError("ComputeResult used directly inside another Compute");
+static bool isKnownLogicalLayout(StringRef layout) { return layout == "nchw"; }
+
+static bool isKnownPhysicalLayout(StringRef layout) {
+  return layout == "dense_nchw" || layout == "nchw_row_strip" || layout == "fragmented";
+}
+
+static LogicalResult verifyPlanTensorTypes(Operation* op, Value input, Value output, StringRef kind) {
+  auto inputType = dyn_cast<RankedTensorType>(input.getType());
+  auto outputType = dyn_cast<RankedTensorType>(output.getType());
+  if (!inputType || !outputType)
+    return op->emitOpError() << kind << " requires ranked tensor input and output types";
+  if (inputType.getElementType() != outputType.getElementType())
+    return op->emitOpError() << kind << " requires matching input/output element types";
+  return success();
+}
+
+LogicalResult SpatConv2DPlanOp::verify() {
+  auto inputType = dyn_cast<RankedTensorType>(getInput().getType());
+  auto weightType = dyn_cast<RankedTensorType>(getWeight().getType());
+  auto outputType = dyn_cast<RankedTensorType>(getOutput().getType());
+  if (!inputType || !weightType || !outputType)
+    return emitError("requires ranked tensor input, weight, and output");
+  if (inputType.getRank() != 4 || weightType.getRank() != 4 || outputType.getRank() != 4)
+    return emitError("requires rank-4 input, weight, and output tensors");
+  if (!isKnownLogicalLayout(getLogicalLayout()))
+    return emitError("requires a known logical layout");
+  if (getPads().size() != 4)
+    return emitError("requires exactly four pad values");
+  if (getStrides().size() != 2)
+    return emitError("requires exactly two stride values");
+  if (getDilations().size() != 2)
+    return emitError("requires exactly two dilation values");
+  if (getGroup() < 1)
+    return emitError("requires group >= 1");
+  if (inputType.getElementType() != weightType.getElementType()
+      || inputType.getElementType() != outputType.getElementType()) {
+    return emitError("requires matching input, weight, and output element types");
+  }
+  if (getBias()) {
+    auto biasType = dyn_cast<RankedTensorType>(getBias().getType());
+    if (!biasType)
+      return emitError("requires ranked tensor bias type");
+    if (biasType.getElementType() != outputType.getElementType())
+      return emitError("requires bias element type to match output element type");
   }
   return success();
 }
 
-LogicalResult SpatCompute::verify() {
-  auto& block = getBody().front();
-  unsigned expectedArgCount = getWeights().size() + getInputs().size();
-  if (block.getNumArguments() != expectedArgCount)
-    return emitError("compute body must have weight and input block arguments");
-
-  for (auto [weightIndex, weight] : llvm::enumerate(getWeights())) {
-    auto blockArg = getWeightArgument(weightIndex);
-    if (!blockArg || blockArg->getType() != weight.getType())
-      return emitError("compute weight block argument types must match weight operand types exactly");
-  }
-  for (auto [inputIndex, input] : llvm::enumerate(getInputs())) {
-    auto blockArg = getInputArgument(inputIndex);
-    if (!blockArg || blockArg->getType() != input.getType())
-      return emitError("compute input block argument types must match input operand types exactly");
+LogicalResult SpatReluPlanOp::verify() {
+  if (failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.relu_plan")))
+    return failure();
+  if (!isKnownLogicalLayout(getLogicalLayout()))
+    return emitError("requires a known logical layout");
+  return success();
 }
 
-  if (block.mightHaveTerminator()) {
-    auto yieldOp = dyn_cast_or_null<SpatYieldOp>(block.getTerminator());
-    if (!yieldOp)
-      return emitError("ComputeOp must have a single yield operation");
+LogicalResult SpatBiasAddPlanOp::verify() {
+  if (failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.bias_add_plan")))
+    return failure();
+  if (!isKnownLogicalLayout(getLogicalLayout()))
+    return emitError("requires a known logical layout");
 
-    auto resultTypes = getResultTypes();
-    auto yieldTypes = yieldOp->getOperandTypes();
-    if (resultTypes.size() != yieldTypes.size())
-      return emitError("ComputeOp must have same number of results as yieldOp operands");
+  auto inputType = dyn_cast<RankedTensorType>(getInput().getType());
+  auto biasType = dyn_cast<RankedTensorType>(getBias().getType());
+  auto outputType = dyn_cast<RankedTensorType>(getOutput().getType());
+  if (!inputType || !biasType || !outputType)
+    return emitError("requires ranked tensor input, bias, and output");
+  if (!inputType.hasStaticShape() || !biasType.hasStaticShape() || !outputType.hasStaticShape())
+    return emitError("requires static tensor input, bias, and output");
+  if (inputType != outputType)
+    return emitError("requires matching input and output tensor types");
+  if (outputType.getRank() != 4)
+    return emitError("requires rank-4 input/output tensors");
+  if (getLogicalLayout() != "nchw")
+    return emitError("requires logical layout \"nchw\"");
+  if (biasType.getElementType() != outputType.getElementType())
+    return emitError("requires bias element type to match the output element type");
 
-    for (auto it : llvm::reverse(llvm::zip(resultTypes, yieldTypes))) {
+  ArrayRef<int64_t> biasShape = biasType.getShape();
+  const int64_t channels = outputType.getDimSize(1);
+  const bool supported = biasShape.empty() || (biasShape.size() == 1 && biasShape[0] == channels)
+                         || (biasShape.size() == 2 && biasShape[0] == 1 && biasShape[1] == channels)
+                         || (biasShape.size() == 4 && biasShape[0] == 1 && biasShape[1] == channels
+                             && biasShape[2] == 1 && biasShape[3] == 1);
+  if (!supported)
+    return emitError("requires scalar or per-channel bias broadcastable over NCHW");
+  return success();
+}
+
+LogicalResult SpatBlueprintOp::verify() {
+  auto modeAttr = getModeAttr();
+  bool isFragmentAssembly = modeAttr && modeAttr.getValue() == "fragment_assembly";
+  if (!isFragmentAssembly && failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.blueprint")))
+    return failure();
+  if (!isKnownLogicalLayout(getLogicalLayout()))
+    return emitError("requires a known logical layout");
+  if (!isKnownPhysicalLayout(getPhysicalLayout()))
+    return emitError("requires a known physical layout");
+
+  auto logicalType = dyn_cast<RankedTensorType>(getOutput().getType());
+  if (!logicalType)
+    return emitError("requires ranked tensor output");
+
+  auto offsets = getFragmentOffsets();
+  auto sizes = getFragmentSizes();
+  if (offsets.size() != sizes.size())
+    return emitError("fragment offset and size arrays must have the same length");
+  int64_t rank = logicalType.getRank();
+  if (offsets.empty())
+    return success();
+  if (rank <= 0 || offsets.size() % rank != 0)
+    return emitError("fragment metadata must be a whole number of rank-sized fragments");
+
+  auto verifyBoundsOnly = [&](ArrayRef<int64_t> strideValues) -> LogicalResult {
+    ArrayRef<int64_t> shape = logicalType.getShape();
+    for (int64_t index = 0; index < static_cast<int64_t>(offsets.size()); ++index) {
+      int64_t dim = index % rank;
+      int64_t offset = offsets[index];
+      int64_t size = sizes[index];
+      int64_t stride = strideValues.empty() ? 1 : strideValues[index];
+      if (offset < 0 || size < 0 || stride < 0)
+        return emitError("fragment offsets, sizes, and strides must be non-negative");
+      int64_t logicalDim = shape[dim];
+      if (!ShapedType::isDynamic(logicalDim) && offset + size > logicalDim)
+        return emitError("fragment bounds must stay within the logical tensor shape");
+      if (stride != 1)
+        return emitError("fragment assembly currently requires unit strides");
+    }
+    return success();
+  };
+
+  if (!isFragmentAssembly) {
+    if (failed(verifyBoundsOnly({})))
+      return failure();
+    if (!getFragments().empty())
+      return emitError("legacy blueprint does not accept extra fragment operands");
+    if (getFragmentSourceOffsetsAttr() || getFragmentStridesAttr() || getConflictPolicyAttr()
+        || getCoveragePolicyAttr())
+      return emitError("legacy blueprint does not accept fragment assembly attributes");
+    return success();
+  }
+
+  auto stridesAttr = getFragmentStridesAttr();
+  auto operandIndicesAttr = getFragmentOperandIndicesAttr();
+  auto sourceOffsetsAttr = getFragmentSourceOffsetsAttr();
+  if (!operandIndicesAttr)
+    return emitError("fragment assembly blueprint requires fragment operand indices");
+  if (!sourceOffsetsAttr)
+    return emitError("fragment assembly blueprint requires fragment source offsets");
+  if (!stridesAttr)
+    return emitError("fragment assembly blueprint requires fragment strides");
+  ArrayRef<int64_t> operandIndices = operandIndicesAttr.asArrayRef();
+  ArrayRef<int64_t> sourceOffsets = sourceOffsetsAttr.asArrayRef();
+  ArrayRef<int64_t> strides = stridesAttr.asArrayRef();
+  if (strides.size() != offsets.size())
+    return emitError("fragment stride and offset arrays must have the same length");
+  if (sourceOffsets.size() != operandIndices.size())
+    return emitError("fragment source offset count must match fragment operand index count");
+  if (!getConflictPolicyAttr() || !getCoveragePolicyAttr())
+    return emitError("fragment assembly blueprint requires conflict and coverage policies");
+  if (getConflictPolicy() != "disjoint")
+    return emitError("fragment assembly blueprint currently supports only conflict_policy=\"disjoint\"");
+  if (getCoveragePolicy() != "complete" && getCoveragePolicy() != "partial")
+    return emitError("fragment assembly blueprint coverage_policy must be \"complete\" or \"partial\"");
+
+  SmallVector<Value> operands;
+  operands.push_back(getInput());
+  llvm::append_range(operands, getFragments());
+  int64_t operandCount = static_cast<int64_t>(operands.size());
+  int64_t fragmentCount = static_cast<int64_t>(operandIndices.size());
+  if (operandCount == 0)
+    return emitError("fragment assembly blueprint requires at least one operand");
+  if (static_cast<int64_t>(offsets.size()) != fragmentCount * rank)
+    return emitError("fragment assembly metadata count must match operand count * result rank");
+  if (failed(verifyBoundsOnly(strides)))
+    return failure();
+
+  SmallVector<std::pair<SmallVector<int64_t, 4>, SmallVector<int64_t, 4>>, 8> slices;
+  slices.reserve(static_cast<size_t>(fragmentCount));
+  SmallVector<int64_t, 8> fragmentCountsByOperand(static_cast<size_t>(operandCount), 0);
+  auto expandFlatElementIndex = [](int64_t flatIndex, ArrayRef<int64_t> shape) {
+    SmallVector<int64_t, 4> 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;
+  };
+  for (int64_t fragmentIndex = 0; fragmentIndex < fragmentCount; ++fragmentIndex) {
+    int64_t operandIndex = operandIndices[fragmentIndex];
+    if (operandIndex < 0 || operandIndex >= operandCount)
+      return emitError("fragment assembly operand index is out of range");
+    if (sourceOffsets[fragmentIndex] < 0)
+      return emitError("fragment assembly source offsets must be nonnegative");
+
+    auto operandType = dyn_cast<RankedTensorType>(operands[operandIndex].getType());
+    if (!operandType || !operandType.hasStaticShape())
+      return emitError("fragment assembly blueprint requires static ranked tensor operands");
+    if (operandType.getRank() != rank)
+      return emitError("fragment assembly blueprint requires operand/result rank match");
+
+    SmallVector<int64_t, 4> fragmentOffsets;
+    SmallVector<int64_t, 4> fragmentSizes;
+    fragmentOffsets.reserve(rank);
+    fragmentSizes.reserve(rank);
+    for (int64_t dim = 0; dim < rank; ++dim) {
+      int64_t flatIndex = fragmentIndex * rank + dim;
+      fragmentOffsets.push_back(offsets[flatIndex]);
+      fragmentSizes.push_back(sizes[flatIndex]);
+    }
+
+    ++fragmentCountsByOperand[static_cast<size_t>(operandIndex)];
+    int64_t fragmentElements = 1;
+    for (int64_t dim = 0; dim < rank; ++dim)
+      fragmentElements *= fragmentSizes[dim];
+    if (sourceOffsets[fragmentIndex] + fragmentElements > operandType.getNumElements())
+      return emitError("fragment assembly source offset exceeds the operand bounds");
+    SmallVector<int64_t, 4> sourceSliceOffsets =
+      expandFlatElementIndex(sourceOffsets[fragmentIndex], operandType.getShape());
+    for (int64_t dim = 0; dim < rank; ++dim)
+      if (sourceSliceOffsets[dim] + fragmentSizes[dim] > operandType.getDimSize(dim))
+        return emitError("fragment assembly source offset must describe a valid unit-stride slice");
+
+    for (const auto& [existingOffsets, existingSizes] : slices) {
+      bool overlaps = true;
+      for (int64_t dim = 0; dim < rank; ++dim) {
+        int64_t begin = fragmentOffsets[dim];
+        int64_t end = begin + fragmentSizes[dim];
+        int64_t existingBegin = existingOffsets[dim];
+        int64_t existingEnd = existingBegin + existingSizes[dim];
+        if (end <= existingBegin || existingEnd <= begin) {
+          overlaps = false;
+          break;
+        }
+      }
+      if (overlaps)
+        return emitError("fragment assembly blueprint requires disjoint static slices");
+    }
+    slices.push_back({std::move(fragmentOffsets), std::move(fragmentSizes)});
+  }
+
+  for (int64_t operandIndex = 0; operandIndex < operandCount; ++operandIndex) {
+    if (fragmentCountsByOperand[static_cast<size_t>(operandIndex)] == 0)
+      return emitError("fragment assembly blueprint requires every operand to contribute at least one fragment");
+  }
+
+  if (getCoveragePolicy() == "complete") {
+    int64_t covered = 0;
+    int64_t logicalElements = 1;
+    for (int64_t dimSize : logicalType.getShape()) {
+      if (ShapedType::isDynamic(dimSize))
+        return emitError("fragment assembly complete coverage requires static result shape");
+      logicalElements *= dimSize;
+    }
+    for (const auto& [ignoredOffsets, fragmentSizes] : slices) {
+      int64_t fragmentElements = 1;
+      for (int64_t dimSize : fragmentSizes)
+        fragmentElements *= dimSize;
+      covered += fragmentElements;
+    }
+    if (covered != logicalElements)
+      return emitError("fragment assembly complete coverage must cover the whole result exactly");
+  }
+
+  return success();
+}
+
+LogicalResult SpatMaterializeLayoutOp::verify() {
+  if (failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.materialize_layout")))
+    return failure();
+  if (!isKnownLogicalLayout(getLogicalLayout()))
+    return emitError("requires a known logical layout");
+  if (!isKnownPhysicalLayout(getSourcePhysicalLayout()))
+    return emitError("requires a known source physical layout");
+  if (!isKnownPhysicalLayout(getTargetPhysicalLayout()))
+    return emitError("requires a known target physical layout");
+  return success();
+}
+
+LogicalResult verifyComputeResultsUses(Operation* op) {
+  if (!isAnySpatialComputeLike(op))
+    return op->emitError("verifyComputeResultUses: op is not a Spatial compute-like operation");
+  if (!llvm::all_of(op->getResults(), [](Value result) {
+        return llvm::all_of(result.getUsers(), [result](Operation* op) {
+          if (isRecordedDeferredCommunicationSource(op, result))
+            return true;
+          return !isAnySpatialComputeLike(op->getParentOp());
+        });
+      })) {
+    return op->emitError("compute result used directly inside another Spatial compute body");
+  }
+  return success();
+}
+
+template <typename ComputeOpTy>
+LogicalResult verifyComputeLikeOp(ComputeOpTy compute, StringRef opName) {
+  unsigned expectedArgCount = compute.getWeights().size() + compute.getInputs().size();
+  bool isScheduled = isa<SpatScheduledCompute>(compute.getOperation());
+  if (compute.getBody().empty())
+    return compute.emitOpError("compute body must have at least one block");
+
+  SmallVector<Type> yieldedTypes;
+  for (Block& block : compute.getBody()) {
+    if ((!isScheduled && block.getNumArguments() != expectedArgCount)
+        || (isScheduled && block.getNumArguments() < expectedArgCount))
+      return compute.emitOpError("compute body must have weight and input block arguments");
+
+    for (auto [weightIndex, weight] : llvm::enumerate(compute.getWeights()))
+      if (block.getArgument(weightIndex).getType() != weight.getType())
+        return compute.emitOpError("compute weight block argument types must match weight operand types exactly");
+    for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs()))
+      if (block.getArgument(compute.getWeights().size() + inputIndex).getType() != input.getType())
+        return compute.emitOpError("compute input block argument types must match input operand types exactly");
+
+    Operation* terminator = block.getTerminator();
+    if (auto yieldOp = dyn_cast_or_null<SpatYieldOp>(terminator)) {
+      if (isScheduled)
+        return compute.emitOpError("scheduled compute blocks must terminate with spat.block_yield");
+      llvm::append_range(yieldedTypes, yieldOp->getOperandTypes());
+      continue;
+    }
+    auto blockYield = dyn_cast_or_null<SpatBlockYieldOp>(terminator);
+    if (!blockYield || !isScheduled)
+      return compute.emitOpError("ComputeOp must have a single yield operation");
+    if (blockYield->getNumSuccessors() == 0)
+      llvm::append_range(yieldedTypes, blockYield->getOperandTypes());
+  }
+
+  auto resultTypes = compute.getResultTypes();
+  if (resultTypes.size() != yieldedTypes.size())
+    return compute.emitOpError("ComputeOp must have same number of results as yielded operands");
+
+  for (auto it : llvm::reverse(llvm::zip(resultTypes, yieldedTypes))) {
     auto resultType = std::get<0>(it);
     auto yieldType = std::get<1>(it);
 
     if (resultType != yieldType || failed(verifyCompatibleShape(resultType, yieldType)))
-        return emitError("ComputeOp output must be of the same type as yieldOp operand");
+      return compute.emitOpError("ComputeOp output must be of the same type as yieldOp operand");
 
     if (auto resultRankedType = dyn_cast<RankedTensorType>(resultType)) {
       if (auto yieldRankedType = dyn_cast<RankedTensorType>(yieldType)) {
         if (resultRankedType.getEncoding() != yieldRankedType.getEncoding())
-            return emitError("ComputeOp output must have the same encoding as yieldOp operand");
+          return compute.emitOpError("ComputeOp output has an encoding while yieldOp operand does not have one");
       }
       else {
-          return emitError("ComputeOp output has an encoding while yieldOp operand does not have one");
+        return compute.emitOpError("ComputeOp output must have the same encoding as yieldOp operand");
       }
     }
     else if (dyn_cast<RankedTensorType>(yieldType)) {
-        return emitError("ComputeOp output must not have an encoding if yieldOp operand has one");
-      }
+      return compute.emitOpError("ComputeOp output must not have an encoding if yieldOp operand has one");
     }
   }
 
-  for (unsigned inputIndex = 0; inputIndex < getInputs().size(); ++inputIndex)
-    if (auto inputArg = getInputArgument(inputIndex); !inputArg || inputArg->use_empty())
-      return emitError("ComputeOp block argument is not used");
-  if (failed(verifyStaticWeights(*this, "compute")))
+  if (compute.getBody().hasOneBlock())
+    for (unsigned inputIndex = 0; inputIndex < compute.getInputs().size(); ++inputIndex)
+      if (auto inputArg = compute.getInputArgument(inputIndex); !inputArg || inputArg->use_empty())
+        return compute.emitOpError("ComputeOp block argument is not used");
+  if (failed(verifyStaticWeights(compute, opName)))
     return failure();
-  if (failed(verifyOnlyConstantExternalValues(this->getOperation(), getBody(), "spat.compute")))
+  if (failed(verifyOnlyConstantExternalValues(compute.getOperation(), compute.getBody(), opName)))
     return failure();
-  if (failed(verifyComputeResultsUses(this->getOperation())))
+  if (failed(verifyComputeResultsUses(compute.getOperation())))
     return failure();
   return success();
 }
 
-LogicalResult SpatComputeBatch::verify() {
-  int32_t count = getLaneCount();
+LogicalResult SpatGraphCompute::verify() { return verifyComputeLikeOp(*this, "spat.graph_compute"); }
+
+LogicalResult SpatScheduledCompute::verify() { return verifyComputeLikeOp(*this, "spat.scheduled_compute"); }
+
+LogicalResult SpatBlockYieldOp::verify() {
+  if (getOperation()->getNumSuccessors() > 1)
+    return emitOpError("may target at most one next scheduled block");
+  Operation* parent = getOperation()->getParentOp();
+  if (!isa_and_nonnull<SpatScheduledCompute>(parent))
+    return emitOpError("expected spat.scheduled_compute parent");
+  if (getOperation()->getNumSuccessors() == 1) {
+    Block* next = getOperation()->getSuccessor(0);
+    if (getOperation()->getNumOperands() != next->getNumArguments())
+      return emitOpError("successor operand count must match next block argument count");
+    for (auto [operand, argument] : llvm::zip(getOperation()->getOperands(), next->getArguments()))
+      if (operand.getType() != argument.getType())
+        return emitOpError("successor operand types must match next block argument types");
+  }
+  return success();
+}
+
+LogicalResult SpatDeferredCommunicationOp::verify() {
+  if (getSources().empty())
+    return emitOpError("requires at least one source");
+  if (failed(verifyRegionArguments(getOperation(), getBody(), getSources(), "spat.deferred_communication")))
+    return failure();
+  return verifyYieldTypes(getOperation(), getBody(), getOperation()->getResultTypes(), "spat.deferred_communication");
+}
+
+template <typename ComputeBatchOpTy>
+LogicalResult verifyComputeBatchLikeOp(ComputeBatchOpTy batch, StringRef opName) {
+  int32_t count = batch.getLaneCount();
   if (count <= 0)
-    return emitError("laneCount must be positive");
+    return batch.emitOpError("laneCount must be positive");
 
   auto laneCountSz = static_cast<size_t>(count);
 
-  if (auto coreIdAttr = (*this)->getAttr(kCoreIdsAttrName)) {
+  if (auto coreIdAttr = batch->getAttr(kCoreIdsAttrName)) {
     auto coreIdsAttr = dyn_cast<DenseI32ArrayAttr>(coreIdAttr);
     if (!coreIdsAttr)
-      return emitError("compute_batch coreIds attribute must be a dense i32 array");
+      return batch.emitOpError("compute_batch coreIds attribute must be a dense i32 array");
     if (coreIdsAttr.size() != static_cast<int64_t>(laneCountSz))
-      return emitError("compute_batch coreIds array length must match laneCount");
+      return batch.emitOpError("compute_batch coreIds array length must match laneCount");
     if (llvm::any_of(coreIdsAttr.asArrayRef(), [](int32_t coreId) { return coreId < 0; }))
-      return emitError("compute_batch coreIds values must be non-negative");
+      return batch.emitOpError("compute_batch coreIds values must be non-negative");
     DenseSet<int32_t> seenCoreIds;
     for (int32_t coreId : coreIdsAttr.asArrayRef())
       if (!seenCoreIds.insert(coreId).second)
-        return emitError("compute_batch coreIds values must be unique");
+        return batch.emitOpError("compute_batch coreIds values must be unique");
   }
 
-  Block& block = getBody().front();
+  if (batch.getBody().empty())
+    return batch.emitOpError("compute_batch body must have at least one block");
+
+  unsigned expectedArgCount = 1 + batch.getWeights().size() + batch.getInputs().size() + batch.getNumResults();
+  bool verifyLaneSliceOffsets = !isa<SpatScheduledComputeBatch>(batch.getOperation());
+  for (Block& block : batch.getBody()) {
     if (block.getNumArguments() == 0)
-    return emitError("compute_batch body must have exactly one lane block argument");
-  unsigned expectedArgCount = 1 + getWeights().size() + getInputs().size() + getNumResults();
+      return batch.emitOpError("compute_batch body must have exactly one lane block argument");
     if (block.getNumArguments() != expectedArgCount)
-    return emitError("compute_batch body block arguments must match lane, weight, input, and output operands/results");
-  auto laneArg = getLaneArgument();
-  if (!laneArg || !laneArg->getType().isIndex())
-    return emitError("compute_batch first block argument must have index type");
+      return batch.emitOpError(
+        "compute_batch body block arguments must match lane, weight, input, and output operands/results");
+    if (!block.getArgument(0).getType().isIndex())
+      return batch.emitOpError("compute_batch first block argument must have index type");
 
-  for (auto [weightIndex, weight] : llvm::enumerate(getWeights())) {
-    auto blockArg = getWeightArgument(weightIndex);
-    if (!blockArg || blockArg->getType() != weight.getType())
-      return emitError("compute_batch weight block argument types must match weight operand types exactly");
-  }
-  for (auto [inputIndex, input] : llvm::enumerate(getInputs())) {
-    auto blockArg = getInputArgument(inputIndex);
-    if (!blockArg || blockArg->getType() != input.getType())
-      return emitError("compute_batch input block argument types must match input operand types exactly");
-  }
-  for (auto [resultIndex, resultType] : llvm::enumerate(getResultTypes())) {
-    auto blockArg = getOutputArgument(resultIndex);
-    if (!blockArg || blockArg->getType() != resultType)
-      return emitError("compute_batch output block argument types must match result types exactly");
+    for (auto [weightIndex, weight] : llvm::enumerate(batch.getWeights()))
+      if (block.getArgument(1 + weightIndex).getType() != weight.getType())
+        return batch.emitOpError("compute_batch weight block argument types must match weight operand types exactly");
+    for (auto [inputIndex, input] : llvm::enumerate(batch.getInputs()))
+      if (block.getArgument(1 + batch.getWeights().size() + inputIndex).getType() != input.getType())
+        return batch.emitOpError("compute_batch input block argument types must match input operand types exactly");
+    for (auto [resultIndex, resultType] : llvm::enumerate(batch.getResultTypes()))
+      if (block.getArgument(1 + batch.getWeights().size() + batch.getInputs().size() + resultIndex).getType()
+          != resultType)
+        return batch.emitOpError("compute_batch output block argument types must match result types exactly");
+
+    if (failed(verifyBatchBody(batch, block, verifyLaneSliceOffsets)))
+      return failure();
   }
 
-  if (failed(verifyComputeResultsUses(this->getOperation())))
+  if (failed(verifyComputeResultsUses(batch.getOperation())))
     return failure();
-  if (failed(verifyStaticWeights(*this, "compute_batch")))
+  if (failed(verifyStaticWeights(batch, opName)))
     return failure();
-  if (failed(verifyOnlyConstantExternalValues(this->getOperation(), getBody(), "spat.compute_batch")))
+  if (failed(verifyOnlyConstantExternalValues(batch.getOperation(), batch.getBody(), opName)))
     return failure();
-  return verifyBatchBody(*this, block);
+  return success();
+}
+
+LogicalResult SpatGraphComputeBatch::verify() { return verifyComputeBatchLikeOp(*this, "spat.graph_compute_batch"); }
+
+LogicalResult SpatScheduledComputeBatch::verify() {
+  return verifyComputeBatchLikeOp(*this, "spat.scheduled_compute_batch");
 }
 
 LogicalResult SpatInParallelOp::verify() {
-  auto batchOp = getOperation()->getParentOfType<SpatComputeBatch>();
-  if (!batchOp)
-    return emitOpError("expected spat.compute_batch parent");
-  if (batchOp.getNumResults() == 0)
+  Operation* parent = getOperation()->getParentOp();
+  if (!isAnySpatialComputeBatchLike(parent))
+    return emitOpError("expected spat.graph_compute_batch or spat.scheduled_compute_batch parent");
+  if (parent->getNumResults() == 0)
     return emitOpError("requires a resultful spat.compute_batch parent");
 
-  auto laneArg = batchOp.getLaneArgument();
+  std::optional<BlockArgument> laneArg;
+  if (auto graphBatch = dyn_cast<SpatGraphComputeBatch>(parent))
+    laneArg = graphBatch.getLaneArgument();
+  else
+    laneArg = cast<SpatScheduledComputeBatch>(parent).getLaneArgument();
   if (!laneArg)
     return emitOpError("expected compute_batch lane block argument");
   for (Operation& op : getRegion().front().getOperations()) {
@@ -494,7 +915,10 @@ LogicalResult SpatInParallelOp::verify() {
 
     MutableOperandRange destinations = insertSliceOp.getUpdatedDestinations();
     for (OpOperand& destination : destinations)
-      if (!isBatchOutputArgument(batchOp, destination.get()))
+      if ((isa<SpatGraphComputeBatch>(parent)
+             && !isBatchOutputArgument(cast<SpatGraphComputeBatch>(parent), destination.get()))
+          || (isa<SpatScheduledComputeBatch>(parent)
+              && !isBatchOutputArgument(cast<SpatScheduledComputeBatch>(parent), destination.get())))
         return op.emitOpError("may only insert into a compute_batch output block argument");
   }
 

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp

@@ -0,0 +1,426 @@
+#include "DeferredCommunicationPlanning.hpp"
+
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+
+namespace onnx_mlir {
+namespace spatial {
+
+using namespace mlir;
+
+namespace {
+
+struct LaneProducerInfo {
+  ProducerValueRef producer;
+  Value originalSource;
+  DeferredEndpoint producerEndpoint;
+  DeferredEndpoint consumerEndpoint;
+};
+
+static Value getBlockOperand(Block &block, ValueRange operands, Value value, unsigned firstArgument = 0) {
+  auto it = llvm::find(operands, value);
+  assert(it != operands.end() && "missing scheduled operand");
+  return block.getArgument(firstArgument + std::distance(operands.begin(), it));
+}
+
+static FailureOr<Value> getOriginalProducerValue(const ProducerValueRef &producer) {
+  auto outputs = getComputeInstanceOutputValues(producer.instance);
+  if (producer.resultIndex >= outputs.size())
+    return failure();
+  return outputs[producer.resultIndex];
+}
+
+static DeferredEndpoint getDeferredEndpoint(const ComputeInstance &instance,
+                                            const MergeScheduleResult &schedule) {
+  size_t cpu = getScheduledCpuForComputeInstance(instance, schedule);
+  return DeferredEndpoint {
+    instance,
+    cpu,
+    getCanonicalPeftClassId(cpu, schedule),
+    instance.laneStart,
+    instance.laneCount
+  };
+}
+
+static void setDeferredTransferMetadata(SpatDeferredCommunicationOp transfer,
+                                        OpBuilder &builder,
+                                        const DeferredTransferMetadata &transferMetadata) {
+  SmallVector<int64_t> sourceCpus;
+  SmallVector<int64_t> sourceClasses;
+  SmallVector<int64_t> sourceLaneRanges;
+  SmallVector<int64_t> targetCpus;
+  SmallVector<int64_t> targetClasses;
+  SmallVector<int64_t> targetLaneRanges;
+  for (const DeferredEndpoint &endpoint : transferMetadata.producers) {
+    sourceCpus.push_back(static_cast<int64_t>(endpoint.cpu));
+    sourceClasses.push_back(static_cast<int64_t>(endpoint.canonicalClassId));
+    sourceLaneRanges.push_back(static_cast<int64_t>(endpoint.laneStart));
+    sourceLaneRanges.push_back(static_cast<int64_t>(endpoint.laneCount));
+  }
+  for (const DeferredEndpoint &endpoint : transferMetadata.consumers) {
+    targetCpus.push_back(static_cast<int64_t>(endpoint.cpu));
+    targetClasses.push_back(static_cast<int64_t>(endpoint.canonicalClassId));
+    targetLaneRanges.push_back(static_cast<int64_t>(endpoint.laneStart));
+    targetLaneRanges.push_back(static_cast<int64_t>(endpoint.laneCount));
+  }
+  transfer->setAttr("source_cpus", builder.getDenseI64ArrayAttr(sourceCpus));
+  transfer->setAttr("source_classes", builder.getDenseI64ArrayAttr(sourceClasses));
+  transfer->setAttr("source_lane_ranges", builder.getDenseI64ArrayAttr(sourceLaneRanges));
+  transfer->setAttr("target_cpus", builder.getDenseI64ArrayAttr(targetCpus));
+  transfer->setAttr("target_classes", builder.getDenseI64ArrayAttr(targetClasses));
+  transfer->setAttr("target_lane_ranges", builder.getDenseI64ArrayAttr(targetLaneRanges));
+  // `batched=true` means this deferred communication belongs to a
+  // spat.scheduled_compute_batch block and the array attrs are indexed by
+  // scheduled lane. The legacy scalar attrs on the op keep representative values
+  // for compatibility; Phase 2 must consume the arrays when batched=true.
+  transfer->setAttr("batched", builder.getBoolAttr(transferMetadata.isBatched));
+  if (transferMetadata.isBatched) {
+    transfer->setAttr("source_operand_for_scheduled_lane",
+                      builder.getDenseI64ArrayAttr(transferMetadata.sourceOperandForScheduledLane));
+    transfer->setAttr("multi_source_payload", builder.getBoolAttr(transferMetadata.multiSourcePayload));
+  }
+}
+
+static FailureOr<Value> projectDeferredPayload(Value selectedSource,
+                                               Value consumerInput,
+                                               OpBuilder &builder,
+                                               Location loc) {
+  if (selectedSource.getType() == consumerInput.getType())
+    return selectedSource;
+
+  auto slice = consumerInput.getDefiningOp<tensor::ExtractSliceOp>();
+  if (!slice || slice.getSource().getType() != selectedSource.getType())
+    return failure();
+
+  IRMapping mapper;
+  mapper.map(slice.getSource(), selectedSource);
+  return cast<tensor::ExtractSliceOp>(builder.clone(*slice, mapper)).getResult();
+}
+
+static FailureOr<Value> buildSelectedDeferredSource(OpBuilder &builder,
+                                                    Location loc,
+                                                    SpatDeferredCommunicationOp transfer,
+                                                    Value scheduledLane,
+                                                    ValueRange sourceBlockArgs,
+                                                    DenseI64ArrayAttr sourceOperandForScheduledLaneAttr) {
+  SmallVector<int64_t> sourceOperandForScheduledLane(sourceOperandForScheduledLaneAttr.asArrayRef().begin(),
+                                                     sourceOperandForScheduledLaneAttr.asArrayRef().end());
+  Value sourceOperandIndexTable =
+    createI64LookupTableConstant(builder, transfer.getOperation(), sourceOperandForScheduledLane);
+  Value sourceIndexI64 =
+    tensor::ExtractOp::create(builder, loc, sourceOperandIndexTable, ValueRange {scheduledLane}).getResult();
+  Value sourceIndex = arith::IndexCastOp::create(builder, loc, builder.getIndexType(), sourceIndexI64).getResult();
+
+  Type elementType = sourceBlockArgs.front().getType();
+  auto sourceTensorType = dyn_cast<RankedTensorType>(elementType);
+  if (!sourceTensorType || !sourceTensorType.hasStaticShape())
+    return transfer.emitOpError(
+             "scheduled batch deferred communication with multiple original graph SSA sources requires compatible static ranked tensor sources"),
+           failure();
+
+  for (Value source : sourceBlockArgs) {
+    auto rankedSourceType = dyn_cast<RankedTensorType>(source.getType());
+    if (!rankedSourceType || !rankedSourceType.hasStaticShape() || rankedSourceType != sourceTensorType)
+      return transfer.emitOpError(
+               "scheduled batch deferred communication with multiple original graph SSA sources requires compatible static ranked tensor sources"),
+             failure();
+  }
+  if (sourceTensorType.getRank() == 0)
+    return transfer.emitOpError(
+             "scheduled batch deferred communication with multiple original graph SSA sources cannot yet build compact tensor source selection"),
+           failure();
+
+  SmallVector<int64_t> stackedShape {static_cast<int64_t>(sourceBlockArgs.size())};
+  llvm::append_range(stackedShape, sourceTensorType.getShape());
+  RankedTensorType stackedType = RankedTensorType::get(stackedShape, sourceTensorType.getElementType());
+  auto stackedEmpty = createEmptyTensorForType(builder, loc, stackedType);
+  if (failed(stackedEmpty))
+    return transfer.emitOpError(
+             "scheduled batch deferred communication with multiple original graph SSA sources cannot yet build compact tensor source selection"),
+           failure();
+  Value stackedSources = *stackedEmpty;
+
+  SmallVector<OpFoldResult> insertSizes;
+  SmallVector<OpFoldResult> unitStrides(stackedType.getRank(), builder.getIndexAttr(1));
+  insertSizes.push_back(builder.getIndexAttr(1));
+  for (int64_t dim : sourceTensorType.getShape())
+    insertSizes.push_back(builder.getIndexAttr(dim));
+
+  for (auto [sourceIndexValue, source] : llvm::enumerate(sourceBlockArgs)) {
+    SmallVector<int64_t> expandedShape {1};
+    llvm::append_range(expandedShape, sourceTensorType.getShape());
+    SmallVector<ReassociationIndices> reassociation;
+    reassociation.push_back({0, 1});
+    for (int64_t dim = 1; dim < sourceTensorType.getRank(); ++dim)
+      reassociation.push_back({dim + 1});
+    Value expandedSource = tensor::ExpandShapeOp::create(builder,
+                                                         loc,
+                                                         RankedTensorType::get(expandedShape, sourceTensorType.getElementType()),
+                                                         source,
+                                                         reassociation)
+                             .getResult();
+    SmallVector<OpFoldResult> insertOffsets(stackedType.getRank(), builder.getIndexAttr(0));
+    insertOffsets[0] = builder.getIndexAttr(sourceIndexValue);
+    stackedSources = tensor::InsertSliceOp::create(
+                       builder, loc, expandedSource, stackedSources, insertOffsets, insertSizes, unitStrides)
+                       .getResult();
+  }
+
+  SmallVector<OpFoldResult> extractOffsets(stackedType.getRank(), builder.getIndexAttr(0));
+  SmallVector<OpFoldResult> extractSizes;
+  SmallVector<OpFoldResult> extractStrides(stackedType.getRank(), builder.getIndexAttr(1));
+  extractOffsets[0] = sourceIndex;
+  extractSizes.push_back(builder.getIndexAttr(1));
+  for (int64_t dim : sourceTensorType.getShape())
+    extractSizes.push_back(builder.getIndexAttr(dim));
+  SmallVector<int64_t> selectedSliceShape {1};
+  llvm::append_range(selectedSliceShape, sourceTensorType.getShape());
+  RankedTensorType selectedSliceType = RankedTensorType::get(selectedSliceShape, sourceTensorType.getElementType());
+  Value selectedSlice = tensor::ExtractSliceOp::create(
+                          builder,
+                          loc,
+                          selectedSliceType,
+                          stackedSources,
+                          extractOffsets,
+                          extractSizes,
+                          extractStrides)
+                          .getResult();
+
+  SmallVector<ReassociationIndices> collapseReassociation;
+  collapseReassociation.push_back({0, 1});
+  for (int64_t dim = 1; dim < sourceTensorType.getRank(); ++dim)
+    collapseReassociation.push_back({dim + 1});
+  return tensor::CollapseShapeOp::create(builder, loc, sourceTensorType, selectedSlice, collapseReassociation)
+    .getResult();
+}
+
+static LogicalResult finalizeDeferredCommunication(OpBuilder &builder,
+                                                   Location loc,
+                                                   Value consumerInput,
+                                                   SpatDeferredCommunicationOp transfer,
+                                                   Value &result) {
+  Block *block = builder.createBlock(&transfer.getBody(),
+                                     transfer.getBody().end(),
+                                     TypeRange {transfer.getSources().getTypes()},
+                                     SmallVector<Location>(transfer.getSources().size(), loc));
+  builder.setInsertionPointToStart(block);
+
+  Value selectedSource = block->getArgument(0);
+  if (auto parentScheduled = dyn_cast<SpatScheduledComputeBatch>(transfer->getParentOp())) {
+    auto sourceOperandForScheduledLaneAttr =
+      transfer->getAttrOfType<DenseI64ArrayAttr>("source_operand_for_scheduled_lane");
+    auto multiSourcePayloadAttr = transfer->getAttrOfType<BoolAttr>("multi_source_payload");
+    if (sourceOperandForScheduledLaneAttr && sourceOperandForScheduledLaneAttr.size() == parentScheduled.getLaneCount()) {
+      if (multiSourcePayloadAttr && multiSourcePayloadAttr.getValue()) {
+        FailureOr<Value> multiSourceSelection = buildSelectedDeferredSource(
+          builder,
+          loc,
+          transfer,
+          transfer->getBlock()->getArgument(0),
+          block->getArguments(),
+          sourceOperandForScheduledLaneAttr);
+        if (failed(multiSourceSelection))
+          return failure();
+        selectedSource = *multiSourceSelection;
+      } else {
+        selectedSource = block->getArgument(sourceOperandForScheduledLaneAttr.asArrayRef().front());
+      }
+    }
+  }
+
+  FailureOr<Value> yielded = projectDeferredPayload(selectedSource, consumerInput, builder, loc);
+  if (failed(yielded) || (*yielded).getType() != consumerInput.getType())
+    return transfer.emitOpError("cannot derive deferred communication payload from original graph producer value");
+
+  SpatYieldOp::create(builder, loc, *yielded);
+  result = transfer.getOutput();
+  builder.setInsertionPointAfter(transfer);
+  return success();
+}
+
+static LogicalResult createSingleLaneDeferredCommunication(OpBuilder &builder,
+                                                           Location loc,
+                                                           Value consumerInput,
+                                                           Value originalSource,
+                                                           const ProducerValueRef &producer,
+                                                           const ComputeInstance &consumer,
+                                                           const MergeScheduleResult &schedule,
+                                                           uint64_t exchangeId,
+                                                           Value &result) {
+  DeferredTransferMetadata transferMetadata;
+  transferMetadata.producers.push_back(getDeferredEndpoint(producer.instance, schedule));
+  transferMetadata.consumers.push_back(getDeferredEndpoint(consumer, schedule));
+  auto transfer = SpatDeferredCommunicationOp::create(
+    builder,
+    loc,
+    consumerInput.getType(),
+    ValueRange {originalSource},
+    builder.getStringAttr(("x" + std::to_string(exchangeId)).c_str()),
+    builder.getStringAttr(getInstanceName(producer.instance)),
+    builder.getStringAttr(getInstanceName(consumer)),
+    builder.getI64IntegerAttr(static_cast<int64_t>(transferMetadata.producers.front().canonicalClassId)),
+    builder.getI64IntegerAttr(static_cast<int64_t>(transferMetadata.consumers.front().canonicalClassId)),
+    builder.getI64IntegerAttr(static_cast<int64_t>(transferMetadata.producers.front().cpu)),
+    builder.getI64IntegerAttr(static_cast<int64_t>(transferMetadata.consumers.front().cpu)),
+    builder.getI64IntegerAttr(producer.instance.laneStart),
+    builder.getI64IntegerAttr(consumer.laneStart),
+    builder.getStringAttr(consumerInput.getDefiningOp<tensor::ExtractSliceOp>() ? "projected" : "ordinary"),
+    builder.getI64IntegerAttr(static_cast<int64_t>(producer.resultIndex)),
+    builder.getDenseI64ArrayAttr({}),
+    builder.getI64IntegerAttr(-1));
+  setDeferredTransferMetadata(transfer, builder, transferMetadata);
+  return finalizeDeferredCommunication(builder, loc, consumerInput, transfer, result);
+}
+
+static LogicalResult createScheduledLaneDeferredCommunication(OpBuilder &builder,
+                                                              Location loc,
+                                                              Value consumerInput,
+                                                              ValueRange originalSources,
+                                                              const ProducerValueRef &producer,
+                                                              const DeferredTransferMetadata &transferMetadata,
+                                                              uint64_t exchangeId,
+                                                              Value &result) {
+  assert(transferMetadata.isBatched && "expected scheduled-lane deferred communication metadata");
+  assert(!transferMetadata.producers.empty() && !transferMetadata.consumers.empty() && "expected non-empty endpoints");
+  auto transfer = SpatDeferredCommunicationOp::create(
+    builder,
+    loc,
+    consumerInput.getType(),
+    originalSources,
+    builder.getStringAttr(("x" + std::to_string(exchangeId)).c_str()),
+    builder.getStringAttr(getInstanceName(producer.instance)),
+    builder.getStringAttr(getInstanceName(transferMetadata.consumers.front().instance)),
+    builder.getI64IntegerAttr(static_cast<int64_t>(transferMetadata.producers.front().canonicalClassId)),
+    builder.getI64IntegerAttr(static_cast<int64_t>(transferMetadata.consumers.front().canonicalClassId)),
+    builder.getI64IntegerAttr(static_cast<int64_t>(transferMetadata.producers.front().cpu)),
+    builder.getI64IntegerAttr(static_cast<int64_t>(transferMetadata.consumers.front().cpu)),
+    builder.getI64IntegerAttr(producer.instance.laneStart),
+    builder.getI64IntegerAttr(transferMetadata.consumers.front().laneStart),
+    builder.getStringAttr(consumerInput.getDefiningOp<tensor::ExtractSliceOp>() ? "projected" : "ordinary"),
+    builder.getI64IntegerAttr(static_cast<int64_t>(producer.resultIndex)),
+    builder.getDenseI64ArrayAttr({}),
+    builder.getI64IntegerAttr(-1));
+  setDeferredTransferMetadata(transfer, builder, transferMetadata);
+  return finalizeDeferredCommunication(builder, loc, consumerInput, transfer, result);
+}
+
+} // namespace
+
+LogicalResult mapSingleCpuInput(OpBuilder &builder,
+                                       Location loc,
+                                       Value input,
+                                       const ComputeInstance &consumerInstance,
+                                       const MergeScheduleResult &schedule,
+                                       ValueRange scheduledInputs,
+                                       Block &block,
+                                       unsigned firstInputArgument,
+                                       ArrayRef<ProducerValueKey> carriedKeys,
+                                       uint64_t &exchangeId,
+                                       Value &mapped) {
+  std::optional<ProducerValueRef> producer = getProducerValueRef(input, &consumerInstance);
+  if (!producer) {
+    mapped = getBlockOperand(block, scheduledInputs, input, firstInputArgument);
+    return success();
+  }
+
+  ProducerValueKey producerKey {producer->instance, producer->resultIndex};
+  auto carriedIt = llvm::find(carriedKeys, producerKey);
+  if (carriedIt != carriedKeys.end()) {
+    mapped = block.getArgument(firstInputArgument + scheduledInputs.size() + std::distance(carriedKeys.begin(), carriedIt));
+    return success();
+  }
+
+  FailureOr<Value> originalSource = getOriginalProducerValue(*producer);
+  if (failed(originalSource))
+    return emitError(loc) << "cannot resolve original graph producer value for " << getInstanceName(producer->instance);
+  return createSingleLaneDeferredCommunication(
+    builder, loc, input, *originalSource, *producer, consumerInstance, schedule, exchangeId++, mapped);
+}
+
+static FailureOr<unsigned> getComputeInstanceInputIndex(const ComputeInstance &instance, Value input) {
+  auto inputs = getComputeInstanceInputs(instance);
+  auto it = llvm::find(inputs, input);
+  if (it == inputs.end())
+    return failure();
+  return static_cast<unsigned>(std::distance(inputs.begin(), it));
+}
+
+LogicalResult mapMultiCpuTupleInput(OpBuilder &builder,
+                                           Location loc,
+                                           Value input,
+                                           const ComputeStepTuple &stepTuple,
+                                           const PeftClassPlan &peftClassPlan,
+                                           const MergeScheduleResult &schedule,
+                                           ValueRange scheduledInputs,
+                                           Block &block,
+                                           unsigned firstInputArgument,
+                                           uint64_t &exchangeId,
+                                           Value &mapped) {
+  assert(!stepTuple.instances.empty() && "expected non-empty step tuple");
+  const ComputeInstance &representative = stepTuple.instances.front();
+  std::optional<ProducerValueRef> representativeProducer = getProducerValueRef(input, &representative);
+  if (!representativeProducer) {
+    FailureOr<unsigned> inputIndex = getComputeInstanceInputIndex(representative, input);
+    if (failed(inputIndex))
+      return emitError(loc) << "cannot resolve scheduled batch step host input index";
+    for (const ComputeInstance &instance : stepTuple.instances) {
+      auto instanceInputs = getComputeInstanceInputs(instance);
+      if (*inputIndex >= instanceInputs.size())
+        return emitError(loc) << "scheduled batch step host input index out of range";
+      if (instanceInputs[*inputIndex] != input || getProducerValueRef(instanceInputs[*inputIndex], &instance))
+        return emitError(loc) << "scheduled batch step requires identical host inputs across lanes";
+    }
+    mapped = getBlockOperand(block, scheduledInputs, input, firstInputArgument);
+    return success();
+  }
+
+  FailureOr<unsigned> inputIndex = getComputeInstanceInputIndex(representative, input);
+  if (failed(inputIndex))
+    return emitError(loc) << "cannot resolve scheduled batch step input index";
+
+  (void)peftClassPlan;
+  DeferredTransferMetadata transferMetadata;
+  transferMetadata.isBatched = true;
+  SmallVector<LaneProducerInfo> laneProducerInfos;
+  SmallVector<Value> uniqueOriginalSources;
+  SmallVector<int64_t> sourceOperandForScheduledLane;
+  for (const ComputeInstance &instance : stepTuple.instances) {
+    auto instanceInputs = getComputeInstanceInputs(instance);
+    if (*inputIndex >= instanceInputs.size())
+      return emitError(loc) << "scheduled batch step input index out of range";
+    Value laneInput = instanceInputs[*inputIndex];
+    std::optional<ProducerValueRef> laneProducer = getProducerValueRef(laneInput, &instance);
+    if (!laneProducer)
+      return emitError(loc) << "scheduled batch step mixes host and producer-derived inputs";
+    FailureOr<Value> laneSource = getOriginalProducerValue(*laneProducer);
+    if (failed(laneSource))
+      return emitError(loc) << "cannot resolve original graph producer value for " << getInstanceName(laneProducer->instance);
+
+    auto sourceIt = llvm::find(uniqueOriginalSources, *laneSource);
+    if (sourceIt == uniqueOriginalSources.end()) {
+      sourceOperandForScheduledLane.push_back(static_cast<int64_t>(uniqueOriginalSources.size()));
+      uniqueOriginalSources.push_back(*laneSource);
+    } else {
+      sourceOperandForScheduledLane.push_back(std::distance(uniqueOriginalSources.begin(), sourceIt));
+    }
+
+    laneProducerInfos.push_back(LaneProducerInfo {
+      *laneProducer,
+      *laneSource,
+      getDeferredEndpoint(laneProducer->instance, schedule),
+      getDeferredEndpoint(instance, schedule),
+    });
+  }
+
+  for (const LaneProducerInfo &laneProducerInfo : laneProducerInfos) {
+    transferMetadata.producers.push_back(laneProducerInfo.producerEndpoint);
+    transferMetadata.consumers.push_back(laneProducerInfo.consumerEndpoint);
+  }
+  transferMetadata.sourceOperandForScheduledLane = sourceOperandForScheduledLane;
+  transferMetadata.multiSourcePayload = uniqueOriginalSources.size() > 1;
+
+  return createScheduledLaneDeferredCommunication(
+    builder, loc, input, uniqueOriginalSources, *representativeProducer, transferMetadata, exchangeId++, mapped);
+}
+
+} // namespace spatial
+} // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.hpp

@@ -0,0 +1,33 @@
+#pragma once
+
+#include "ScheduledComputePlan.hpp"
+
+namespace onnx_mlir {
+namespace spatial {
+
+LogicalResult mapSingleCpuInput(OpBuilder &builder,
+                                Location loc,
+                                Value input,
+                                const ComputeInstance &consumerInstance,
+                                const MergeScheduleResult &schedule,
+                                ValueRange scheduledInputs,
+                                Block &block,
+                                unsigned firstInputArgument,
+                                ArrayRef<ProducerValueKey> carriedKeys,
+                                uint64_t &exchangeId,
+                                Value &mapped);
+
+LogicalResult mapMultiCpuTupleInput(OpBuilder &builder,
+                                    Location loc,
+                                    Value input,
+                                    const ComputeStepTuple &stepTuple,
+                                    const PeftClassPlan &peftClassPlan,
+                                    const MergeScheduleResult &schedule,
+                                    ValueRange scheduledInputs,
+                                    Block &block,
+                                    unsigned firstInputArgument,
+                                    uint64_t &exchangeId,
+                                    Value &mapped);
+
+} // namespace spatial
+} // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.hpp

@@ -1,17 +0,0 @@
-#pragma once
-
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Support/LogicalResult.h"
-
-#include "Scheduling/MergeSchedule.hpp"
-
-namespace onnx_mlir {
-namespace spatial {
-
-class MergeScheduleMaterializer {
-public:
-  mlir::LogicalResult run(mlir::func::FuncOp func, const MergeScheduleResult& schedule, int64_t& nextChannelId);
-};
-
-} // namespace spatial
-} // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp

@@ -1,359 +1,94 @@
-#include "mlir/Analysis/TopologicalSortUtils.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/IR/IRMapping.h"
-#include "mlir/IR/Location.h"
-#include "mlir/IR/PatternMatch.h"
-#include "mlir/IR/Region.h"
-#include "mlir/IR/Value.h"
-#include "mlir/IR/ValueRange.h"
+#include "ScheduledComputeMaterialization.hpp"
+#include "ScheduledComputeReport.hpp"
+#include "ScheduledComputeVerification.hpp"
+
 #include "mlir/Pass/Pass.h"
-#include "mlir/Support/LLVM.h"
 
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SmallSet.h"
-#include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/raw_os_ostream.h"
-
-#include <algorithm>
-#include <cstddef>
-#include <cstdint>
-#include <fstream>
-#include <memory>
-#include <optional>
-#include <utility>
-#include <vector>
-
-#include "MaterializeMergeSchedule.hpp"
-#include "Scheduling/ComputeGraph.hpp"
-#include "Scheduling/ComputeInstanceUtils.hpp"
 #include "Scheduling/MergeSchedulingAnalysis.hpp"
-#include "src/Accelerators/PIM/Common/IR/CompactAsmUtils.hpp"
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
-#include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Pass/PIMPasses.h"
 
 using namespace mlir;
 
 namespace onnx_mlir {
+namespace spatial {
 namespace {
-using namespace onnx_mlir::compact_asm;
-using SpatCompute = spatial::SpatCompute;
-using SpatComputeBatch = spatial::SpatComputeBatch;
-using spatial::getProducerValueRef;
 
-static std::optional<int32_t> getComputeCoreId(SpatCompute compute) {
-  if (auto coreIdAttr = compute->getAttrOfType<IntegerAttr>(onnx_mlir::kCoreIdAttrName)) {
-    auto checkedCoreId = pim::checkedI32(coreIdAttr.getInt(), compute, "merge compute core id");
-    if (failed(checkedCoreId))
-      return std::nullopt;
-    return *checkedCoreId;
-  }
-  return std::nullopt;
-}
-
-bool isTrivialSerialMergeCandidate(SpatCompute compute) {
-  if (!compute->hasOneUse())
-    return false;
-
-  auto& use = *compute->getUses().begin();
-  auto user = dyn_cast<SpatCompute>(use.getOwner());
-  return user && user.getInputs().size() == 1 && use.getOperandNumber() >= user.getWeights().size();
-}
-
-SmallVector<size_t> appendMissingWeightsAndBuildIndexMap(SmallVectorImpl<Value>& targetWeights,
-                                                         ValueRange sourceWeights) {
-  DenseMap<Value, SmallVector<size_t, 4>> targetWeightIndices;
-  for (auto [weightIndex, weight] : llvm::enumerate(targetWeights))
-    targetWeightIndices[weight].push_back(weightIndex);
-
-  DenseMap<Value, size_t> usedSourceWeightOccurrences;
-  SmallVector<size_t> sourceToTargetIndex;
-  sourceToTargetIndex.reserve(sourceWeights.size());
-  for (Value weight : sourceWeights) {
-    size_t occurrence = usedSourceWeightOccurrences[weight]++;
-    auto& matchingIndices = targetWeightIndices[weight];
-    if (occurrence >= matchingIndices.size()) {
-      size_t newIndex = targetWeights.size();
-      targetWeights.push_back(weight);
-      matchingIndices.push_back(newIndex);
-      sourceToTargetIndex.push_back(newIndex);
-      continue;
-    }
-    sourceToTargetIndex.push_back(matchingIndices[occurrence]);
-  }
-  return sourceToTargetIndex;
-}
-
-void mergeTriviallyConnectedComputes(func::FuncOp funcOp) {
-  Location loc = funcOp.getLoc();
-  IRRewriter rewriter(funcOp->getContext());
-  SmallVector<SpatCompute> trivialComputes;
-  llvm::SmallSet<SpatCompute, 8> toErase;
-
-  for (auto compute : funcOp.getOps<SpatCompute>())
-    if (isTrivialSerialMergeCandidate(compute))
-      trivialComputes.push_back(compute);
-
-  while (!trivialComputes.empty()) {
-    auto compute = trivialComputes.front();
-
-    if (compute.use_empty()) {
-      std::swap(trivialComputes.front(), trivialComputes.back());
-      trivialComputes.pop_back();
-      continue;
-    }
-
-    auto& computeUse = *compute->getUses().begin();
-    auto child = cast<SpatCompute>(computeUse.getOwner());
-    auto usedResult = cast<OpResult>(computeUse.get()).getResultNumber();
-    auto childInputIndex = computeUse.getOperandNumber() - child.getWeights().size();
-
-    rewriter.setInsertionPointAfter(compute.getOperation());
-    SmallVector<Value> mergedWeights(compute.getWeights().begin(), compute.getWeights().end());
-    SmallVector<size_t> childWeightToNewIndex = appendMissingWeightsAndBuildIndexMap(mergedWeights, child.getWeights());
-    SmallVector<Value> mergedInputs(compute.getInputs().begin(), compute.getInputs().end());
-    auto newCompute = SpatCompute::create(rewriter, loc, child.getResultTypes(), mergedWeights, mergedInputs);
-    Block* newBody = rewriter.createBlock(&newCompute.getBodyRegion());
-    for (Value weight : mergedWeights)
-      newBody->addArgument(weight.getType(), loc);
-    for (Value input : mergedInputs)
-      newBody->addArgument(input.getType(), loc);
-
-    IRMapping mapper;
-    for (auto [weightIndex, _] : llvm::enumerate(compute.getWeights())) {
-      auto oldWeightArg = compute.getWeightArgument(weightIndex);
-      auto newWeightArg = newCompute.getWeightArgument(weightIndex);
-      assert(oldWeightArg && newWeightArg && "expected compute weight block arguments");
-      mapper.map(*oldWeightArg, *newWeightArg);
-    }
-    for (auto [inputIndex, _] : llvm::enumerate(compute.getInputs())) {
-      auto oldInputArg = compute.getInputArgument(inputIndex);
-      auto newInputArg = newCompute.getInputArgument(inputIndex);
-      assert(oldInputArg && newInputArg && "expected compute input block arguments");
-      mapper.map(*oldInputArg, *newInputArg);
-    }
-    for (auto [oldIndex, weight] : llvm::enumerate(child.getWeights())) {
-      auto oldWeightArg = child.getWeightArgument(oldIndex);
-      auto newWeightArg = newCompute.getWeightArgument(childWeightToNewIndex[oldIndex]);
-      assert(oldWeightArg && newWeightArg && "expected child compute weight block arguments");
-      mapper.map(*oldWeightArg, *newWeightArg);
-    }
-
-    rewriter.setInsertionPointToEnd(newBody);
-    auto computeYield = cast<spatial::SpatYieldOp>(compute.getBody().front().getTerminator());
-    for (Operation& op : compute.getBody().front().without_terminator())
-      rewriter.clone(op, mapper);
-    auto childInputArg = child.getInputArgument(childInputIndex);
-    assert(childInputArg && "expected child compute input block argument");
-    mapper.map(*childInputArg, mapper.lookupOrDefault(computeYield.getOperand(usedResult)));
-
-    rewriter.setInsertionPointToEnd(newBody);
-    for (auto& op : child.getBody().front())
-      rewriter.clone(op, mapper);
-
-    child.replaceAllUsesWith(newCompute);
-    toErase.insert(child);
-
-    std::swap(trivialComputes.front(), trivialComputes.back());
-    trivialComputes.pop_back();
-    toErase.insert(compute);
-
-    if (isTrivialSerialMergeCandidate(newCompute))
-      trivialComputes.push_back(newCompute);
-  }
-
-  for (auto compute : toErase) {
-    for (Value result : compute->getResults())
-      result.dropAllUses();
-    compute.erase();
-  }
-}
-
-void generateReport(func::FuncOp funcOp, const std::string& name, size_t usedCpuCount = 0) {
-  std::fstream file = openReportFile(name);
-  if (!file.is_open())
-    return;
-  llvm::raw_os_ostream os(file);
-
-  struct ReportRow {
-    uint64_t id = 0;
-    uint64_t logicalComputeCount = 0;
-    uint64_t crossbarCount = 0;
-    uint64_t instructionCount = 0;
-    bool isRebatched = false;
-    SmallVector<int32_t> coreIds;
-  };
-
-  uint64_t totalComputeOps = 0;
-  uint64_t totalLogicalComputes = 0;
-  uint64_t totalBatchComputeOps = 0;
-  uint64_t totalInstructionCount = 0;
-  uint64_t totalCrossbarCount = 0;
-  uint64_t nextBatchId = 0;
-  std::vector<ReportRow> collectedData;
-
-  auto getPerInstanceCrossbarCount = [&](Operation* op) -> uint64_t {
-    return static_cast<uint64_t>(spatial::collectDistinctCrossbarWeights(op).size());
-  };
-
-  for (Operation& op : funcOp.getBody().front()) {
-    if (auto spatCompute = dyn_cast<SpatCompute>(&op)) {
-      uint64_t numInst = spatial::countComputeBodyInstructions(spatCompute.getBody());
-      uint64_t perInstanceCrossbarCount = getPerInstanceCrossbarCount(spatCompute.getOperation());
-      SmallVector<int32_t> coreIds;
-      if (auto coreId = getComputeCoreId(spatCompute))
-        coreIds.push_back(*coreId);
-      collectedData.push_back({totalComputeOps++, 1, perInstanceCrossbarCount, numInst, false, coreIds});
-      totalLogicalComputes += 1;
-      totalInstructionCount += numInst;
-      totalCrossbarCount += perInstanceCrossbarCount;
-      continue;
-    }
-    if (auto batch = dyn_cast<SpatComputeBatch>(&op)) {
-      uint64_t numInst = spatial::countComputeBodyInstructions(batch.getBody());
-      uint64_t logicalCount = static_cast<uint64_t>(batch.getLaneCount());
-      uint64_t perInstanceCrossbarCount = getPerInstanceCrossbarCount(batch.getOperation());
-      SmallVector<int32_t> coreIds;
-      if (auto coreIdsAttr = batch->getAttrOfType<DenseI32ArrayAttr>(onnx_mlir::kCoreIdsAttrName))
-        llvm::append_range(coreIds, coreIdsAttr.asArrayRef());
-      collectedData.push_back(
-        {nextBatchId++, logicalCount, perInstanceCrossbarCount * logicalCount, numInst, true, coreIds});
-      totalComputeOps += 1;
-      totalLogicalComputes += logicalCount;
-      totalBatchComputeOps += 1;
-      totalInstructionCount += numInst * logicalCount;
-      totalCrossbarCount += perInstanceCrossbarCount * logicalCount;
-    }
-  }
-
-  llvm::SmallVector<ReportField, 6> totalFields = {
-    {"Used cores",                            std::to_string(usedCpuCount)         },
-    {"Number of top-level compute ops",       std::to_string(totalComputeOps)      },
-    {"Number of logical computes",            std::to_string(totalLogicalComputes) },
-    {"Number of top-level batch compute ops", std::to_string(totalBatchComputeOps) },
-    {"Number of instructions",                std::to_string(totalInstructionCount)},
-    {"Number of used crossbars",              std::to_string(totalCrossbarCount)   }
-  };
-  printReportTotalsBlock(os, totalFields);
-  if (!collectedData.empty())
-    os << "\n";
-
-  sortReportEntriesByFirstCore(collectedData);
-
-  for (uint64_t cI = 0; cI < totalComputeOps; ++cI) {
-    uint64_t lastIndex = cI;
-    ReportRow current = collectedData[cI];
-
-    for (uint64_t nI = cI + 1; nI < totalComputeOps; ++nI) {
-      ReportRow next = collectedData[nI];
-      if (current.isRebatched == next.isRebatched && current.crossbarCount == next.crossbarCount
-          && current.instructionCount == next.instructionCount
-          && current.logicalComputeCount == next.logicalComputeCount)
-        lastIndex = nI;
-      else
-        break;
-    }
-
-    if (current.isRebatched) {
-      os << "Batch ";
-      for (uint64_t index = cI; index <= lastIndex; ++index) {
-        if (index != cI)
-          os << ",\n      ";
-        os << collectedData[index].id << " (cores ";
-        if (collectedData[index].coreIds.empty())
-          os << "unknown";
-        else
-          printCompressedIntegerEntries(os, ArrayRef<int32_t>(collectedData[index].coreIds));
-        os << ")";
-      }
-    }
-    else {
-      os << "Compute ";
-      SmallVector<uint64_t> opIds;
-      opIds.reserve(lastIndex - cI + 1);
-      for (uint64_t index = cI; index <= lastIndex; ++index)
-        opIds.push_back(collectedData[index].id);
-      printCompressedIntegerEntries(os, ArrayRef<uint64_t>(opIds));
-    }
-
-    os << ":\n";
-    uint64_t perCoreLogicalComputeCount = current.isRebatched ? 1 : current.logicalComputeCount;
-    uint64_t perCoreInstructionCount = current.instructionCount;
-    uint64_t perCoreCrossbarCount =
-      current.logicalComputeCount == 0 ? 0 : current.crossbarCount / current.logicalComputeCount;
-    uint64_t totalEntryInstructionCount = current.instructionCount * current.logicalComputeCount;
-
-    llvm::SmallVector<ReportField, 3> perCoreFields = {
-      {"Number of logical computes", std::to_string(perCoreLogicalComputeCount)},
-      {"Number of instructions",     std::to_string(perCoreInstructionCount)   },
-      {"Number of used crossbars",   std::to_string(perCoreCrossbarCount)      }
-    };
-    if (current.isRebatched) {
-      llvm::SmallVector<ReportField, 3> totalEntryFields = {
-        {"Number of logical computes", std::to_string(current.logicalComputeCount)},
-        {"Number of instructions",     std::to_string(totalEntryInstructionCount) },
-        {"Number of used crossbars",   std::to_string(current.crossbarCount)      }
-      };
-      printReportPerCoreAndTotalFields(os, perCoreFields, totalEntryFields);
-    }
-    else {
-      printReportFlatFields(os, perCoreFields);
-    }
-    printReportEntrySeparator(os, lastIndex + 1 < totalComputeOps);
-    cI = lastIndex;
-  }
-
-  os.flush();
-  file.close();
-}
-
-struct MergeComputeNodesPass : PassWrapper<MergeComputeNodesPass, OperationPass<func::FuncOp>> {
-
-private:
-  int64_t nextChannelId = 0;
-
-public:
+struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MergeComputeNodesPass)
 
-  StringRef getArgument() const override { return "pim-merge-compute-nodes-pass"; }
+  StringRef getArgument() const override { return "pim-merge-compute-nodes"; }
   StringRef getDescription() const override {
-    return "Merge Spatial-Compute-Nodes in order to reduce the total "
-           "execution time";
+    return "Materialize scheduled Spatial compute with deferred communication placeholders.";
   }
 
-  LogicalResult initialize(MLIRContext* context) override { return success(); }
-
   void runOnOperation() override {
-    func::FuncOp func = getOperation();
-    mergeTriviallyConnectedComputes(func);
-
-    const spatial::MergeScheduleResult* analysisResult = nullptr;
-    analysisResult = &getAnalysis<spatial::MergeSchedulingAnalysis>().getResult();
-    if (failed(spatial::MergeScheduleMaterializer().run(func, *analysisResult, nextChannelId))) {
+    ModuleOp moduleOp = getOperation();
+    auto entryFunc = getPimEntryFunc(moduleOp);
+    if (failed(entryFunc)) {
+      moduleOp.emitError("failed to locate the PIM entry function during MergeComputeNodes");
       signalPassFailure();
       return;
     }
 
-    if (!sortTopologically(&func.getBody().front())) {
-      func.emitOpError("failed to topologically order merged Spatial IR");
+    func::FuncOp funcOp = *entryFunc;
+    MergeScheduleResult schedule = MergeSchedulingAnalysis(funcOp).getResult();
+    PatternRewriter rewriter(moduleOp.getContext());
+    FailureOr<ScheduledComputeMaterializationResult> materialization =
+      materializeScheduledCompute(funcOp, schedule, rewriter);
+    if (failed(materialization)) {
       signalPassFailure();
       return;
     }
-    if (failed(verifySpatialCommunicationInvariants(func))) {
-      func.emitOpError("merged Spatial communication invariant verification failed");
+
+    if (failed(verifyMaterializedScheduleMapping(funcOp,
+                                                 schedule,
+                                                 materialization->peftClassPlans,
+                                                 materialization->graphComputeToBlockMap,
+                                                 materialization->materializedSchedules))) {
+      moduleOp.emitError("scheduled Spatial materialization mapping verification failed");
       signalPassFailure();
       return;
     }
-    dumpModule(cast<ModuleOp>(func->getParentOp()), "spatial1_merged");
-    generateReport(func, "spatial_merge_report", analysisResult->cpuToLastComputeMap.size());
+    if (failed(verifyDeferredTransferPhase1Invariants(funcOp))) {
+      moduleOp.emitError("scheduled Spatial deferred communication verification failed");
+      signalPassFailure();
+      return;
+    }
+    if (failed(verifyScheduledMaterializationRecords(materialization->materializedSchedules))) {
+      moduleOp.emitError("scheduled Spatial materialization record verification failed");
+      signalPassFailure();
+      return;
+    }
+    if (failed(verifyPeftMaterializationReportSummary(funcOp,
+                                                      schedule,
+                                                      materialization->peftClassPlans,
+                                                      materialization->materializedSchedules))) {
+      moduleOp.emitError("scheduled Spatial report verification failed");
+      signalPassFailure();
+      return;
+    }
+    if (failed(verifyScheduledSpatialInvariants(funcOp))) {
+      moduleOp.emitError("scheduled Spatial phase 1 verification failed");
+      signalPassFailure();
+      return;
+    }
+
+    dumpScheduledComputeReportAndModule(moduleOp,
+                                        funcOp,
+                                        schedule,
+                                        materialization->peftClassPlans,
+                                        materialization->materializedSchedules);
+    moduleOp.emitError("MergeComputeNodes stopped after spatial1_scheduled_no_comm; "
+                       "Phase 2 communication realization is not implemented");
+    signalPassFailure();
   }
 };
 
 } // namespace
+} // namespace spatial
 
-std::unique_ptr<Pass> createMergeComputeNodesPass() { return std::make_unique<MergeComputeNodesPass>(); }
+std::unique_ptr<Pass> createMergeComputeNodesPass() { return std::make_unique<spatial::MergeComputeNodesPass>(); }
 
 } // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp

@@ -0,0 +1,954 @@
+#include "ScheduledComputeMaterialization.hpp"
+#include "DeferredCommunicationPlanning.hpp"
+
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+
+#include <map>
+
+namespace onnx_mlir {
+namespace spatial {
+
+using namespace mlir;
+namespace {
+
+struct BatchFragmentSpec {
+  RankedTensorType resultType;
+  RankedTensorType sourceSliceType;
+  unsigned laneDim = 0;
+  SmallVector<OpFoldResult> offsets;
+  SmallVector<OpFoldResult> sizes;
+  SmallVector<OpFoldResult> strides;
+};
+
+static SmallVector<OpFoldResult> remapMixedOffsets(ArrayRef<OpFoldResult> mixedOffsets, IRMapping &mapper) {
+  SmallVector<OpFoldResult> remapped;
+  remapped.reserve(mixedOffsets.size());
+  for (OpFoldResult ofr : mixedOffsets) {
+    if (auto value = dyn_cast<Value>(ofr))
+      remapped.push_back(cast<Value>(mapper.lookupOrDefault(value)));
+    else
+      remapped.push_back(cast<Attribute>(ofr));
+  }
+  return remapped;
+}
+
+static void appendUnique(SmallVectorImpl<Value> &values, Value value) {
+  if (!llvm::is_contained(values, value))
+    values.push_back(value);
+}
+
+static Value getBlockOperand(Block &block, ValueRange operands, Value value, unsigned firstArgument = 0) {
+  auto it = llvm::find(operands, value);
+  assert(it != operands.end() && "missing scheduled operand");
+  return block.getArgument(firstArgument + std::distance(operands.begin(), it));
+}
+
+static Value getScheduledComputeOutputArgument(Block &block, ValueRange scheduledWeights, ValueRange scheduledInputs,
+                                               ArrayRef<ProducerValueKey> carriedKeys, ProducerValueKey key) {
+  unsigned base = scheduledWeights.size() + scheduledInputs.size();
+  auto it = llvm::find(carriedKeys, key);
+  assert(it != carriedKeys.end() && "missing carried output");
+  return block.getArgument(base + std::distance(carriedKeys.begin(), it));
+}
+
+static unsigned getScheduledComputeResultArgBase(SpatScheduledCompute scheduled) {
+  return scheduled.getWeights().size() + scheduled.getInputs().size();
+}
+
+static void appendComputeBlockArguments(SmallVectorImpl<Type> &argTypes,
+                                        SmallVectorImpl<Location> &argLocs,
+                                        ValueRange weights,
+                                        ValueRange inputs,
+                                        ArrayRef<ProducerValueKey> carriedKeys,
+                                        Location loc) {
+  for (Value weight : weights)
+    argTypes.push_back(weight.getType());
+  for (Value input : inputs)
+    argTypes.push_back(input.getType());
+  for (ProducerValueKey key : carriedKeys) {
+    auto outputs = getComputeInstanceOutputValues(key.instance);
+    assert(key.resultIndex < outputs.size() && "missing carried result");
+    argTypes.push_back(outputs[key.resultIndex].getType());
+  }
+  argLocs.append(argTypes.size(), loc);
+}
+
+static Block *createScheduledComputeBlock(PatternRewriter &rewriter,
+                                          SpatScheduledCompute scheduled,
+                                          ArrayRef<ProducerValueKey> carriedKeys,
+                                          Location loc) {
+  SmallVector<Type> argTypes;
+  SmallVector<Location> argLocs;
+  appendComputeBlockArguments(argTypes, argLocs, scheduled.getWeights(), scheduled.getInputs(), carriedKeys, loc);
+  return rewriter.createBlock(&scheduled.getBody(), scheduled.getBody().end(), TypeRange(argTypes), argLocs);
+}
+
+static void appendBlockYieldBaseAndCarriedOperands(Block &block,
+                                                   unsigned baseArgCount,
+                                                   size_t carriedCount,
+                                                   SmallVectorImpl<Value> &operands) {
+  for (unsigned index = 0; index < baseArgCount; ++index)
+    operands.push_back(block.getArgument(index));
+  for (size_t index = 0; index < carriedCount; ++index)
+    operands.push_back(block.getArgument(baseArgCount + index));
+}
+
+static void createBlockYield(PatternRewriter &rewriter, Location loc, ValueRange outputs, Block *next = nullptr) {
+  OperationState state(loc, SpatBlockYieldOp::getOperationName());
+  state.addOperands(outputs);
+  if (next)
+    state.addSuccessors(next);
+  rewriter.create(state);
+}
+
+static FailureOr<BatchFragmentSpec> getBatchFragmentSpec(SpatComputeBatch batch,
+                                                         unsigned resultIndex,
+                                                         uint32_t fragmentLaneCount) {
+  auto inParallel = dyn_cast<SpatInParallelOp>(batch.getBody().front().getTerminator());
+  if (!inParallel)
+    return batch.emitOpError("scheduled materialization only supports resultful spat.graph_compute_batch");
+
+  auto outputArg = batch.getOutputArgument(resultIndex);
+  if (!outputArg)
+    return batch.emitOpError("scheduled materialization could not locate batch output block argument");
+
+  for (Operation &op : inParallel.getRegion().front()) {
+    auto insert = dyn_cast<tensor::ParallelInsertSliceOp>(&op);
+    if (!insert)
+      return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+    if (insert.getDest() != *outputArg)
+      continue;
+
+    RankedTensorType destType = insert.getDestType();
+    RankedTensorType sourceType = insert.getSourceType();
+    if (!destType || !sourceType || !destType.hasStaticShape() || !sourceType.hasStaticShape())
+      return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+    if (destType.getRank() == 0 || sourceType.getRank() == 0 || destType.getRank() != sourceType.getRank())
+      return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+    if (!insert.hasUnitStride())
+      return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+
+    auto offsets = insert.getMixedOffsets();
+    auto sizes = insert.getMixedSizes();
+    auto strides = insert.getMixedStrides();
+    if (offsets.size() != static_cast<size_t>(destType.getRank()) || sizes.size() != static_cast<size_t>(destType.getRank())
+        || strides.size() != static_cast<size_t>(destType.getRank()))
+      return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+
+    std::optional<unsigned> laneDim;
+    for (unsigned dim = 0; dim < offsets.size(); ++dim) {
+      bool dimLooksLaneLike = sourceType.getShape()[dim] != destType.getShape()[dim];
+      if (auto value = dyn_cast<Value>(offsets[dim]))
+        dimLooksLaneLike = dimLooksLaneLike || valueTransitivelyDependsOn(value, *batch.getLaneArgument());
+      if (dimLooksLaneLike) {
+        if (laneDim && *laneDim != dim)
+          return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+        laneDim = dim;
+      }
+    }
+    if (!laneDim)
+      return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+
+    SmallVector<int64_t> shape(destType.getShape().begin(), destType.getShape().end());
+    shape[*laneDim] = fragmentLaneCount;
+    return BatchFragmentSpec {
+      RankedTensorType::get(shape, destType.getElementType()),
+      sourceType,
+      *laneDim,
+      SmallVector<OpFoldResult>(offsets.begin(), offsets.end()),
+      SmallVector<OpFoldResult>(sizes.begin(), sizes.end()),
+      SmallVector<OpFoldResult>(strides.begin(), strides.end())
+    };
+  }
+
+  return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+}
+
+
+static SourceLaneSelector buildSourceLaneSelector(PatternRewriter &rewriter,
+                                                  const ComputeStepTuple &stepTuple,
+                                                  Operation *constantAnchor,
+                                                  std::map<std::vector<uint32_t>, Value> &laneStartTableCache) {
+  if (std::optional<SourceLaneAffineMapping> affineMapping = getSourceLaneAffineMapping(stepTuple)) {
+    SourceLaneSelector selector;
+    selector.kind = SourceLaneSelector::Kind::Affine;
+    selector.affine = *affineMapping;
+    return selector;
+  }
+
+  SmallVector<uint32_t> tableValues = collectSourceLaneStarts(stepTuple);
+  std::vector<uint32_t> cacheKey(tableValues.begin(), tableValues.end());
+  auto cacheIt = laneStartTableCache.find(cacheKey);
+  if (cacheIt != laneStartTableCache.end()) {
+    SourceLaneSelector selector;
+    selector.kind = SourceLaneSelector::Kind::Table;
+    selector.table = cacheIt->second;
+    selector.tableValues = tableValues;
+    return selector;
+  }
+
+  SmallVector<int64_t> tableValuesI64;
+  tableValuesI64.reserve(tableValues.size());
+  for (uint32_t value : tableValues)
+    tableValuesI64.push_back(value);
+  Value table = createI64LookupTableConstant(rewriter, constantAnchor, tableValuesI64);
+  laneStartTableCache.emplace(std::move(cacheKey), table);
+  SourceLaneSelector selector;
+  selector.kind = SourceLaneSelector::Kind::Table;
+  selector.table = table;
+  selector.tableValues = tableValues;
+  return selector;
+}
+
+static FailureOr<Value> buildSourceLaneStartForScheduledLane(OpBuilder &builder,
+                                                             Location loc,
+                                                             Value scheduledLane,
+                                                             const SourceLaneSelector &selector,
+                                                             Operation *constantAnchor) {
+  (void)constantAnchor;
+  if (selector.kind == SourceLaneSelector::Kind::Affine) {
+    if (selector.affine.baseLaneStart == 0 && selector.affine.laneCount == 1)
+      return scheduledLane;
+    AffineExpr d0 = builder.getAffineDimExpr(0);
+    AffineExpr expr = d0;
+    if (selector.affine.laneCount != 1)
+      expr = d0 * selector.affine.laneCount;
+    if (selector.affine.baseLaneStart != 0)
+      expr = expr + selector.affine.baseLaneStart;
+    return affine::AffineApplyOp::create(
+             builder, loc, AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, expr), ValueRange {scheduledLane})
+      .getResult();
+  }
+
+  if (!selector.table)
+    return failure();
+  Value sourceLaneStartI64 =
+    tensor::ExtractOp::create(builder, loc, selector.table, ValueRange {scheduledLane}).getResult();
+  return arith::IndexCastOp::create(builder, loc, builder.getIndexType(), sourceLaneStartI64).getResult();
+}
+
+static LogicalResult verifyPeftClassPlan(Operation *diagnosticAnchor,
+                                         const PeftClassPlan &peftClassPlan,
+                                         const MergeScheduleResult &schedule) {
+  if (peftClassPlan.cpus.empty())
+    return diagnosticAnchor->emitOpError("PEFT materialization class has no CPUs");
+
+  SmallVector<const SmallVector<ComputeInstance> *> schedules;
+  for (size_t cpu : peftClassPlan.cpus) {
+    auto it = peftClassPlan.instancesByCpu.find(cpu);
+    if (it == peftClassPlan.instancesByCpu.end())
+      return diagnosticAnchor->emitOpError("PEFT materialization class is missing a per-CPU schedule");
+    schedules.push_back(&it->second);
+    for (const ComputeInstance &instance : it->second)
+      if (!schedule.computeToCpuSlotMap.count(instance))
+        return diagnosticAnchor->emitOpError("PEFT materialization class references a compute instance without a scheduler position");
+  }
+
+  if (peftClassPlan.cpus.size() == 1)
+    return success();
+
+  auto emitNonIso = [&](size_t stepPosition) -> LogicalResult {
+    std::string cpus;
+    llvm::raw_string_ostream os(cpus);
+    llvm::interleaveComma(peftClassPlan.cpus, os, [&](size_t cpu) { os << cpu; });
+    diagnosticAnchor->emitOpError("PEFT equivalence class has non-isomorphic per-CPU schedules")
+      << " class " << peftClassPlan.canonicalClassId << " cpus [" << os.str() << "] step " << stepPosition;
+    return failure();
+  };
+
+  size_t tupleCount = schedules.front()->size();
+  for (const SmallVector<ComputeInstance> *cpuSchedule : schedules)
+    if (cpuSchedule->size() != tupleCount)
+      return emitNonIso(0);
+
+  for (size_t stepPosition = 0; stepPosition < tupleCount; ++stepPosition) {
+    const ComputeInstance &reference = (*schedules.front())[stepPosition];
+    bool refIsScalar = isa<SpatCompute>(reference.op);
+    for (size_t cpuIndex = 1; cpuIndex < schedules.size(); ++cpuIndex) {
+      const ComputeInstance &instance = (*schedules[cpuIndex])[stepPosition];
+      if (instance.op != reference.op || instance.laneCount != reference.laneCount)
+        return emitNonIso(stepPosition);
+      if (isa<SpatCompute>(instance.op) != refIsScalar)
+        return emitNonIso(stepPosition);
+    }
+  }
+
+  return success();
+}
+
+static LogicalResult collectPeftClassOperandsAndResults(PeftClassPlan &peftClassPlan,
+                                                        const MergeScheduleResult &schedule) {
+  peftClassPlan.weights.clear();
+  peftClassPlan.inputs.clear();
+  peftClassPlan.resultTypes.clear();
+
+  if (peftClassPlan.cpus.size() == 1) {
+    size_t cpu = peftClassPlan.cpus.front();
+    for (const ComputeInstance &instance : peftClassPlan.instancesByCpu.lookup(cpu)) {
+      if (auto compute = dyn_cast<SpatCompute>(instance.op)) {
+        llvm::append_range(peftClassPlan.resultTypes, compute.getResultTypes());
+      } else {
+        auto batch = cast<SpatComputeBatch>(instance.op);
+        for (unsigned resultIndex = 0; resultIndex < batch.getNumResults(); ++resultIndex) {
+          auto spec = getBatchFragmentSpec(batch, resultIndex, instance.laneCount);
+          if (failed(spec))
+            return failure();
+          peftClassPlan.resultTypes.push_back(spec->resultType);
+        }
+      }
+
+      for (Value weight : getComputeInstanceWeights(instance))
+        appendUnique(peftClassPlan.weights, weight);
+      for (Value input : getComputeInstanceInputs(instance))
+        if (!getProducerValueRef(input, &instance))
+          appendUnique(peftClassPlan.inputs, input);
+    }
+    return success();
+  }
+
+  for (const ScheduledStepPlan &stepPlan : buildScheduledStepPlans(peftClassPlan)) {
+    const ComputeStepTuple &stepTuple = stepPlan.stepTuple;
+    const ComputeInstance &representative = stepTuple.instances.front();
+    if (auto compute = dyn_cast<SpatCompute>(representative.op)) {
+      for (Type type : compute.getResultTypes()) {
+        auto tensorType = dyn_cast<RankedTensorType>(type);
+        if (!tensorType || !tensorType.hasStaticShape())
+          return compute.emitOpError("scheduled materialization only supports static ranked tensor scalar results");
+        SmallVector<int64_t> shape;
+        shape.push_back(static_cast<int64_t>(peftClassPlan.cpus.size()));
+        llvm::append_range(shape, tensorType.getShape());
+        peftClassPlan.resultTypes.push_back(RankedTensorType::get(shape, tensorType.getElementType()));
+      }
+    } else {
+      auto batch = cast<SpatComputeBatch>(representative.op);
+      uint32_t totalLanes = static_cast<uint32_t>(peftClassPlan.cpus.size()) * representative.laneCount;
+      for (unsigned resultIndex = 0; resultIndex < batch.getNumResults(); ++resultIndex) {
+        auto spec = getBatchFragmentSpec(batch, resultIndex, totalLanes);
+        if (failed(spec))
+          return failure();
+        peftClassPlan.resultTypes.push_back(spec->resultType);
+      }
+    }
+
+    for (const ComputeInstance &instance : stepTuple.instances) {
+      for (Value weight : getComputeInstanceWeights(instance))
+        appendUnique(peftClassPlan.weights, weight);
+      for (Value input : getComputeInstanceInputs(instance))
+        if (!getProducerValueRef(input, &instance))
+          appendUnique(peftClassPlan.inputs, input);
+    }
+  }
+
+  return success();
+}
+
+static void cloneComputeBody(OpBuilder &builder, Block &source, IRMapping &mapper, SmallVectorImpl<Value> &yieldedValues) {
+  for (Operation &op : source.without_terminator())
+    builder.clone(op, mapper);
+  auto yield = cast<SpatYieldOp>(source.getTerminator());
+  for (Value output : yield.getOutputs())
+    yieldedValues.push_back(mapper.lookup(output));
+}
+
+static LogicalResult materializeResultfulBatchChunkAsScalar(PatternRewriter &rewriter,
+                                                            SpatComputeBatch batch,
+                                                            const ComputeInstance &instance,
+                                                            ValueRange scheduledWeights,
+                                                            ValueRange scheduledInputs,
+                                                            Block &block,
+                                                            const MergeScheduleResult &schedule,
+                                                            uint64_t &exchangeId,
+                                                            SmallVectorImpl<Value> &yieldedValues) {
+  SmallVector<Value> initResults;
+  SmallVector<BatchFragmentSpec> fragmentSpecs;
+  for (unsigned resultIndex = 0; resultIndex < batch.getNumResults(); ++resultIndex) {
+    auto spec = getBatchFragmentSpec(batch, resultIndex, instance.laneCount);
+    if (failed(spec))
+      return failure();
+    fragmentSpecs.push_back(*spec);
+    auto empty = createEmptyTensorForType(rewriter, batch.getLoc(), spec->resultType);
+    if (failed(empty))
+      return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+    initResults.push_back(*empty);
+  }
+
+  Value lower = getOrCreateIndexConstant(rewriter, batch.getOperation(), instance.laneStart);
+  Value upper = getOrCreateIndexConstant(rewriter, batch.getOperation(), instance.laneStart + instance.laneCount);
+  Value step = getOrCreateIndexConstant(rewriter, batch.getOperation(), 1);
+  auto loop = buildNormalizedScfFor(
+    rewriter,
+    batch.getLoc(),
+    lower,
+    upper,
+    step,
+    initResults,
+    [&](OpBuilder &builder, Location bodyLoc, Value originalLane, ValueRange iterArgs, SmallVectorImpl<Value> &yielded) -> LogicalResult {
+
+      IRMapping mapper;
+      mapper.map(*batch.getLaneArgument(), originalLane);
+      for (auto [index, weight] : llvm::enumerate(batch.getWeights()))
+        mapper.map(*batch.getWeightArgument(index), getBlockOperand(block, scheduledWeights, weight));
+      for (auto [index, input] : llvm::enumerate(batch.getInputs())) {
+        Value mapped;
+        if (failed(mapSingleCpuInput(builder,
+                                     input.getLoc(),
+                                     input,
+                                     instance,
+                                     schedule,
+                                     scheduledInputs,
+                                     block,
+                                     scheduledWeights.size(),
+                                     ArrayRef<ProducerValueKey> {},
+                                     exchangeId,
+                                     mapped)))
+          return failure();
+        mapper.map(*batch.getInputArgument(index), mapped);
+      }
+      for (auto [index, outputArg] : llvm::enumerate(batch.getOutputs()))
+        (void)outputArg, mapper.map(*batch.getOutputArgument(index), iterArgs[index]);
+
+      Block &source = batch.getBody().front();
+      for (Operation &op : source.without_terminator())
+        builder.clone(op, mapper);
+
+      auto inParallel = dyn_cast<SpatInParallelOp>(source.getTerminator());
+      if (!inParallel)
+          return batch.emitOpError("expected spat.in_parallel in resultful spat.graph_compute_batch"), failure();
+      Value localLane = arith::SubIOp::create(builder,
+                                              bodyLoc,
+                                              originalLane,
+                                              getOrCreateIndexConstant(rewriter, batch.getOperation(), instance.laneStart))
+                          .getResult();
+      DenseMap<BlockArgument, size_t> outputIndexByArg;
+      for (size_t index = 0; index < batch.getNumResults(); ++index)
+        outputIndexByArg[*batch.getOutputArgument(index)] = index;
+
+      SmallVector<Value> current(iterArgs.begin(), iterArgs.end());
+      for (Operation &op : inParallel.getRegion().front()) {
+        auto insert = dyn_cast<tensor::ParallelInsertSliceOp>(&op);
+        if (!insert)
+          return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+        auto oldDest = dyn_cast<BlockArgument>(insert.getDest());
+        if (!oldDest || !outputIndexByArg.count(oldDest))
+          return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"), failure();
+        size_t resultIndex = outputIndexByArg.lookup(oldDest);
+        SmallVector<OpFoldResult> offsets = remapMixedOffsets(insert.getMixedOffsets(), mapper);
+        offsets[fragmentSpecs[resultIndex].laneDim] = localLane;
+        current[resultIndex] = tensor::InsertSliceOp::create(builder,
+                                                             insert.getLoc(),
+                                                             mapper.lookup(insert.getSource()),
+                                                             current[resultIndex],
+                                                             offsets,
+                                                             remapMixedOffsets(insert.getMixedSizes(), mapper),
+                                                             remapMixedOffsets(insert.getMixedStrides(), mapper))
+                                 .getResult();
+      }
+      llvm::append_range(yielded, current);
+      return success();
+    });
+  if (failed(loop))
+    return failure();
+  llvm::append_range(yieldedValues, loop->results);
+  return success();
+}
+
+static LogicalResult materializeSingleCpuPeftClass(
+  PatternRewriter &rewriter,
+  SpatScheduledCompute scheduled,
+  const PeftClassPlan &peftClassPlan,
+  const MergeScheduleResult &schedule,
+  DenseMap<GraphComputeBlockKey, Block *> &graphComputeToBlockMap,
+  ScheduledMaterializationRecord &record,
+  uint64_t &exchangeId) {
+  size_t cpu = peftClassPlan.cpus.front();
+  auto instancesIt = peftClassPlan.instancesByCpu.find(cpu);
+  assert(instancesIt != peftClassPlan.instancesByCpu.end() && "missing single-cpu schedule");
+  const SmallVector<ComputeInstance> &instances = instancesIt->second;
+
+  SmallVector<ProducerValueKey> carriedKeys;
+  Block *block = nullptr;
+  for (auto [ordinal, instance] : llvm::enumerate(instances)) {
+    if (!block)
+      block = createScheduledComputeBlock(rewriter, scheduled, carriedKeys, instance.op->getLoc());
+
+    GraphComputeBlockKey key = getGraphComputeBlockKey(instance);
+    graphComputeToBlockMap[key] = block;
+    record.computeKeys.push_back(key);
+    record.blocks.push_back(block);
+
+    rewriter.setInsertionPointToStart(block);
+    SmallVector<Value> yieldedValues;
+    if (auto compute = dyn_cast<SpatCompute>(instance.op)) {
+      IRMapping mapper;
+      for (auto [index, weight] : llvm::enumerate(compute.getWeights()))
+        mapper.map(*compute.getWeightArgument(index), getBlockOperand(*block, scheduled.getWeights(), weight));
+      for (auto [index, input] : llvm::enumerate(compute.getInputs())) {
+        Value mapped;
+        if (failed(mapSingleCpuInput(rewriter,
+                                     input.getLoc(),
+                                     input,
+                                     instance,
+                                     schedule,
+                                     scheduled.getInputs(),
+                                     *block,
+                                     scheduled.getWeights().size(),
+                                     carriedKeys,
+                                     exchangeId,
+                                     mapped)))
+          return failure();
+        mapper.map(*compute.getInputArgument(index), mapped);
+      }
+      cloneComputeBody(rewriter, compute.getBody().front(), mapper, yieldedValues);
+    } else {
+      auto batch = cast<SpatComputeBatch>(instance.op);
+      if (failed(materializeResultfulBatchChunkAsScalar(rewriter,
+                                                        batch,
+                                                        instance,
+                                                        scheduled.getWeights(),
+                                                        scheduled.getInputs(),
+                                                        *block,
+                                                        schedule,
+                                                        exchangeId,
+                                                        yieldedValues)))
+        return failure();
+    }
+
+    SmallVector<ProducerValueKey> currentKeys;
+    for (size_t index = 0; index < yieldedValues.size(); ++index)
+      currentKeys.push_back({instance, index});
+    unsigned baseArgCount = getScheduledComputeResultArgBase(scheduled);
+    SmallVector<Value> blockYieldOperands;
+    bool hasNextBlock = ordinal + 1 < instances.size();
+    if (hasNextBlock) {
+      SmallVector<ProducerValueKey> nextCarriedKeys(carriedKeys);
+      llvm::append_range(nextCarriedKeys, currentKeys);
+      Block *nextBlock = createScheduledComputeBlock(rewriter, scheduled, nextCarriedKeys, instance.op->getLoc());
+      appendBlockYieldBaseAndCarriedOperands(*block, baseArgCount, carriedKeys.size(), blockYieldOperands);
+      llvm::append_range(blockYieldOperands, yieldedValues);
+      rewriter.setInsertionPointToEnd(block);
+      createBlockYield(rewriter, instance.op->getLoc(), blockYieldOperands, nextBlock);
+      carriedKeys = std::move(nextCarriedKeys);
+      block = nextBlock;
+    } else {
+      for (ProducerValueKey carried : carriedKeys)
+        blockYieldOperands.push_back(getScheduledComputeOutputArgument(*block,
+                                                                      scheduled.getWeights(),
+                                                                      scheduled.getInputs(),
+                                                                      carriedKeys,
+                                                                      carried));
+      llvm::append_range(blockYieldOperands, yieldedValues);
+      rewriter.setInsertionPointToEnd(block);
+      createBlockYield(rewriter, instance.op->getLoc(), blockYieldOperands);
+    }
+  }
+  return success();
+}
+
+// Builds offsets for inserting one per-CPU local fragment into the
+// scheduled_compute_batch output. The lane offset is in scheduled-output
+// lane space, not local fragment lane space.
+static SmallVector<OpFoldResult> buildScheduledOutputInsertOffsets(OpBuilder &builder,
+                                                                   Location loc,
+                                                                   Value scheduledLane,
+                                                                   int64_t lanesPerScheduledLane,
+                                                                   RankedTensorType localFragmentType,
+                                                                   unsigned laneDim) {
+  SmallVector<OpFoldResult> offsets;
+  Value scheduledOutputLane = scheduledLane;
+  if (lanesPerScheduledLane != 1) {
+    scheduledOutputLane =
+      affine::AffineApplyOp::create(builder,
+                                    loc,
+                                    AffineMap::get(/*dimCount=*/1,
+                                                   /*symbolCount=*/0,
+                                                   builder.getAffineDimExpr(0) * lanesPerScheduledLane),
+                                    ValueRange {scheduledLane})
+        .getResult();
+  }
+  for (unsigned dim = 0; dim < static_cast<unsigned>(localFragmentType.getRank()); ++dim)
+    offsets.push_back(dim == laneDim ? OpFoldResult(scheduledOutputLane) : OpFoldResult(builder.getIndexAttr(0)));
+  return offsets;
+}
+
+static LogicalResult materializeMultiCpuPeftClass(
+  PatternRewriter &rewriter,
+  SpatScheduledComputeBatch scheduled,
+  const PeftClassPlan &peftClassPlan,
+  const MergeScheduleResult &schedule,
+  DenseMap<GraphComputeBlockKey, Block *> &graphComputeToBlockMap,
+  ScheduledMaterializationRecord &record,
+  uint64_t &exchangeId) {
+  std::map<std::vector<uint32_t>, Value> laneStartTableCache;
+  ArrayRef<ScheduledStepPlan> stepPlans = record.stepPlans;
+  for (const ScheduledStepPlan &stepPlan : stepPlans) {
+    const ComputeStepTuple &stepTuple = stepPlan.stepTuple;
+    SourceLaneSelector sourceLaneSelector =
+      buildSourceLaneSelector(rewriter, stepTuple, scheduled.getOperation(), laneStartTableCache);
+    SmallVector<Type> blockArgTypes {rewriter.getIndexType()};
+    SmallVector<Location> blockArgLocs {scheduled.getLoc()};
+    for (Value weight : scheduled.getWeights()) {
+      blockArgTypes.push_back(weight.getType());
+      blockArgLocs.push_back(weight.getLoc());
+    }
+    for (Value input : scheduled.getInputs()) {
+      blockArgTypes.push_back(input.getType());
+      blockArgLocs.push_back(input.getLoc());
+    }
+    for (Type resultType : scheduled.getResultTypes()) {
+      blockArgTypes.push_back(resultType);
+      blockArgLocs.push_back(scheduled.getLoc());
+    }
+    Block *block = rewriter.createBlock(&scheduled.getBody(), scheduled.getBody().end(), blockArgTypes, blockArgLocs);
+    for (const ComputeInstance &instance : stepTuple.instances) {
+      GraphComputeBlockKey key = getGraphComputeBlockKey(instance);
+      graphComputeToBlockMap[key] = block;
+      record.computeKeys.push_back(key);
+      record.blocks.push_back(block);
+    }
+
+    rewriter.setInsertionPointToStart(block);
+    Value scheduledLane = block->getArgument(0);
+    const ComputeInstance &representative = stepTuple.instances.front();
+    SmallVector<Value> finalLocalFragments;
+    SmallVector<unsigned> finalLaneDims;
+
+    if (auto compute = dyn_cast<SpatCompute>(representative.op)) {
+      IRMapping mapper;
+      for (auto [index, weight] : llvm::enumerate(compute.getWeights()))
+        mapper.map(*compute.getWeightArgument(index), block->getArgument(1 + index));
+      unsigned firstInputArg = 1 + scheduled.getWeights().size();
+      for (auto [index, input] : llvm::enumerate(compute.getInputs())) {
+        Value mapped;
+        if (failed(mapMultiCpuTupleInput(rewriter,
+                                         input.getLoc(),
+                                         input,
+                                         stepTuple,
+                                         peftClassPlan,
+                                         schedule,
+                                         scheduled.getInputs(),
+                                         *block,
+                                         firstInputArg,
+                                         exchangeId,
+                                         mapped)))
+          return failure();
+        mapper.map(*compute.getInputArgument(index), mapped);
+      }
+      SmallVector<Value> yieldedValues;
+      cloneComputeBody(rewriter, compute.getBody().front(), mapper, yieldedValues);
+      for (Value yielded : yieldedValues) {
+        auto tensorType = dyn_cast<RankedTensorType>(yielded.getType());
+        if (!tensorType || !tensorType.hasStaticShape() || tensorType.getRank() == 0)
+          return compute.emitOpError("scheduled materialization only supports static ranked tensor scalar step results");
+        SmallVector<ReassociationIndices> reassociation;
+        reassociation.push_back({0, 1});
+        for (int64_t dim = 1; dim < tensorType.getRank(); ++dim)
+          reassociation.push_back({static_cast<int64_t>(dim + 1)});
+        SmallVector<int64_t> expandedShape {1};
+        llvm::append_range(expandedShape, tensorType.getShape());
+        finalLocalFragments.push_back(tensor::ExpandShapeOp::create(rewriter,
+                                                                    scheduled.getLoc(),
+                                                                    RankedTensorType::get(expandedShape, tensorType.getElementType()),
+                                                                    yielded,
+                                                                    reassociation)
+                                       .getResult());
+        finalLaneDims.push_back(0);
+      }
+    } else {
+      auto batch = cast<SpatComputeBatch>(representative.op);
+      SmallVector<Value> localFragments;
+      SmallVector<BatchFragmentSpec> fragmentSpecs;
+      for (unsigned resultIndex = 0; resultIndex < batch.getNumResults(); ++resultIndex) {
+        auto spec = getBatchFragmentSpec(batch, resultIndex, representative.laneCount);
+        if (failed(spec))
+          return failure();
+        fragmentSpecs.push_back(*spec);
+        auto empty = createEmptyTensorForType(rewriter, batch.getLoc(), spec->resultType);
+        if (failed(empty))
+          return failure();
+        localFragments.push_back(*empty);
+      }
+
+      Value lower = getOrCreateIndexConstant(rewriter, scheduled.getOperation(), 0);
+      Value upper = getOrCreateIndexConstant(rewriter, scheduled.getOperation(), representative.laneCount);
+      Value step = getOrCreateIndexConstant(rewriter, scheduled.getOperation(), 1);
+      auto loop = buildNormalizedScfFor(
+        rewriter,
+        batch.getLoc(),
+        lower,
+        upper,
+        step,
+        localFragments,
+        [&](OpBuilder &builder, Location bodyLoc, Value innerLane, ValueRange iterArgs, SmallVectorImpl<Value> &yielded) -> LogicalResult {
+
+          IRMapping mapper;
+          FailureOr<Value> sourceLaneStart =
+            buildSourceLaneStartForScheduledLane(builder, bodyLoc, scheduledLane, sourceLaneSelector, scheduled.getOperation());
+          if (failed(sourceLaneStart))
+            return failure();
+          Value sourceLane =
+            affine::AffineApplyOp::create(builder,
+                                          bodyLoc,
+                                          AffineMap::get(/*dimCount=*/2,
+                                                         /*symbolCount=*/0,
+                                                         builder.getAffineDimExpr(0) + builder.getAffineDimExpr(1)),
+                                          ValueRange {*sourceLaneStart, innerLane})
+              .getResult();
+          mapper.map(*batch.getLaneArgument(), sourceLane);
+          for (auto [index, weight] : llvm::enumerate(batch.getWeights()))
+            mapper.map(*batch.getWeightArgument(index), block->getArgument(1 + index));
+          unsigned firstInputArg = 1 + scheduled.getWeights().size();
+          for (auto [index, input] : llvm::enumerate(batch.getInputs())) {
+            Value mapped;
+            if (failed(mapMultiCpuTupleInput(builder,
+                                             input.getLoc(),
+                                             input,
+                                             stepTuple,
+                                             peftClassPlan,
+                                             schedule,
+                                             scheduled.getInputs(),
+                                             *block,
+                                             firstInputArg,
+                                             exchangeId,
+                                             mapped)))
+              return failure();
+            mapper.map(*batch.getInputArgument(index), mapped);
+          }
+          for (unsigned index = 0; index < batch.getNumResults(); ++index)
+            mapper.map(*batch.getOutputArgument(index), iterArgs[index]);
+          for (Operation &op : batch.getBody().front().without_terminator())
+            builder.clone(op, mapper);
+
+          auto inParallel = dyn_cast<SpatInParallelOp>(batch.getBody().front().getTerminator());
+          if (!inParallel)
+            return batch.emitOpError("expected spat.in_parallel in resultful spat.graph_compute_batch"), failure();
+
+          DenseMap<BlockArgument, size_t> outputIndexByArg;
+          for (size_t index = 0; index < batch.getNumResults(); ++index)
+            outputIndexByArg[*batch.getOutputArgument(index)] = index;
+
+          SmallVector<Value> current(iterArgs.begin(), iterArgs.end());
+          for (Operation &op : inParallel.getRegion().front()) {
+            auto insert = dyn_cast<tensor::ParallelInsertSliceOp>(&op);
+            if (!insert)
+              return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+            auto oldDest = dyn_cast<BlockArgument>(insert.getDest());
+            if (!oldDest || !outputIndexByArg.count(oldDest))
+              return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments"), failure();
+            size_t resultIndex = outputIndexByArg.lookup(oldDest);
+            SmallVector<OpFoldResult> offsets = remapMixedOffsets(insert.getMixedOffsets(), mapper);
+            offsets[fragmentSpecs[resultIndex].laneDim] = innerLane;
+            current[resultIndex] = tensor::InsertSliceOp::create(builder,
+                                                                 insert.getLoc(),
+                                                                 mapper.lookup(insert.getSource()),
+                                                                 current[resultIndex],
+                                                                 offsets,
+                                                                 remapMixedOffsets(insert.getMixedSizes(), mapper),
+                                                                 remapMixedOffsets(insert.getMixedStrides(), mapper))
+                                     .getResult();
+          }
+          llvm::append_range(yielded, current);
+          return success();
+        });
+      if (failed(loop))
+        return failure();
+      finalLocalFragments.assign(loop->results.begin(), loop->results.end());
+      for (const BatchFragmentSpec &spec : fragmentSpecs)
+        finalLaneDims.push_back(spec.laneDim);
+    }
+
+    auto inParallel = SpatInParallelOp::create(rewriter, scheduled.getLoc());
+    rewriter.setInsertionPointToStart(&inParallel.getRegion().front());
+    for (auto [resultIndex, localFragment] : llvm::enumerate(finalLocalFragments)) {
+      auto localFragmentType = cast<RankedTensorType>(localFragment.getType());
+      int64_t lanesPerScheduledLane = isa<SpatCompute>(representative.op) ? 1 : representative.laneCount;
+      SmallVector<OpFoldResult> offsets = buildScheduledOutputInsertOffsets(
+        rewriter,
+        scheduled.getLoc(),
+        scheduledLane,
+        lanesPerScheduledLane,
+        localFragmentType,
+        finalLaneDims[resultIndex]);
+      SmallVector<OpFoldResult> sizes;
+      SmallVector<OpFoldResult> strides;
+      for (int64_t dim : localFragmentType.getShape()) {
+        sizes.push_back(rewriter.getIndexAttr(dim));
+        strides.push_back(rewriter.getIndexAttr(1));
+      }
+      tensor::ParallelInsertSliceOp::create(
+        rewriter,
+        scheduled.getLoc(),
+        localFragment,
+        block->getArgument(getScheduledBatchResultArgBase(scheduled) + stepPlan.resultOffset + resultIndex),
+        offsets,
+        sizes,
+        strides);
+    }
+  }
+  return success();
+}
+
+
+} // namespace
+
+
+FailureOr<ScheduledComputeMaterializationResult>
+materializeScheduledCompute(func::FuncOp funcOp,
+                            const MergeScheduleResult &schedule,
+                            PatternRewriter &rewriter) {
+  llvm::MapVector<size_t, PeftClassPlan> peftClassPlans;
+  for (const ComputeInstance &instance : schedule.dominanceOrderCompute) {
+    size_t cpu = schedule.computeToCpuMap.lookup(instance);
+    size_t canonicalPeftClassId = getCanonicalPeftClassId(cpu, schedule);
+    auto &peftClassPlan = peftClassPlans[canonicalPeftClassId];
+    peftClassPlan.canonicalClassId = canonicalPeftClassId;
+    if (!llvm::is_contained(peftClassPlan.cpus, cpu))
+      peftClassPlan.cpus.push_back(cpu);
+    peftClassPlan.instancesByCpu[cpu].push_back(instance);
+  }
+  for (auto &entry : peftClassPlans) {
+    PeftClassPlan &peftClassPlan = entry.second;
+    llvm::sort(peftClassPlan.cpus);
+    for (size_t cpu : peftClassPlan.cpus)
+      llvm::sort(peftClassPlan.instancesByCpu[cpu], [&](const ComputeInstance &lhs, const ComputeInstance &rhs) {
+        return schedule.computeToCpuSlotMap.lookup(lhs) < schedule.computeToCpuSlotMap.lookup(rhs);
+      });
+    if (failed(verifyPeftClassPlan(funcOp.getOperation(), peftClassPlan, schedule)))
+      return failure();
+    if (failed(collectPeftClassOperandsAndResults(peftClassPlan, schedule)))
+      return failure();
+  }
+
+  Operation *insertionPoint = funcOp.getBody().front().getTerminator();
+  DenseMap<GraphComputeBlockKey, Block *> graphComputeToBlockMap;
+  DenseMap<size_t, SpatScheduledCompute> scheduledComputes;
+  DenseMap<size_t, SpatScheduledComputeBatch> scheduledComputeBatches;
+  DenseMap<size_t, size_t> classToRecordIndex;
+  std::vector<ScheduledMaterializationRecord> materializedSchedules;
+  uint64_t exchangeId = 0;
+
+  for (auto &entry : peftClassPlans) {
+    PeftClassPlan &peftClassPlan = entry.second;
+    rewriter.setInsertionPoint(insertionPoint);
+
+    ScheduledMaterializationRecord record;
+    record.canonicalPeftClassId = peftClassPlan.canonicalClassId;
+    record.cpus = peftClassPlan.cpus;
+    record.stepPlans = buildScheduledStepPlans(peftClassPlan);
+
+    if (peftClassPlan.cpus.size() == 1) {
+      auto scheduled = SpatScheduledCompute::create(
+        rewriter,
+        peftClassPlan.instancesByCpu.lookup(peftClassPlan.cpus.front()).front().op->getLoc(),
+        TypeRange(peftClassPlan.resultTypes),
+        peftClassPlan.weights,
+        peftClassPlan.inputs);
+      scheduled->setAttr(kCoreIdAttrName, rewriter.getI32IntegerAttr(static_cast<int32_t>(peftClassPlan.cpus.front())));
+      scheduled->setAttr("scheduled.peft_cpus", rewriter.getDenseI64ArrayAttr(toI64Array(peftClassPlan.cpus)));
+      SmallVector<Attribute> stepSources;
+      SmallVector<Attribute> sourceLaneSelectors;
+      SmallVector<int64_t> stepResultOffsets;
+      SmallVector<int64_t> stepResultCounts;
+      SmallVector<int64_t> sourceLaneStarts;
+      SmallVector<int64_t> sourceLaneCounts;
+      size_t resultOffset = 0;
+      for (const ComputeInstance &instance : peftClassPlan.instancesByCpu.lookup(peftClassPlan.cpus.front())) {
+        stepSources.push_back(rewriter.getStringAttr(getInstanceName(instance)));
+        sourceLaneSelectors.push_back(rewriter.getStringAttr(isa<SpatCompute>(instance.op) ? "scalar" : "affine"));
+        size_t resultCount = getComputeInstanceResultValueCount(instance);
+        stepResultOffsets.push_back(static_cast<int64_t>(resultOffset));
+        stepResultCounts.push_back(static_cast<int64_t>(resultCount));
+        resultOffset += resultCount;
+        if (isa<SpatCompute>(instance.op)) {
+          sourceLaneStarts.push_back(0);
+          sourceLaneCounts.push_back(0);
+        } else {
+          sourceLaneStarts.push_back(instance.laneStart);
+          sourceLaneCounts.push_back(instance.laneCount);
+        }
+      }
+      scheduled->setAttr("scheduled.step_sources", rewriter.getArrayAttr(stepSources));
+      scheduled->setAttr("scheduled.step_result_offsets", rewriter.getDenseI64ArrayAttr(stepResultOffsets));
+      scheduled->setAttr("scheduled.step_result_counts", rewriter.getDenseI64ArrayAttr(stepResultCounts));
+      scheduled->setAttr("scheduled.source_lane_starts", rewriter.getDenseI64ArrayAttr(sourceLaneStarts));
+      scheduled->setAttr("scheduled.source_lane_counts", rewriter.getDenseI64ArrayAttr(sourceLaneCounts));
+      scheduled->setAttr("scheduled.source_lane_selector", rewriter.getArrayAttr(sourceLaneSelectors));
+      record.scheduledOp = scheduled.getOperation();
+      scheduledComputes[peftClassPlan.canonicalClassId] = scheduled;
+    } else {
+      auto scheduled = SpatScheduledComputeBatch::create(rewriter,
+                                                         peftClassPlan.instancesByCpu.lookup(peftClassPlan.cpus.front()).front().op->getLoc(),
+                                                         TypeRange(peftClassPlan.resultTypes),
+                                                         rewriter.getI32IntegerAttr(static_cast<int32_t>(peftClassPlan.cpus.size())),
+                                                         peftClassPlan.weights,
+                                                         peftClassPlan.inputs);
+      scheduled->setAttr(kCoreIdsAttrName, rewriter.getDenseI32ArrayAttr(toI32Array(peftClassPlan.cpus)));
+      scheduled->setAttr("scheduled.peft_cpus", rewriter.getDenseI64ArrayAttr(toI64Array(peftClassPlan.cpus)));
+      SmallVector<Attribute> stepSources;
+      SmallVector<Attribute> sourceLaneSelectors;
+      SmallVector<int64_t> resultOffsets;
+      SmallVector<int64_t> resultCounts;
+      SmallVector<int64_t> sourceLaneStarts;
+      SmallVector<int64_t> sourceLaneCounts;
+      for (const ScheduledStepPlan &stepPlan : record.stepPlans) {
+        stepSources.push_back(rewriter.getStringAttr(getInstanceName(stepPlan.stepTuple.instances.front())));
+        sourceLaneSelectors.push_back(rewriter.getStringAttr(usesAffineSourceLaneMapping(stepPlan.stepTuple) ? "affine" : "table"));
+        resultOffsets.push_back(static_cast<int64_t>(stepPlan.resultOffset));
+        resultCounts.push_back(static_cast<int64_t>(stepPlan.resultCount));
+        for (const ComputeInstance &instance : stepPlan.stepTuple.instances) {
+          sourceLaneStarts.push_back(instance.laneStart);
+          sourceLaneCounts.push_back(instance.laneCount);
+        }
+      }
+      RankedTensorType sourceLaneTableType = RankedTensorType::get(
+        {static_cast<int64_t>(record.stepPlans.size()), static_cast<int64_t>(peftClassPlan.cpus.size())},
+        rewriter.getI64Type());
+      scheduled->setAttr("scheduled.step_sources", rewriter.getArrayAttr(stepSources));
+      scheduled->setAttr("scheduled.step_result_offsets", rewriter.getDenseI64ArrayAttr(resultOffsets));
+      scheduled->setAttr("scheduled.step_result_counts", rewriter.getDenseI64ArrayAttr(resultCounts));
+      scheduled->setAttr("scheduled.source_lane_starts", DenseElementsAttr::get(sourceLaneTableType, ArrayRef<int64_t>(sourceLaneStarts)));
+      scheduled->setAttr("scheduled.source_lane_counts", DenseElementsAttr::get(sourceLaneTableType, ArrayRef<int64_t>(sourceLaneCounts)));
+      scheduled->setAttr("scheduled.source_lane_selector", rewriter.getArrayAttr(sourceLaneSelectors));
+      record.scheduledOp = scheduled.getOperation();
+      scheduledComputeBatches[peftClassPlan.canonicalClassId] = scheduled;
+    }
+
+    classToRecordIndex[peftClassPlan.canonicalClassId] = materializedSchedules.size();
+    materializedSchedules.push_back(std::move(record));
+  }
+
+  for (auto &entry : peftClassPlans) {
+    PeftClassPlan &peftClassPlan = entry.second;
+    ScheduledMaterializationRecord &record =
+      materializedSchedules[classToRecordIndex.lookup(peftClassPlan.canonicalClassId)];
+    if (peftClassPlan.cpus.size() == 1) {
+      if (failed(materializeSingleCpuPeftClass(rewriter,
+                                               scheduledComputes.lookup(peftClassPlan.canonicalClassId),
+                                               peftClassPlan,
+                                               schedule,
+                                               graphComputeToBlockMap,
+                                               record,
+                                               exchangeId)))
+        return failure();
+    } else {
+      if (failed(materializeMultiCpuPeftClass(rewriter,
+                                              scheduledComputeBatches.lookup(peftClassPlan.canonicalClassId),
+                                              peftClassPlan,
+                                              schedule,
+                                              graphComputeToBlockMap,
+                                              record,
+                                              exchangeId)))
+        return failure();
+    }
+  }
+
+  return ScheduledComputeMaterializationResult {std::move(peftClassPlans), std::move(materializedSchedules), std::move(graphComputeToBlockMap)};
+}
+
+
+} // namespace spatial
+} // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.hpp

@@ -0,0 +1,20 @@
+#pragma once
+
+#include "ScheduledComputePlan.hpp"
+
+namespace onnx_mlir {
+namespace spatial {
+
+struct ScheduledComputeMaterializationResult {
+  llvm::MapVector<size_t, PeftClassPlan> peftClassPlans;
+  std::vector<ScheduledMaterializationRecord> materializedSchedules;
+  DenseMap<GraphComputeBlockKey, Block *> graphComputeToBlockMap;
+};
+
+FailureOr<ScheduledComputeMaterializationResult>
+materializeScheduledCompute(func::FuncOp funcOp,
+                            const MergeScheduleResult &schedule,
+                            PatternRewriter &rewriter);
+
+} // namespace spatial
+} // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp

@@ -0,0 +1,337 @@
+#pragma once
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/AsmState.h"
+#include "mlir/IR/IRMapping.h"
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <limits>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "Scheduling/ComputeInstanceUtils.hpp"
+#include "Scheduling/MergeSchedulingAnalysis.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+namespace onnx_mlir {
+namespace spatial {
+
+using namespace mlir;
+
+struct ProducerValueKey {
+  ComputeInstance instance;
+  size_t resultIndex = 0;
+
+  bool operator==(const ProducerValueKey &other) const {
+    return instance == other.instance && resultIndex == other.resultIndex;
+  }
+};
+
+struct GraphComputeBlockKey {
+  Operation *op = nullptr;
+  uint32_t laneStart = 0;
+  uint32_t laneCount = 1;
+
+  bool operator==(const GraphComputeBlockKey &other) const {
+    return op == other.op && laneStart == other.laneStart && laneCount == other.laneCount;
+  }
+};
+
+struct PeftClassPlan {
+  size_t canonicalClassId = 0;
+  SmallVector<size_t> cpus;
+  llvm::MapVector<size_t, SmallVector<ComputeInstance>> instancesByCpu;
+
+  SmallVector<Value> weights;
+  SmallVector<Value> inputs;
+  SmallVector<Type> resultTypes;
+};
+
+struct ComputeStepTuple {
+  SmallVector<ComputeInstance> instances;
+};
+
+struct ScheduledStepPlan {
+  ComputeStepTuple stepTuple;
+  size_t stepIndex = 0;
+  size_t resultOffset = 0;
+  size_t resultCount = 0;
+};
+
+struct SourceLaneAffineMapping {
+  uint32_t baseLaneStart = 0;
+  uint32_t laneCount = 1;
+};
+
+struct SourceLaneSelector {
+  enum class Kind { Affine, Table } kind = Kind::Affine;
+  SourceLaneAffineMapping affine;
+  Value table;
+  SmallVector<uint32_t> tableValues;
+};
+
+struct DeferredEndpoint {
+  ComputeInstance instance;
+  size_t cpu = 0;
+  size_t canonicalClassId = 0;
+  uint32_t laneStart = 0;
+  uint32_t laneCount = 1;
+};
+
+struct DeferredTransferMetadata {
+  SmallVector<DeferredEndpoint> producers;
+  SmallVector<DeferredEndpoint> consumers;
+  bool isBatched = false;
+  SmallVector<int64_t> sourceOperandForScheduledLane;
+  bool multiSourcePayload = false;
+};
+
+struct ScheduledMaterializationRecord {
+  Operation *scheduledOp = nullptr;
+  size_t canonicalPeftClassId = 0;
+  SmallVector<size_t> cpus;
+  SmallVector<ScheduledStepPlan> stepPlans;
+  SmallVector<GraphComputeBlockKey> computeKeys;
+  SmallVector<Block *> blocks;
+};
+
+struct ScheduledComputePrintContext {
+  mlir::AsmState asmState;
+  explicit ScheduledComputePrintContext(ModuleOp module, const OpPrintingFlags &flags = OpPrintingFlags())
+    : asmState(module.getOperation(), flags) {}
+};
+
+inline GraphComputeBlockKey getGraphComputeBlockKey(const ComputeInstance &instance) {
+  return {instance.op, instance.laneStart, instance.laneCount};
+}
+
+inline SmallVector<size_t> getPeftClassCpus(size_t cpu, const MergeScheduleResult &schedule) {
+  llvm::SmallDenseSet<size_t, 8> seen;
+  SmallVector<size_t> cpus;
+  auto append = [&](size_t value) {
+    if (seen.insert(value).second)
+      cpus.push_back(value);
+  };
+  append(cpu);
+  if (auto it = schedule.equivalentClass.find(cpu); it != schedule.equivalentClass.end())
+    for (size_t equivalentCpu : it->second)
+      append(equivalentCpu);
+  llvm::sort(cpus);
+  return cpus;
+}
+
+inline size_t getCanonicalPeftClassId(size_t cpu, const MergeScheduleResult &schedule) {
+  SmallVector<size_t> cpus = getPeftClassCpus(cpu, schedule);
+  return cpus.empty() ? cpu : cpus.front();
+}
+
+inline size_t getScheduledCpuForComputeInstance(const ComputeInstance &instance, const MergeScheduleResult &schedule) {
+  if (auto it = schedule.computeToCpuMap.find(instance); it != schedule.computeToCpuMap.end())
+    return it->second;
+
+  auto batch = dyn_cast<SpatComputeBatch>(instance.op);
+  assert(batch && instance.laneCount != 0 && "missing scheduled CPU for non-batch compute instance");
+  assert(instance.laneStart < static_cast<uint32_t>(batch.getLaneCount()) && "batch lane start out of range");
+  ComputeInstance chunk = getBatchChunkForLane(batch, instance.laneStart);
+  auto it = schedule.computeToCpuMap.find(chunk);
+  assert(it != schedule.computeToCpuMap.end() && "missing scheduled CPU for batch chunk");
+  return it->second;
+}
+
+inline std::string getInstanceName(const ComputeInstance &instance) {
+  return llvm::formatv("{0}[lanes={1}:{2}]",
+                       instance.op->getName().getStringRef(),
+                       instance.laneStart,
+                       instance.laneStart + instance.laneCount)
+    .str();
+}
+
+inline SmallVector<int64_t> toI64Array(ArrayRef<size_t> values) {
+  SmallVector<int64_t> converted;
+  converted.reserve(values.size());
+  for (size_t value : values)
+    converted.push_back(static_cast<int64_t>(value));
+  return converted;
+}
+
+inline SmallVector<int32_t> toI32Array(ArrayRef<size_t> values) {
+  SmallVector<int32_t> converted;
+  converted.reserve(values.size());
+  for (size_t value : values)
+    converted.push_back(static_cast<int32_t>(value));
+  return converted;
+}
+
+inline unsigned getScheduledBatchResultArgBase(SpatScheduledComputeBatch scheduled) {
+  unsigned weightArgBase = 1;
+  unsigned inputArgBase = weightArgBase + scheduled.getWeights().size();
+  return inputArgBase + scheduled.getInputs().size();
+}
+
+inline SmallVector<GraphComputeBlockKey> collectExpectedGraphComputeBlockKeys(func::FuncOp funcOp) {
+  SmallVector<GraphComputeBlockKey> keys;
+  for (Operation &op : funcOp.getOps()) {
+    if (auto compute = dyn_cast<SpatGraphCompute>(&op))
+      keys.push_back(getGraphComputeBlockKey({compute.getOperation(), 0, 1}));
+    else if (auto batch = dyn_cast<SpatGraphComputeBatch>(&op))
+      for (ComputeInstance chunk : getBatchChunksForRange(batch, 0, static_cast<uint32_t>(batch.getLaneCount())))
+        keys.push_back(getGraphComputeBlockKey(chunk));
+  }
+  return keys;
+}
+
+inline size_t countPeftEquivalenceClasses(const MergeScheduleResult &schedule) {
+  llvm::SmallDenseSet<size_t, 16> classes;
+  for (const ComputeInstance &instance : schedule.dominanceOrderCompute) {
+    auto cpuIt = schedule.computeToCpuMap.find(instance);
+    if (cpuIt == schedule.computeToCpuMap.end())
+      continue;
+    classes.insert(getCanonicalPeftClassId(cpuIt->second, schedule));
+  }
+  return classes.size();
+}
+
+inline SmallVector<ComputeStepTuple> buildComputeStepTuples(const PeftClassPlan &peftClassPlan) {
+  SmallVector<ComputeStepTuple> stepTuples;
+  if (peftClassPlan.cpus.empty())
+    return stepTuples;
+  size_t stepCount = peftClassPlan.instancesByCpu.lookup(peftClassPlan.cpus.front()).size();
+  stepTuples.resize(stepCount);
+  for (size_t stepIndex = 0; stepIndex < stepCount; ++stepIndex)
+    for (size_t cpu : peftClassPlan.cpus)
+      stepTuples[stepIndex].instances.push_back(peftClassPlan.instancesByCpu.lookup(cpu)[stepIndex]);
+  return stepTuples;
+}
+
+inline size_t getComputeInstanceResultValueCount(const ComputeInstance &instance) {
+  return TypeSwitch<Operation *, size_t>(instance.op)
+    .Case<SpatCompute>([](SpatCompute compute) { return compute.getNumResults(); })
+    .Case<SpatComputeBatch>([](SpatComputeBatch batch) { return batch.getNumResults(); })
+    .Default([](Operation *) -> size_t {
+      llvm_unreachable("expected graph compute or graph compute batch");
+    });
+}
+
+inline SmallVector<ScheduledStepPlan> buildScheduledStepPlans(const PeftClassPlan &peftClassPlan) {
+  SmallVector<ScheduledStepPlan> stepPlans;
+  size_t resultOffset = 0;
+  for (auto [stepIndex, stepTuple] : llvm::enumerate(buildComputeStepTuples(peftClassPlan))) {
+    assert(!stepTuple.instances.empty() && "expected non-empty step tuple");
+    size_t resultCount = getComputeInstanceResultValueCount(stepTuple.instances.front());
+    stepPlans.push_back(ScheduledStepPlan {stepTuple, stepIndex, resultOffset, resultCount});
+    resultOffset += resultCount;
+  }
+  return stepPlans;
+}
+
+inline bool valueTransitivelyDependsOn(Value value, Value dependency) {
+  SmallVector<Value> worklist {value};
+  DenseSet<Value> visited;
+  while (!worklist.empty()) {
+    Value current = worklist.pop_back_val();
+    if (!visited.insert(current).second)
+      continue;
+    if (current == dependency)
+      return true;
+    Operation *def = current.getDefiningOp();
+    if (!def)
+      continue;
+    for (Value operand : def->getOperands())
+      worklist.push_back(operand);
+  }
+  return false;
+}
+
+inline std::optional<SourceLaneAffineMapping> getSourceLaneAffineMapping(const ComputeStepTuple &stepTuple) {
+  if (stepTuple.instances.empty())
+    return std::nullopt;
+  const ComputeInstance &reference = stepTuple.instances.front();
+  for (const ComputeInstance &instance : stepTuple.instances) {
+    if (instance.op != reference.op || instance.laneCount != reference.laneCount)
+      return std::nullopt;
+  }
+
+  for (size_t index = 0; index < stepTuple.instances.size(); ++index) {
+    uint32_t expectedLaneStart = reference.laneStart + static_cast<uint32_t>(index) * reference.laneCount;
+    if (stepTuple.instances[index].laneStart != expectedLaneStart)
+      return std::nullopt;
+  }
+
+  return SourceLaneAffineMapping {reference.laneStart, reference.laneCount};
+}
+
+inline bool usesAffineSourceLaneMapping(const ComputeStepTuple &stepTuple) {
+  return getSourceLaneAffineMapping(stepTuple).has_value();
+}
+
+inline SmallVector<uint32_t> collectSourceLaneStarts(const ComputeStepTuple &stepTuple) {
+  SmallVector<uint32_t> sourceLaneStarts;
+  sourceLaneStarts.reserve(stepTuple.instances.size());
+  for (const ComputeInstance &instance : stepTuple.instances)
+    sourceLaneStarts.push_back(instance.laneStart);
+  return sourceLaneStarts;
+}
+
+inline Value createI64LookupTableConstant(OpBuilder &builder, Operation *constantAnchor, ArrayRef<int64_t> values) {
+  OpBuilder::InsertionGuard guard(builder);
+  builder.setInsertionPoint(constantAnchor);
+  RankedTensorType tableType = RankedTensorType::get({static_cast<int64_t>(values.size())}, builder.getI64Type());
+  DenseElementsAttr tableAttr = DenseElementsAttr::get(tableType, values);
+  return arith::ConstantOp::create(builder, constantAnchor->getLoc(), tableType, tableAttr).getResult();
+}
+
+inline FailureOr<Value> createEmptyTensorForType(OpBuilder &builder, Location loc, Type type) {
+  auto tensorType = dyn_cast<RankedTensorType>(type);
+  if (!tensorType || !tensorType.hasStaticShape())
+    return failure();
+  return tensor::EmptyOp::create(builder, loc, tensorType.getShape(), tensorType.getElementType()).getResult();
+}
+
+} // namespace spatial
+} // namespace onnx_mlir
+
+namespace llvm {
+template <>
+struct DenseMapInfo<onnx_mlir::spatial::ProducerValueKey> {
+  static onnx_mlir::spatial::ProducerValueKey getEmptyKey() {
+    return {DenseMapInfo<onnx_mlir::spatial::ComputeInstance>::getEmptyKey(), std::numeric_limits<size_t>::max()};
+  }
+  static onnx_mlir::spatial::ProducerValueKey getTombstoneKey() {
+    return {DenseMapInfo<onnx_mlir::spatial::ComputeInstance>::getTombstoneKey(), std::numeric_limits<size_t>::max()};
+  }
+  static unsigned getHashValue(const onnx_mlir::spatial::ProducerValueKey &key) {
+    return hash_combine(DenseMapInfo<onnx_mlir::spatial::ComputeInstance>::getHashValue(key.instance), key.resultIndex);
+  }
+  static bool isEqual(const onnx_mlir::spatial::ProducerValueKey &lhs,
+                      const onnx_mlir::spatial::ProducerValueKey &rhs) {
+    return lhs == rhs;
+  }
+};
+
+template <>
+struct DenseMapInfo<onnx_mlir::spatial::GraphComputeBlockKey> {
+  static onnx_mlir::spatial::GraphComputeBlockKey getEmptyKey() {
+    return {DenseMapInfo<mlir::Operation *>::getEmptyKey(), UINT32_MAX, UINT32_MAX};
+  }
+  static onnx_mlir::spatial::GraphComputeBlockKey getTombstoneKey() {
+    return {DenseMapInfo<mlir::Operation *>::getTombstoneKey(), UINT32_MAX, UINT32_MAX - 1};
+  }
+  static unsigned getHashValue(const onnx_mlir::spatial::GraphComputeBlockKey &key) {
+    return hash_combine(key.op, key.laneStart, key.laneCount);
+  }
+  static bool isEqual(const onnx_mlir::spatial::GraphComputeBlockKey &lhs,
+                      const onnx_mlir::spatial::GraphComputeBlockKey &rhs) {
+    return lhs == rhs;
+  }
+};
+} // namespace llvm

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp

@@ -0,0 +1,312 @@
+#include "ScheduledComputeReport.hpp"
+
+#include "llvm/Support/raw_os_ostream.h"
+
+#include <fstream>
+
+#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
+#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
+
+namespace onnx_mlir {
+namespace spatial {
+
+using namespace mlir;
+namespace {
+
+static std::string formatValueLabel(Value value, AsmState &asmState) {
+  std::string storage;
+  llvm::raw_string_ostream os(storage);
+  value.printAsOperand(os, asmState);
+  return storage;
+}
+
+static std::string formatOperationLabel(Operation *op, AsmState &asmState) {
+  if (op->getNumResults() == 0)
+    return op->getName().getStringRef().str();
+  std::string storage;
+  llvm::raw_string_ostream os(storage);
+  llvm::interleaveComma(op->getResults(), os, [&](Value result) { os << formatValueLabel(result, asmState); });
+  return os.str();
+}
+
+static std::string formatGraphComputeBlockKey(const GraphComputeBlockKey &key, AsmState &asmState) {
+  return llvm::formatv("{0} {1}", formatOperationLabel(key.op, asmState), key.op->getName().getStringRef()).str();
+}
+
+static std::string formatComputeInstanceForReport(const ComputeInstance &instance, AsmState &asmState) {
+  std::string opLabel = formatGraphComputeBlockKey(getGraphComputeBlockKey(instance), asmState);
+  if (isa<SpatCompute>(instance.op))
+    return opLabel;
+  return llvm::formatv("{0} sourceLanes [{1}:{2}]",
+                       opLabel,
+                       instance.laneStart,
+                       instance.laneStart + instance.laneCount)
+    .str();
+}
+
+template <typename T>
+static void printIndexedList(raw_ostream &os, ArrayRef<T> values) {
+  os << "[";
+  llvm::interleaveComma(llvm::enumerate(values), os, [&](auto indexedValue) {
+    os << indexedValue.index() << ":" << indexedValue.value();
+  });
+  os << "]";
+}
+
+struct PeftMaterializationReportSummary {
+  size_t scalarGraphCompute = 0;
+  size_t graphComputeBatchOps = 0;
+  size_t scalarGraphComputeInstances = 0;
+  size_t graphComputeBatchInstances = 0;
+  size_t peftClasses = 0;
+  size_t singleCpuClasses = 0;
+  size_t multiCpuClasses = 0;
+  size_t scheduledCompute = 0;
+  size_t scheduledComputeBatch = 0;
+  size_t deferredCommunication = 0;
+  size_t deferredCommunicationScalarMetadata = 0;
+  size_t deferredCommunicationScheduledLaneMetadata = 0;
+  size_t deferredCommunicationMultiSourcePayloads = 0;
+};
+
+static PeftMaterializationReportSummary buildPeftMaterializationReportSummary(
+  func::FuncOp funcOp,
+  const MergeScheduleResult &schedule,
+  const llvm::MapVector<size_t, PeftClassPlan> &peftClassPlans,
+  ArrayRef<ScheduledMaterializationRecord> materializedSchedules) {
+  PeftMaterializationReportSummary summary;
+  for (Operation &op : funcOp.getOps()) {
+    if (isa<SpatGraphCompute>(op))
+      summary.scalarGraphCompute++;
+    else if (isa<SpatGraphComputeBatch>(op)) {
+      summary.graphComputeBatchOps++;
+    }
+  }
+  for (const ComputeInstance &instance : schedule.dominanceOrderCompute)
+    (isa<SpatCompute>(instance.op) ? summary.scalarGraphComputeInstances : summary.graphComputeBatchInstances)++;
+  summary.peftClasses = peftClassPlans.size();
+  for (const auto &entry : peftClassPlans)
+    (entry.second.cpus.size() == 1 ? summary.singleCpuClasses : summary.multiCpuClasses)++;
+  for (const ScheduledMaterializationRecord &record : materializedSchedules) {
+    if (isa<SpatScheduledCompute>(record.scheduledOp))
+      summary.scheduledCompute++;
+    else
+      summary.scheduledComputeBatch++;
+  }
+  funcOp.walk([&](SpatDeferredCommunicationOp transfer) {
+    summary.deferredCommunication++;
+    if (auto batchedAttr = transfer->getAttrOfType<BoolAttr>("batched"); batchedAttr && batchedAttr.getValue())
+      summary.deferredCommunicationScheduledLaneMetadata++;
+    else
+      summary.deferredCommunicationScalarMetadata++;
+    if (transfer.getSources().size() > 1)
+      summary.deferredCommunicationMultiSourcePayloads++;
+  });
+  return summary;
+}
+
+} // namespace
+
+LogicalResult verifyPeftMaterializationReportSummary(func::FuncOp funcOp,
+                                                            const MergeScheduleResult &schedule,
+                                                            const llvm::MapVector<size_t, PeftClassPlan> &peftClassPlans,
+                                                            ArrayRef<ScheduledMaterializationRecord> materializedSchedules) {
+  PeftMaterializationReportSummary summary =
+    buildPeftMaterializationReportSummary(funcOp, schedule, peftClassPlans, materializedSchedules);
+  pim::CappedDiagnosticReporter diagnostics;
+  if (summary.peftClasses != peftClassPlans.size())
+    diagnostics.report(funcOp.getOperation(), [&](Operation *illegalOp) {
+      illegalOp->emitOpError() << "phase-check report PEFT total " << summary.peftClasses
+                               << " does not match classes.size() " << peftClassPlans.size();
+    });
+  if (summary.scalarGraphComputeInstances + summary.graphComputeBatchInstances != schedule.dominanceOrderCompute.size())
+    diagnostics.report(funcOp.getOperation(), [&](Operation *illegalOp) {
+      illegalOp->emitOpError() << "phase-check report total compute instances "
+                               << (summary.scalarGraphComputeInstances + summary.graphComputeBatchInstances)
+                               << " does not match schedule size " << schedule.dominanceOrderCompute.size();
+    });
+  if (summary.scheduledCompute + summary.scheduledComputeBatch != materializedSchedules.size())
+    diagnostics.report(funcOp.getOperation(), [&](Operation *illegalOp) {
+      illegalOp->emitOpError() << "phase-check report scheduled total "
+                               << (summary.scheduledCompute + summary.scheduledComputeBatch)
+                               << " does not match materialized scheduled ops " << materializedSchedules.size();
+    });
+  diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial report verification failed");
+  return success(!diagnostics.hasFailure());
+}
+
+namespace {
+
+static std::string formatStepResultRange(size_t resultOffset, size_t resultCount) {
+  if (resultCount == 1)
+    return llvm::formatv("result[{0}]", resultOffset).str();
+  return llvm::formatv("result[{0}:{1}]", resultOffset, resultOffset + resultCount).str();
+}
+
+static void printMultiSourceDeferredInputs(raw_ostream &os, Block &block) {
+  unsigned deferredInputIndex = 0;
+  for (Operation &op : block.getOperations()) {
+    auto transfer = dyn_cast<SpatDeferredCommunicationOp>(&op);
+    if (!transfer)
+      continue;
+    auto multiSourcePayload = transfer->getAttrOfType<BoolAttr>("multi_source_payload");
+    auto sourceOperandForScheduledLane =
+      transfer->getAttrOfType<DenseI64ArrayAttr>("source_operand_for_scheduled_lane");
+    if (multiSourcePayload && multiSourcePayload.getValue() && sourceOperandForScheduledLane) {
+      SmallVector<size_t> sourceOperandIndexes;
+      for (int64_t sourceOperandIndex : sourceOperandForScheduledLane.asArrayRef())
+        sourceOperandIndexes.push_back(static_cast<size_t>(sourceOperandIndex));
+      os << "    deferred input " << deferredInputIndex << ": multi-source uniqueSources="
+         << transfer.getSources().size() << " sourceOperandForScheduledLane=";
+      printIndexedList(os, ArrayRef<size_t>(sourceOperandIndexes));
+      os << "\n";
+    }
+    deferredInputIndex++;
+  }
+}
+
+static void dumpPeftMaterializationReport(ModuleOp moduleOp,
+                                          func::FuncOp funcOp,
+                                          const MergeScheduleResult &schedule,
+                                          const llvm::MapVector<size_t, PeftClassPlan> &peftClassPlans,
+                                          ArrayRef<ScheduledMaterializationRecord> materializedSchedules,
+                                          ScheduledComputePrintContext &printContext) {
+  std::fstream file = openDialectDumpFileWithExtension("spatial2_scheduled_no_comm", "/reports", "txt");
+  if (!file.is_open())
+    return;
+
+  llvm::raw_os_ostream os(file);
+  AsmState &asmState = printContext.asmState;
+  PeftMaterializationReportSummary summary =
+    buildPeftMaterializationReportSummary(funcOp, schedule, peftClassPlans, materializedSchedules);
+
+  os << "Summary\n";
+  os << "=======\n";
+  os << "Graph computes:\n";
+  os << "  total: " << (summary.scalarGraphCompute + summary.graphComputeBatchOps) << "\n";
+  os << "  scalar graph_compute: " << summary.scalarGraphCompute << "\n";
+  os << "  graph_compute_batch: " << summary.graphComputeBatchOps << "\n";
+  os << "Compute instances:\n";
+  os << "  total: " << (summary.scalarGraphComputeInstances + summary.graphComputeBatchInstances) << "\n";
+  os << "  scalar graph_compute instances: " << summary.scalarGraphComputeInstances << "\n";
+  os << "  graph_compute_batch instances: " << summary.graphComputeBatchInstances << "\n";
+  os << "PEFT classes:\n";
+  os << "  total: " << summary.peftClasses << "\n";
+  os << "  single-cpu: " << summary.singleCpuClasses << "\n";
+  os << "  multi-cpu: " << summary.multiCpuClasses << "\n";
+  os << "Scheduled ops:\n";
+  os << "  total: " << (summary.scheduledCompute + summary.scheduledComputeBatch) << "\n";
+  os << "  scheduled_compute: " << summary.scheduledCompute << "\n";
+  os << "  scheduled_compute_batch: " << summary.scheduledComputeBatch << "\n";
+  os << "Deferred communications:\n";
+  os << "  total: " << summary.deferredCommunication << "\n";
+  os << "  scalar metadata: " << summary.deferredCommunicationScalarMetadata << "\n";
+  os << "  scheduled-lane metadata: " << summary.deferredCommunicationScheduledLaneMetadata << "\n";
+  os << "  multi-source payloads: " << summary.deferredCommunicationMultiSourcePayloads << "\n\n";
+
+  os << "PEFT Classes\n";
+  os << "============\n";
+  for (const auto &entry : peftClassPlans) {
+    const PeftClassPlan &peftClassPlan = entry.second;
+    os << "C" << peftClassPlan.canonicalClassId << " "
+       << (peftClassPlan.cpus.size() == 1 ? "single-cpu" : "multi-cpu") << " PEFT class\n";
+    if (peftClassPlan.cpus.size() == 1) {
+      size_t cpu = peftClassPlan.cpus.front();
+      os << "  cpu: " << cpu << "\n";
+      os << "  steps: " << peftClassPlan.instancesByCpu.lookup(cpu).size() << "\n";
+      for (auto [stepIndex, instance] : llvm::enumerate(peftClassPlan.instancesByCpu.lookup(cpu)))
+        os << "  step " << stepIndex << ": " << formatComputeInstanceForReport(instance, asmState) << "\n";
+    } else {
+      os << "  scheduled lanes: " << peftClassPlan.cpus.size() << "\n";
+      os << "  steps: " << peftClassPlan.instancesByCpu.lookup(peftClassPlan.cpus.front()).size() << "\n";
+      os << "  cpus by scheduled lane:\n";
+      os << "    ";
+      printIndexedList(os, ArrayRef<size_t>(peftClassPlan.cpus));
+      os << "\n";
+      os << "  step sources:\n";
+      for (auto [stepIndex, stepTuple] : llvm::enumerate(buildComputeStepTuples(peftClassPlan)))
+        os << "    step " << stepIndex << ": "
+           << formatGraphComputeBlockKey(getGraphComputeBlockKey(stepTuple.instances.front()), asmState) << "\n";
+    }
+    os << "\n";
+  }
+
+  os << "Materialized Scheduled Ops\n";
+  os << "=========================\n";
+  for (const ScheduledMaterializationRecord &record : materializedSchedules) {
+    os << "C" << record.canonicalPeftClassId << " -> " << formatOperationLabel(record.scheduledOp, asmState) << " "
+       << record.scheduledOp->getName().getStringRef() << "\n";
+    os << "  kind: "
+       << (isa<SpatScheduledCompute>(record.scheduledOp) ? "single-cpu scheduled_compute"
+                                                         : "multi-cpu scheduled_compute_batch")
+       << "\n";
+    if (isa<SpatScheduledCompute>(record.scheduledOp))
+      os << "  cpu: " << record.cpus.front() << "\n";
+    else
+      os << "  scheduled lanes: " << record.cpus.size() << "\n";
+    os << "  results: " << record.scheduledOp->getNumResults() << "\n";
+    os << "  steps: "
+       << (isa<SpatScheduledCompute>(record.scheduledOp)
+             ? peftClassPlans.lookup(record.canonicalPeftClassId).instancesByCpu.lookup(record.cpus.front()).size()
+             : record.stepPlans.size())
+       << "\n";
+    if (isa<SpatScheduledComputeBatch>(record.scheduledOp)) {
+      os << "  cpus by scheduled lane:\n";
+      os << "    ";
+      printIndexedList(os, ArrayRef<size_t>(record.cpus));
+      os << "\n\n";
+    }
+    if (isa<SpatScheduledCompute>(record.scheduledOp)) {
+      const PeftClassPlan &peftClassPlan = peftClassPlans.lookup(record.canonicalPeftClassId);
+      size_t cpu = peftClassPlan.cpus.front();
+      size_t resultOffset = 0;
+      for (auto [stepIndex, instance] : llvm::enumerate(peftClassPlan.instancesByCpu.lookup(cpu))) {
+        size_t resultCount = getComputeInstanceResultValueCount(instance);
+        os << "  step " << stepIndex << " " << formatStepResultRange(resultOffset, resultCount) << " "
+           << formatComputeInstanceForReport(instance, asmState) << "\n";
+        resultOffset += resultCount;
+      }
+    } else {
+      auto scheduledBatch = cast<SpatScheduledComputeBatch>(record.scheduledOp);
+      for (auto [stepIndex, stepPlan] : llvm::enumerate(record.stepPlans)) {
+        const ComputeInstance &representative = stepPlan.stepTuple.instances.front();
+        SmallVector<uint32_t> sourceLaneStarts = collectSourceLaneStarts(stepPlan.stepTuple);
+        os << "  step " << stepIndex << " " << formatStepResultRange(stepPlan.resultOffset, stepPlan.resultCount) << " "
+           << formatGraphComputeBlockKey(getGraphComputeBlockKey(representative), asmState)
+           << " lanesPerScheduledLane=" << representative.laneCount << " sourceLaneSelector="
+           << (usesAffineSourceLaneMapping(stepPlan.stepTuple) ? "affine" : "table") << "\n";
+        os << "    source lanes by scheduled lane:\n";
+        os << "      ";
+        printIndexedList(os, ArrayRef<uint32_t>(sourceLaneStarts));
+        os << "\n";
+        Block &stepBlock = *std::next(scheduledBatch.getBody().begin(), stepIndex);
+        printMultiSourceDeferredInputs(os, stepBlock);
+      }
+    }
+    os << "\n";
+  }
+}
+
+
+} // namespace
+
+void dumpScheduledComputeReportAndModule(ModuleOp moduleOp,
+                                         func::FuncOp funcOp,
+                                         const MergeScheduleResult &schedule,
+                                         const llvm::MapVector<size_t, PeftClassPlan> &peftClassPlans,
+                                         ArrayRef<ScheduledMaterializationRecord> materializedSchedules) {
+  OpPrintingFlags flags;
+  flags.elideLargeElementsAttrs().enableDebugInfo(false, false).assumeVerified();
+  ScheduledComputePrintContext printContext(moduleOp, flags);
+  dumpPeftMaterializationReport(moduleOp, funcOp, schedule, peftClassPlans, materializedSchedules, printContext);
+
+  std::fstream file = openDialectDumpFileWithExtension("spatial2_scheduled_no_comm", "/dialects", "mlir");
+  if (!file.is_open())
+    return;
+  llvm::raw_os_ostream os(file);
+  moduleOp.getOperation()->print(os, printContext.asmState);
+  os.flush();
+}
+
+} // namespace spatial
+} // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeReport.hpp

@@ -0,0 +1,21 @@
+#pragma once
+
+#include "ScheduledComputeMaterialization.hpp"
+
+namespace onnx_mlir {
+namespace spatial {
+
+LogicalResult verifyPeftMaterializationReportSummary(
+  func::FuncOp funcOp,
+  const MergeScheduleResult &schedule,
+  const llvm::MapVector<size_t, PeftClassPlan> &peftClassPlans,
+  ArrayRef<ScheduledMaterializationRecord> materializedSchedules);
+
+void dumpScheduledComputeReportAndModule(ModuleOp moduleOp,
+                                         func::FuncOp funcOp,
+                                         const MergeScheduleResult &schedule,
+                                         const llvm::MapVector<size_t, PeftClassPlan> &peftClassPlans,
+                                         ArrayRef<ScheduledMaterializationRecord> materializedSchedules);
+
+} // namespace spatial
+} // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp

@@ -0,0 +1,316 @@
+#include "ScheduledComputeVerification.hpp"
+
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
+
+namespace onnx_mlir {
+namespace spatial {
+
+using namespace mlir;
+
+LogicalResult verifyMaterializedScheduleMapping(
+  func::FuncOp funcOp,
+  const MergeScheduleResult &schedule,
+  const llvm::MapVector<size_t, PeftClassPlan> &peftClassPlans,
+  const DenseMap<GraphComputeBlockKey, Block *> &graphComputeToBlockMap,
+  ArrayRef<ScheduledMaterializationRecord> materializedSchedules) {
+  pim::CappedDiagnosticReporter diagnostics;
+  size_t expectedClassCount = countPeftEquivalenceClasses(schedule);
+  if (expectedClassCount != materializedSchedules.size()) {
+    diagnostics.report(funcOp.getOperation(), [&](Operation *illegalOp) {
+      illegalOp->emitOpError() << "phase-check expected " << expectedClassCount
+                               << " PEFT equivalence classes but materialized " << materializedSchedules.size()
+                               << " scheduled computes";
+    });
+  }
+
+  llvm::SmallDenseSet<size_t, 16> seenClasses;
+  for (const ScheduledMaterializationRecord &record : materializedSchedules) {
+    if (!seenClasses.insert(record.canonicalPeftClassId).second) {
+      diagnostics.report(record.scheduledOp, [&](Operation *illegalOp) {
+        illegalOp->emitOpError("phase-check multiple scheduled ops own the same PEFT equivalence class");
+      });
+    }
+    if (!peftClassPlans.count(record.canonicalPeftClassId)) {
+      diagnostics.report(record.scheduledOp, [&](Operation *illegalOp) {
+        illegalOp->emitOpError("phase-check scheduled op refers to a missing PEFT materialization class");
+      });
+    }
+  }
+
+  for (const auto &entry : peftClassPlans) {
+    if (!seenClasses.count(entry.first)) {
+      diagnostics.report(funcOp.getOperation(), [&](Operation *illegalOp) {
+        illegalOp->emitOpError() << "phase-check PEFT equivalence class " << entry.first
+                                 << " was not materialized by any scheduled op";
+      });
+    }
+  }
+
+  for (GraphComputeBlockKey key : collectExpectedGraphComputeBlockKeys(funcOp)) {
+    if (graphComputeToBlockMap.count(key))
+      continue;
+    diagnostics.report(key.op, [&](Operation *illegalOp) {
+      illegalOp->emitOpError() << "phase-check graph compute is missing a scheduled MLIR block mapping for lanes ["
+                               << key.laneStart << ":" << (key.laneStart + key.laneCount) << "]";
+    });
+  }
+
+  for (const auto &entry : graphComputeToBlockMap) {
+    Block *block = entry.second;
+    if (!block || !isa<SpatScheduledCompute, SpatScheduledComputeBatch>(block->getParentOp())) {
+      diagnostics.report(entry.first.op, [&](Operation *illegalOp) {
+        illegalOp->emitOpError("phase-check graph compute block mapping does not target a scheduled compute block");
+      });
+    }
+  }
+
+  if (graphComputeToBlockMap.size() != collectExpectedGraphComputeBlockKeys(funcOp).size()) {
+    diagnostics.report(funcOp.getOperation(), [&](Operation *illegalOp) {
+      illegalOp->emitOpError() << "phase-check expected "
+                               << collectExpectedGraphComputeBlockKeys(funcOp).size()
+                               << " graph compute block mappings but saw " << graphComputeToBlockMap.size();
+    });
+  }
+
+  diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial materialization verification failed");
+  return success(!diagnostics.hasFailure());
+}
+
+LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp) {
+  pim::CappedDiagnosticReporter diagnostics;
+  funcOp.walk([&](SpatDeferredCommunicationOp transfer) {
+    for (Value source : transfer.getSources()) {
+      if (isa_and_nonnull<SpatScheduledCompute, SpatScheduledComputeBatch>(source.getDefiningOp())) {
+        diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) {
+          illegalOp->emitOpError("phase-check deferred communication source operand must remain an original graph SSA producer value");
+        });
+      }
+    }
+
+    if (auto scheduled = dyn_cast<SpatScheduledComputeBatch>(transfer->getParentOp())) {
+      auto batchedAttr = transfer->getAttrOfType<BoolAttr>("batched");
+      auto sourceOperandForScheduledLane =
+        transfer->getAttrOfType<DenseI64ArrayAttr>("source_operand_for_scheduled_lane");
+      auto multiSourcePayload = transfer->getAttrOfType<BoolAttr>("multi_source_payload");
+      auto sourceCpus = transfer->getAttrOfType<DenseI64ArrayAttr>("source_cpus");
+      auto sourceClasses = transfer->getAttrOfType<DenseI64ArrayAttr>("source_classes");
+      auto sourceLaneRanges = transfer->getAttrOfType<DenseI64ArrayAttr>("source_lane_ranges");
+      auto targetCpus = transfer->getAttrOfType<DenseI64ArrayAttr>("target_cpus");
+      auto targetClasses = transfer->getAttrOfType<DenseI64ArrayAttr>("target_classes");
+      auto targetLaneRanges = transfer->getAttrOfType<DenseI64ArrayAttr>("target_lane_ranges");
+      int64_t laneCount = scheduled.getLaneCount();
+      auto hasExpectedSize = [&](DenseI64ArrayAttr attr, int64_t expected, StringRef name) {
+        if (!attr || attr.size() != expected)
+          diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) {
+            illegalOp->emitOpError() << "phase-check scheduled batch deferred communication requires " << name
+                                     << " length " << expected;
+          });
+      };
+      if (!batchedAttr || !batchedAttr.getValue())
+        diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) {
+          illegalOp->emitOpError("phase-check scheduled batch deferred communication requires batched metadata");
+        });
+      hasExpectedSize(sourceCpus, laneCount, "source_cpus");
+      hasExpectedSize(sourceClasses, laneCount, "source_classes");
+      hasExpectedSize(targetCpus, laneCount, "target_cpus");
+      hasExpectedSize(targetClasses, laneCount, "target_classes");
+      hasExpectedSize(sourceLaneRanges, laneCount * 2, "source_lane_ranges");
+      hasExpectedSize(targetLaneRanges, laneCount * 2, "target_lane_ranges");
+      hasExpectedSize(sourceOperandForScheduledLane, laneCount, "source_operand_for_scheduled_lane");
+      if (!multiSourcePayload || multiSourcePayload.getValue() != (transfer.getSources().size() > 1))
+        diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) {
+          illegalOp->emitOpError("phase-check scheduled batch deferred communication has inconsistent multi_source_payload metadata");
+        });
+      if (sourceOperandForScheduledLane) {
+        for (int64_t sourceOperandIndex : sourceOperandForScheduledLane.asArrayRef()) {
+          if (sourceOperandIndex < 0 || sourceOperandIndex >= static_cast<int64_t>(transfer.getSources().size())) {
+            diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) {
+              illegalOp->emitOpError(
+                "phase-check scheduled batch deferred communication source_operand_for_scheduled_lane index is out of range");
+            });
+            break;
+          }
+        }
+      }
+    }
+  });
+
+  diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial deferred communication verification failed");
+  return success(!diagnostics.hasFailure());
+}
+
+LogicalResult verifyMultiCpuStepResultRouting(SpatScheduledComputeBatch scheduled,
+                                                     ArrayRef<ScheduledStepPlan> stepPlans) {
+  pim::CappedDiagnosticReporter diagnostics;
+  unsigned resultArgBase = getScheduledBatchResultArgBase(scheduled);
+  if (scheduled.getBody().getBlocks().size() != stepPlans.size()) {
+    diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) {
+      illegalOp->emitOpError() << "scheduled batch step routing expected " << stepPlans.size()
+                               << " blocks but saw " << scheduled.getBody().getBlocks().size();
+    });
+    diagnostics.emitSuppressedSummary(scheduled.getOperation(),
+                                      "scheduled batch step routing verification failed");
+    return failure();
+  }
+
+  SmallVector<unsigned> globalResultWrites(scheduled.getNumResults(), 0);
+  size_t stepIndex = 0;
+  for (Block &block : scheduled.getBody().getBlocks()) {
+    const ScheduledStepPlan &stepPlan = stepPlans[stepIndex++];
+    SmallVector<bool> localWrites(stepPlan.resultCount, false);
+    auto inParallel = dyn_cast<SpatInParallelOp>(block.getTerminator());
+    if (!inParallel) {
+      diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) {
+        illegalOp->emitOpError() << "scheduled batch step " << stepPlan.stepIndex
+                                 << " is missing spat.in_parallel";
+      });
+      continue;
+    }
+
+    for (Operation &op : inParallel.getRegion().front()) {
+      auto insert = dyn_cast<tensor::ParallelInsertSliceOp>(&op);
+      if (!insert)
+        continue;
+      auto dest = dyn_cast<BlockArgument>(insert.getDest());
+      if (!dest || dest.getOwner() != &block) {
+        diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) {
+          illegalOp->emitOpError() << "scheduled batch step " << stepPlan.stepIndex
+                                   << " writes to a non-block result destination";
+        });
+        continue;
+      }
+      unsigned resultIndex = dest.getArgNumber() - resultArgBase;
+      if (dest.getArgNumber() < resultArgBase || resultIndex >= scheduled.getNumResults()) {
+        diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) {
+          illegalOp->emitOpError() << "scheduled batch step " << stepPlan.stepIndex
+                                   << " writes to invalid result block argument " << dest.getArgNumber();
+        });
+        continue;
+      }
+      if (resultIndex < stepPlan.resultOffset
+          || resultIndex >= stepPlan.resultOffset + stepPlan.resultCount) {
+        diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) {
+          illegalOp->emitOpError() << "scheduled batch step " << stepPlan.stepIndex
+                                   << " expected result range [" << stepPlan.resultOffset << ":"
+                                   << (stepPlan.resultOffset + stepPlan.resultCount)
+                                   << ") but wrote result " << resultIndex;
+        });
+        continue;
+      }
+      localWrites[resultIndex - stepPlan.resultOffset] = true;
+      globalResultWrites[resultIndex]++;
+    }
+
+    for (size_t index = 0; index < localWrites.size(); ++index)
+      if (!localWrites[index])
+        diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) {
+          illegalOp->emitOpError() << "scheduled batch step " << stepPlan.stepIndex
+                                   << " did not write expected result " << (stepPlan.resultOffset + index);
+        });
+  }
+
+  for (size_t resultIndex = 0; resultIndex < globalResultWrites.size(); ++resultIndex)
+    if (globalResultWrites[resultIndex] != 1)
+      diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) {
+        illegalOp->emitOpError() << "scheduled batch result " << resultIndex << " expected one producing step but saw "
+                                 << globalResultWrites[resultIndex];
+      });
+
+  diagnostics.emitSuppressedSummary(scheduled.getOperation(), "scheduled batch step routing verification failed");
+  return success(!diagnostics.hasFailure());
+}
+
+LogicalResult verifyMultiCpuLocalFragmentOffsets(SpatScheduledComputeBatch scheduled) {
+  pim::CappedDiagnosticReporter diagnostics;
+  unsigned resultArgBase = getScheduledBatchResultArgBase(scheduled);
+  for (auto enumeratedBlock : llvm::enumerate(scheduled.getBody().getBlocks())) {
+    size_t stepIndex = enumeratedBlock.index();
+    Block &block = enumeratedBlock.value();
+    Value scheduledLane = block.getArgument(0);
+    auto inParallel = dyn_cast<SpatInParallelOp>(block.getTerminator());
+    if (!inParallel) {
+      diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) {
+        illegalOp->emitOpError() << "phase-check scheduled batch step " << stepIndex
+                                 << " is missing spat.in_parallel";
+      });
+      continue;
+    }
+    auto isFinalScheduledOutputInsert = [&](Operation *op) {
+      auto insert = dyn_cast<tensor::ParallelInsertSliceOp>(op);
+      if (!insert || op->getParentOp() != inParallel.getOperation())
+        return false;
+      auto dest = dyn_cast<BlockArgument>(insert.getDest());
+      return dest && dest.getOwner() == &block && dest.getArgNumber() >= resultArgBase;
+    };
+
+    block.walk([&](Operation *op) {
+      if (op == block.getTerminator())
+        return;
+      if (isFinalScheduledOutputInsert(op)) {
+        if (scheduled.getLaneCount() > 1) {
+          auto insert = cast<tensor::ParallelInsertSliceOp>(op);
+          bool dependsOnScheduledLane = false;
+          for (OpFoldResult offset : insert.getMixedOffsets()) {
+            if (auto value = dyn_cast<Value>(offset); value && valueTransitivelyDependsOn(value, scheduledLane)) {
+              dependsOnScheduledLane = true;
+              break;
+            }
+          }
+          if (!dependsOnScheduledLane)
+            diagnostics.report(insert.getOperation(), [&](Operation *illegalOp) {
+              illegalOp->emitOpError(
+                "phase-check scheduled batch final output insert must be indexed by scheduled lane");
+            });
+        }
+        return;
+      }
+
+      auto insertSlice = dyn_cast<tensor::InsertSliceOp>(op);
+      if (!insertSlice)
+        return;
+      auto dest = dyn_cast<BlockArgument>(insertSlice.getDest());
+      if (dest && dest.getOwner() == &block && dest.getArgNumber() >= resultArgBase)
+        return;
+
+      auto destType = dyn_cast<RankedTensorType>(insertSlice.getDestType());
+      if (!destType || !destType.hasStaticShape() || destType.getRank() == 0)
+        return;
+
+      for (OpFoldResult offset : insertSlice.getMixedOffsets()) {
+        auto value = dyn_cast<Value>(offset);
+        if (!value)
+          continue;
+        if (!valueTransitivelyDependsOn(value, scheduledLane))
+          continue;
+        diagnostics.report(insertSlice.getOperation(), [&](Operation *illegalOp) {
+          illegalOp->emitOpError()
+            << "phase-check scheduled batch local fragment insert offset must use the source-instance inner lane, not the scheduled lane"
+            << " step " << stepIndex;
+        });
+        break;
+      }
+    });
+  }
+
+  diagnostics.emitSuppressedSummary(scheduled.getOperation(),
+                                    "scheduled batch local fragment offset verification failed");
+  return success(!diagnostics.hasFailure());
+}
+
+
+LogicalResult verifyScheduledMaterializationRecords(ArrayRef<ScheduledMaterializationRecord> materializedSchedules) {
+  for (const ScheduledMaterializationRecord &record : materializedSchedules) {
+    auto scheduled = dyn_cast<SpatScheduledComputeBatch>(record.scheduledOp);
+    if (!scheduled)
+      continue;
+    if (failed(verifyMultiCpuStepResultRouting(scheduled, record.stepPlans)))
+      return failure();
+    if (failed(verifyMultiCpuLocalFragmentOffsets(scheduled)))
+      return failure();
+  }
+  return success();
+}
+
+} // namespace spatial
+} // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/ScheduledComputeVerification.hpp

@@ -0,0 +1,22 @@
+#pragma once
+
+#include "ScheduledComputeMaterialization.hpp"
+
+namespace onnx_mlir {
+namespace spatial {
+
+LogicalResult verifyMaterializedScheduleMapping(
+  func::FuncOp funcOp,
+  const MergeScheduleResult &schedule,
+  const llvm::MapVector<size_t, PeftClassPlan> &peftClassPlans,
+  const DenseMap<GraphComputeBlockKey, Block *> &graphComputeToBlockMap,
+  ArrayRef<ScheduledMaterializationRecord> materializedSchedules);
+
+LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp);
+LogicalResult verifyMultiCpuStepResultRouting(SpatScheduledComputeBatch scheduled,
+                                              ArrayRef<ScheduledStepPlan> stepPlans);
+LogicalResult verifyMultiCpuLocalFragmentOffsets(SpatScheduledComputeBatch scheduled);
+LogicalResult verifyScheduledMaterializationRecords(ArrayRef<ScheduledMaterializationRecord> materializedSchedules);
+
+} // namespace spatial
+} // namespace onnx_mlir

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp

@@ -12,6 +12,7 @@
 #include "llvm/Support/Casting.h"
 
 #include <algorithm>
+#include <cmath>
 #include <iterator>
 #include <limits>
 #include <optional>
@@ -21,6 +22,7 @@
 
 #include "ComputeGraph.hpp"
 #include "ComputeInstanceUtils.hpp"
+#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Support/TypeUtilities.hpp"
@@ -35,9 +37,223 @@ uint64_t countComputeBodyOperationInstances(Region& body);
 
 namespace {
 
-Cost getComputeBodyCost(Region& body) {
-  constexpr Cost kOperationCost = 100;
-  return checkedMultiply(static_cast<Cost>(countComputeBodyOperationInstances(body)), kOperationCost);
+struct PimsimSchedulerCostModel {
+  static constexpr Cost kDefaultBitwidth = 8;
+  static constexpr Cost kCorePeriodNs = 1;
+  static constexpr Cost kLocalMemoryWidthBytes = 64;
+  static constexpr Cost kLocalMemoryLatencyCycles = 1;
+  static constexpr Cost kNetworkBusWidthBytes = 8;
+  static constexpr Cost kNetworkBaseLatencyNs = 2;
+  static constexpr Cost kNetworkPerHopLatencyNs = 1;
+  static constexpr Cost kVectorWidth = 16;
+  static constexpr Cost kVectorLatencyCycles = 4;
+  static constexpr Cost kDacResolutionBits = 1;
+  static constexpr Cost kDacLatencyCycles = 1;
+  static constexpr Cost kDacCount = 128;
+  static constexpr Cost kXbarReadLatencyNs = 30;
+  static constexpr Cost kSampleHoldLatencyCycles = 1;
+  static constexpr Cost kAdcLatencyCycles = 10;
+  static constexpr Cost kAdcCount = 2;
+  static constexpr Cost kShiftAdderLatencyCycles = 1;
+  static constexpr Cost kOutputBufferLatencyCycles = 1;
+  static constexpr Cost kInputBufferLatencyCycles = 0;
+  static constexpr Cost kFallbackOperationCost = 1;
+
+  static Cost ceilDiv(Cost numerator, Cost denominator) {
+    assert(denominator > 0 && "denominator must be positive");
+    return (numerator + denominator - 1) / denominator;
+  }
+
+  static std::optional<Cost> getStaticElementCount(Type type) {
+    auto shaped = dyn_cast<ShapedType>(type);
+    if (!shaped || !shaped.hasStaticShape())
+      return std::nullopt;
+    return static_cast<Cost>(shaped.getNumElements());
+  }
+
+  static Cost getBitwidthOrDefault(Type type) {
+    if (auto shaped = dyn_cast<ShapedType>(type))
+      type = shaped.getElementType();
+    if (auto intType = dyn_cast<IntegerType>(type))
+      return intType.getWidth();
+    if (auto floatType = dyn_cast<FloatType>(type))
+      return floatType.getWidth();
+    if (isa<IndexType>(type))
+      return 64;
+    return kDefaultBitwidth;
+  }
+
+  static Cost getComputeBitwidth(Type type) {
+    return std::min(getBitwidthOrDefault(type), kDefaultBitwidth);
+  }
+
+  static Cost getByteSize(Type type, Cost fallbackBitwidth = kDefaultBitwidth) {
+    auto elementCount = getStaticElementCount(type);
+    if (!elementCount)
+      return kFallbackOperationCost;
+    Cost bitwidth = fallbackBitwidth;
+    if (bitwidth <= 0)
+      bitwidth = getBitwidthOrDefault(type);
+    if (bitwidth <= 0)
+      bitwidth = kDefaultBitwidth;
+    return ceilDiv(checkedMultiply(*elementCount, bitwidth), static_cast<Cost>(8));
+  }
+
+  static Cost getVectorReadWriteCost(Cost readBytes, Cost writeBytes) {
+    Cost totalBytes = checkedAdd(readBytes, writeBytes);
+    return checkedMultiply(ceilDiv(totalBytes, kLocalMemoryWidthBytes),
+                           checkedMultiply(kLocalMemoryLatencyCycles, kCorePeriodNs));
+  }
+
+  static Cost getVectorComputeCost(Cost elementCount) {
+    return checkedMultiply(ceilDiv(elementCount, kVectorWidth),
+                           checkedMultiply(kVectorLatencyCycles, kCorePeriodNs));
+  }
+
+  static Cost getTensorMoveCost(Type type) {
+    return getVectorReadWriteCost(getByteSize(type), 0);
+  }
+
+  static std::pair<Cost, Cost> estimateMeshShape() {
+    Cost coreCount = static_cast<Cost>(std::max<long>(1, coresCount.getValue()));
+    Cost rows = static_cast<Cost>(std::sqrt(static_cast<long double>(coreCount)));
+    if (rows == 0)
+      rows = 1;
+    while (rows > 1 && coreCount % rows != 0)
+      --rows;
+    Cost cols = ceilDiv(coreCount, rows);
+    return {rows, cols};
+  }
+
+  static Cost getAverageInterCoreLatencyNs() {
+    auto [rows, cols] = estimateMeshShape();
+    auto averageAxisDistance = [](Cost size) -> Cost {
+      if (size <= 1)
+        return 0;
+      return checkedMultiply(size, size) - 1;
+    };
+    Cost avgRow = averageAxisDistance(rows) / (static_cast<Cost>(3) * rows);
+    Cost avgCol = averageAxisDistance(cols) / (static_cast<Cost>(3) * cols);
+    return checkedAdd(kNetworkBaseLatencyNs, checkedMultiply(kNetworkPerHopLatencyNs, checkedAdd(avgRow, avgCol)));
+  }
+
+  static Cost getInterCoreTransferCostFromBytes(Cost bytes) {
+    Cost localRead = checkedMultiply(ceilDiv(bytes, kLocalMemoryWidthBytes),
+                                     checkedMultiply(kLocalMemoryLatencyCycles, kCorePeriodNs));
+    Cost localWrite = checkedMultiply(ceilDiv(bytes, kLocalMemoryWidthBytes),
+                                      checkedMultiply(kLocalMemoryLatencyCycles, kCorePeriodNs));
+    Cost payloadFlits = ceilDiv(bytes, kNetworkBusWidthBytes);
+    Cost averageNoCLatency = getAverageInterCoreLatencyNs();
+    Cost network = checkedMultiply(checkedAdd(static_cast<Cost>(2), payloadFlits), averageNoCLatency);
+    return checkedAdd(checkedAdd(localRead, localWrite), network);
+  }
+
+  static Cost getUnaryVectorCost(Type inputType, Type outputType, bool scalarOutput = false) {
+    auto maybeElements = getStaticElementCount(inputType);
+    if (!maybeElements)
+      return kFallbackOperationCost;
+    Cost inputBytes = getByteSize(inputType, getComputeBitwidth(inputType));
+    Cost outputBytes = scalarOutput ? ceilDiv(getComputeBitwidth(outputType), static_cast<Cost>(8))
+                                    : getByteSize(outputType, getComputeBitwidth(outputType));
+    return checkedAdd(getVectorReadWriteCost(inputBytes, outputBytes), getVectorComputeCost(*maybeElements));
+  }
+
+  static Cost getBinaryVectorCost(Type lhsType, Type rhsType, Type outputType, bool scalarOutput = false) {
+    auto maybeElements = getStaticElementCount(lhsType);
+    if (!maybeElements)
+      return kFallbackOperationCost;
+    Cost readBytes = checkedAdd(getByteSize(lhsType, getComputeBitwidth(lhsType)),
+                                getByteSize(rhsType, getComputeBitwidth(rhsType)));
+    Cost outputBytes = scalarOutput ? ceilDiv(getComputeBitwidth(outputType), static_cast<Cost>(8))
+                                    : getByteSize(outputType, getComputeBitwidth(outputType));
+    return checkedAdd(getVectorReadWriteCost(readBytes, outputBytes), getVectorComputeCost(*maybeElements));
+  }
+
+  static Cost getMatrixComputeLatency(Cost inputBitwidth) {
+    Cost xbarDim = static_cast<Cost>(crossbarSize.getValue());
+    Cost inputTimes = ceilDiv(inputBitwidth, kDacResolutionBits);
+    Cost dacTimes = ceilDiv(xbarDim, kDacCount);
+    Cost adcTimes = ceilDiv(xbarDim, kAdcCount);
+    Cost frontStage = kInputBufferLatencyCycles + kDacLatencyCycles + kXbarReadLatencyNs + kSampleHoldLatencyCycles;
+    Cost backPipe = std::max(kAdcLatencyCycles, checkedAdd(kShiftAdderLatencyCycles, kOutputBufferLatencyCycles));
+    Cost backStage = checkedAdd(checkedAdd(kAdcLatencyCycles, kShiftAdderLatencyCycles), kOutputBufferLatencyCycles);
+    backStage = checkedAdd(backStage, checkedMultiply(adcTimes - 1, backPipe));
+    Cost totalTimes = checkedMultiply(inputTimes, dacTimes);
+    Cost stagePipe = std::max(frontStage, backStage);
+    return checkedAdd(checkedAdd(frontStage, backStage),
+                      checkedMultiply(totalTimes - 1, stagePipe));
+  }
+
+  static Cost getWvmmCost(Type inputType, Type outputType) {
+    Cost inputBitwidth = getComputeBitwidth(inputType);
+    Cost inputBytes = checkedMultiply(static_cast<Cost>(crossbarSize.getValue()),
+                                      ceilDiv(inputBitwidth, static_cast<Cost>(8)));
+    inputBytes = checkedMultiply(inputBytes, static_cast<Cost>(8));
+    Cost outputBytes = getByteSize(outputType, getComputeBitwidth(outputType));
+    return checkedAdd(getVectorReadWriteCost(inputBytes, outputBytes), getMatrixComputeLatency(inputBitwidth));
+  }
+};
+
+std::optional<uint64_t> getStaticTripCount(scf::ForOp loop);
+[[maybe_unused]] Cost getOperationCost(Operation& op);
+
+[[maybe_unused]] Cost getRegionCost(Region& body) {
+  Cost cost = 0;
+  for (Block& block : body)
+    for (Operation& op : block)
+      cost = checkedAdd(cost, getOperationCost(op));
+  return cost;
+}
+
+[[maybe_unused]] Cost getOperationCost(Operation& op) {
+  if (auto loop = dyn_cast<scf::ForOp>(&op)) {
+    std::optional<uint64_t> tripCount = getStaticTripCount(loop);
+    if (!tripCount)
+      return PimsimSchedulerCostModel::kFallbackOperationCost;
+    return checkedMultiply(getRegionCost(loop.getRegion()), static_cast<Cost>(*tripCount));
+  }
+
+  if (isa<SpatYieldOp, SpatInParallelOp, affine::AffineApplyOp, arith::ConstantOp,
+          tensor::EmptyOp, tensor::CollapseShapeOp, tensor::ExpandShapeOp>(&op))
+    return 0;
+
+  if (auto wvmm = dyn_cast<SpatVMMOp>(&op))
+    return PimsimSchedulerCostModel::getWvmmCost(wvmm.getInput().getType(), wvmm.getOutput().getType());
+  if (auto vvdmul = dyn_cast<SpatVVDMulOp>(&op))
+    return PimsimSchedulerCostModel::getBinaryVectorCost(
+      vvdmul.getLhs().getType(), vvdmul.getRhs().getType(), vvdmul.getOutput().getType(), /*scalarOutput=*/true);
+  if (auto vadd = dyn_cast<SpatVAddOp>(&op))
+    return PimsimSchedulerCostModel::getBinaryVectorCost(vadd.getLhs().getType(), vadd.getRhs().getType(),
+                                                         vadd.getOutput().getType());
+  if (auto vsub = dyn_cast<SpatVSubOp>(&op))
+    return PimsimSchedulerCostModel::getBinaryVectorCost(vsub.getLhs().getType(), vsub.getRhs().getType(),
+                                                         vsub.getOutput().getType());
+  if (auto vmul = dyn_cast<SpatVMulOp>(&op))
+    return PimsimSchedulerCostModel::getBinaryVectorCost(vmul.getLhs().getType(), vmul.getRhs().getType(),
+                                                         vmul.getOutput().getType());
+  if (auto vmax = dyn_cast<SpatVMaxOp>(&op))
+    return PimsimSchedulerCostModel::getBinaryVectorCost(vmax.getLhs().getType(), vmax.getRhs().getType(),
+                                                         vmax.getOutput().getType());
+  if (auto vavg = dyn_cast<SpatVAvgOp>(&op))
+    return PimsimSchedulerCostModel::getUnaryVectorCost(vavg.getInput().getType(), vavg.getOutput().getType(),
+                                                        /*scalarOutput=*/true);
+  if (auto relu = dyn_cast<SpatReluOp>(&op))
+    return PimsimSchedulerCostModel::getUnaryVectorCost(relu.getInput().getType(), relu.getOutput().getType());
+  if (auto sigm = dyn_cast<SpatSigmoidOp>(&op))
+    return PimsimSchedulerCostModel::getUnaryVectorCost(sigm.getInput().getType(), sigm.getOutput().getType());
+  if (auto softmax = dyn_cast<SpatSoftmaxOp>(&op)) {
+    Cost unary = PimsimSchedulerCostModel::getUnaryVectorCost(softmax.getInput().getType(), softmax.getOutput().getType());
+    return checkedMultiply(unary, static_cast<Cost>(4));
+  }
+  if (auto extract = dyn_cast<tensor::ExtractSliceOp>(&op))
+    return PimsimSchedulerCostModel::getTensorMoveCost(extract.getResult().getType());
+  if (auto insert = dyn_cast<tensor::InsertSliceOp>(&op))
+    return PimsimSchedulerCostModel::getTensorMoveCost(insert.getSource().getType());
+
+  Cost nestedCost = 0;
+  for (Region& region : op.getRegions())
+    nestedCost = checkedAdd(nestedCost, getRegionCost(region));
+  return checkedAdd(PimsimSchedulerCostModel::kFallbackOperationCost, nestedCost);
 }
 
 std::optional<uint64_t> getStaticTripCount(scf::ForOp loop) {
@@ -54,6 +270,11 @@ std::optional<uint64_t> getStaticTripCount(scf::ForOp loop) {
   return (distance + stride - 1) / stride;
 }
 
+Cost getComputeBodyCost(Region& body) {
+  constexpr Cost kOperationCost = 100;
+  return checkedMultiply(static_cast<Cost>(countComputeBodyOperationInstances(body)), kOperationCost);
+}
+
 uint64_t countOperationInstances(Operation& op) {
   if (auto loop = dyn_cast<scf::ForOp>(&op)) {
     std::optional<uint64_t> tripCount = getStaticTripCount(loop);
@@ -149,7 +370,8 @@ std::optional<Cost> getBatchProjectedInputTransferCost(SpatComputeBatch batch, V
     auto resultType = dyn_cast<ShapedType>(extract.getResult().getType());
     if (!resultType || !resultType.hasStaticShape())
       return std::nullopt;
-    projectedCost = checkedAdd(projectedCost, static_cast<Cost>(getSizeInBytes(resultType)));
+    projectedCost = checkedAdd(
+      projectedCost, PimsimSchedulerCostModel::getInterCoreTransferCostFromBytes(static_cast<Cost>(getSizeInBytes(resultType))));
   }
 
   if (projectedCost == 0)
@@ -162,7 +384,7 @@ Cost getInputTransferCost(const ComputeInstance& consumerInstance, Value input)
   if (auto batch = dyn_cast<SpatComputeBatch>(consumerInstance.op))
     if (std::optional<Cost> projectedCost = getBatchProjectedInputTransferCost(batch, input))
       return *projectedCost;
-  return static_cast<Cost>(getSizeInBytes(inputType));
+  return PimsimSchedulerCostModel::getInterCoreTransferCostFromBytes(static_cast<Cost>(getSizeInBytes(inputType)));
 }
 
 uint32_t getLaneOverlapCount(const ComputeInstance& lhs, const ComputeInstance& rhs) {
@@ -451,7 +673,7 @@ std::vector<ComputeGraphEdge> aggregateEdges(llvm::ArrayRef<ComputeGraphEdge> ed
       continue;
     auto inserted = edgeCosts.try_emplace({edge.source, edge.target}, edge.transferCost);
     if (!inserted.second)
-      inserted.first->second = std::max(inserted.first->second, edge.transferCost);
+      inserted.first->second = checkedAdd(inserted.first->second, edge.transferCost);
   }
 
   std::vector<ComputeGraphEdge> aggregatedEdges;

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp

@@ -106,10 +106,11 @@ static std::optional<uint32_t> getConstantExtractLane(tensor::ExtractSliceOp ext
 
 static std::optional<ProducerValueRef> getResultfulBatchProducerValueRef(SpatComputeBatch batch,
                                                                          const ComputeInstance* consumerInstance) {
-  if (!consumerInstance)
-    return std::nullopt;
-  if (!isa<SpatComputeBatch>(consumerInstance->op))
-    return std::nullopt;
+  if (!consumerInstance || !isa<SpatComputeBatch>(consumerInstance->op))
+    return ProducerValueRef {
+      {batch.getOperation(), 0, static_cast<uint32_t>(batch.getLaneCount())},
+      0
+    };
   if (consumerInstance->laneStart + consumerInstance->laneCount > static_cast<uint32_t>(batch.getLaneCount()))
     return std::nullopt;
   return ProducerValueRef {

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp

@@ -23,7 +23,10 @@ MergeSchedulerKind getSchedulerKind() {
   llvm_unreachable("unknown merge scheduler kind");
 }
 
-void verifySchedule(const ComputeGraph& graph, const MergeScheduleResult& result, unsigned long crossbarCapacity) {
+void verifySchedule(const ComputeGraph& graph,
+                    const MergeScheduleResult& result,
+                    unsigned long crossbarCapacity,
+                    size_t processorCount) {
   llvm::DenseMap<size_t, std::vector<std::pair<size_t, size_t>>> tasksByCpu;
   tasksByCpu.reserve(result.cpuToLastComputeMap.size());
 
@@ -79,7 +82,8 @@ void verifySchedule(const ComputeGraph& graph, const MergeScheduleResult& result
 
     Time earliestTargetStart = addOrMax(sourceStart, graph.nodes[edge.source].cost);
     if (sourceCpu != targetCpu)
-      earliestTargetStart = addOrMax(earliestTargetStart, edge.transferCost);
+      earliestTargetStart = addOrMax(
+        earliestTargetStart, getPeftTransferTime(edge.transferCost, sourceCpu, targetCpu, processorCount));
     if (targetStart < earliestTargetStart) {
       std::string message = llvm::formatv("merge scheduling: dependency legality failed between tasks {0} and {1}",
                                           graph.nodes[edge.source].originalOrder,
@@ -115,7 +119,10 @@ MergeScheduleResult MergeSchedulingAnalysis::run() {
                                                      static_cast<unsigned long>(crossbarCountInCore.getValue()),
                                                      entryOp->getContext()});
   }
-  verifySchedule(graph, schedule, static_cast<unsigned long>(crossbarCountInCore.getValue()));
+  verifySchedule(graph,
+                 schedule,
+                 static_cast<unsigned long>(crossbarCountInCore.getValue()),
+                 options.processorCount);
   return schedule;
 }