Better implementation of MergeComputeNodesPass

AGENTS.md

@@ -1,4 +1,5 @@
 * Always read the full README.md before doing anything
+* Always read the full invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md before modifying Spatial graph IR, Blueprint handling, or MergeComputeNodes.
 * Build commands:
   * `cmake --build ./build_release`
   * `cmake --build ./build_debug`

README.md

@@ -105,6 +105,9 @@ Pass these to `onnx-mlir` when compiling for PIM:
   the codegen tail.
 - `--pim-emit-json` - also emit `core_*.json` instruction files alongside
   `core_*.pim`.
+- `--pim-export-spatial-dataflow=<none|spatial1|spatial2|spatial3|all>` - control Spatial
+  dataflow CSV reports. The default `all` emits graph, scheduled, and realized
+  snapshots under `reports/`.
 - `--use-experimental-conv-impl` - use the alternate convolution lowering.
 - `--ignore-concat-error` - soft-fail a ConcatOp corner case.
 
@@ -167,7 +170,8 @@ Each validation run writes artifacts in the model workspace, for example under
 - `simulation/out.bin` - raw simulator output used for comparison.
 
 The compiler currently dumps dialect snapshots such as `spatial0.mlir`,
-`spatial1_dcp_merged.mlir`, `pim0.mlir`, `pim1_buff.mlir`,
+`spatial1_graph.mlir`, `spatial2_scheduled_no_comm.mlir`,
+`spatial3_scheduled.mlir`, `pim0.mlir`, `pim1_buff.mlir`,
 `pim2_coalesced.mlir`, and `pim3_folded.mlir` when an output directory is
 available.
 

invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md

@@ -0,0 +1,362 @@
+# Graph Compute Batch Physical-Fragment Invariant
+
+## Status
+
+This document is **normative** for Raptor's Spatial graph IR.
+
+Every developer or coding agent modifying Spatial graph construction, graph
+verification, Blueprint handling, or `MergeComputeNodes` must read this file
+after `README.md` and `AGENTS.md`.
+
+`AGENTS.md` must contain this instruction:
+
+```text
+* Always read the full invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md before modifying Spatial graph IR, Blueprint handling, or MergeComputeNodes.
+```
+
+## Scope
+
+This invariant applies to:
+
+- `spat.graph_compute_batch`;
+- graph-level values produced by it;
+- `tensor.parallel_insert_slice` operations that publish its lane results;
+- `spat.blueprint` operations that describe logical reconstruction;
+- graph analyses and transformations that consume those values;
+- the graph-to-scheduled transition in `MergeComputeNodes`.
+
+It does **not** impose the same representation on:
+
+- `spat.scheduled_compute`;
+- `spat.scheduled_compute_batch`;
+- `pim.core` or `pim.core_batch`;
+- values whose cross-core movement is already represented by explicit
+  `spat.channel_send` and `spat.channel_receive` operations.
+
+Scheduled IR represents execution on assigned cores. Communication and value
+availability there are defined by local SSA forwarding and explicit
+send/receive operations, not by the graph physical-fragment invariant.
+
+## Core invariant
+
+For every result of a `spat.graph_compute_batch` with `N` graph lanes:
+
+1. Every graph lane produces exactly one fragment for that result.
+2. All lanes produce fragments with the same exact ranked tensor type `F`.
+3. The graph result is a physical collection of those fragments with type:
+
+   ```text
+   tensor<N x shape(F) x element-type(F)>
+   ```
+
+   Conceptually, the result is `N × F`: one leading physical fragment-slot
+   dimension followed by the complete per-lane fragment shape.
+4. Physical slot `i` identifies a fragment publication. It does not, by itself,
+   identify a row, column, channel, tile, or any other logical tensor position.
+5. The result type carries no logical reconstruction order.
+
+The leading dimension is therefore a **physical fragment-slot dimension**, not
+a logical tensor dimension.
+
+## Per-lane computation is unrestricted
+
+The invariant constrains the published result representation, not what a lane
+may compute.
+
+A graph lane may:
+
+- read several input slices;
+- perform reductions;
+- add or combine multiple columns;
+- execute matrix/vector operations;
+- produce a fragment that corresponds to any logical region;
+- participate in a multi-stage or logarithmic reduction tree implemented by
+  following `spat.graph_compute` or `spat.graph_compute_batch` operations.
+
+Arithmetic combination is graph computation. `spat.blueprint` is not an
+arithmetic reduction operation.
+
+### Example: `16×4 -> 16×2`
+
+Two graph lanes may compute:
+
+```text
+lane 0: input[:, 0] + input[:, 1] -> tensor<16x1>
+lane 1: input[:, 2] + input[:, 3] -> tensor<16x1>
+```
+
+The physical graph result is:
+
+```text
+tensor<2x16x1>
+```
+
+A Blueprint then maps:
+
+```text
+physical slot 0 -> logical output[:, 0:1]
+physical slot 1 -> logical output[:, 1:2]
+```
+
+and describes the logical result `tensor<16x2>`.
+
+For a larger reduction, following graph compute batches may reduce fragments in
+`ceil(log2(N))` stages. Every intermediate batch still publishes a physical
+`batch × fragment` collection.
+
+## Physical publication inside `spat.graph_compute_batch`
+
+The batch body must publish each lane's fragment into the physical result.
+
+For one result with fragment type `F`, the corresponding
+`tensor.parallel_insert_slice` must insert the fragment into one slot of the
+physical `N × F` destination:
+
+```text
+physical offsets = [slot, 0, 0, ...]
+physical sizes   = [1, shape(F)...]
+physical strides = [1, 1, 1, ...]
+```
+
+The slot may be the graph lane directly or a statically analyzable permutation
+of it. The insertion describes physical slot placement only. It must not use a
+logical output dimension as the physical batch dimension.
+
+For each graph result, the body must contain exactly one physical publication
+per graph lane. Since the body executes once per lane, this normally means one
+`tensor.parallel_insert_slice` operation targeting that result.
+
+## Logical reconstruction
+
+Logical reconstruction is separate from physical publication.
+
+The reconstruction descriptor defines, for every physical fragment slot:
+
+- which physical batch operand owns the fragment;
+- which physical slot contains it;
+- its destination offsets in the logical tensor;
+- its destination sizes;
+- its destination strides;
+- coverage and conflict policy where relevant.
+
+The persistent owner of this information is `spat.blueprint` or an equivalent
+explicit graph-level reconstruction operation.
+
+A logical consumer must not infer reconstruction from the physical tensor type
+or assume that physical slot order equals logical order.
+
+The logical mapping may be arbitrary. For example:
+
+```text
+physical slot 0 -> logical row 13
+physical slot 1 -> logical row 4
+physical slot 2 -> logical row 10
+```
+
+The physical result remains a regular `batch × fragment` tensor.
+
+## Relationship between `parallel_insert_slice` and Blueprint
+
+During graph construction, an algorithm may naturally describe logical
+placement with `tensor.parallel_insert_slice` geometry. Before the graph is in
+its canonical form:
+
+1. that geometry must be separated from physical fragment publication;
+2. the graph batch result must be normalized to `N × F`;
+3. the logical insertion geometry must be transferred to a persistent
+   `spat.blueprint` reconstruction descriptor.
+
+After normalization:
+
+- `parallel_insert_slice` inside `spat.graph_compute_batch` publishes into
+  physical fragment slots;
+- `spat.blueprint` describes reconstruction into the logical tensor.
+
+The original graph operation may be erased only after all reconstruction
+information needed by later stages has a persistent owner.
+
+## Blueprint semantics
+
+Blueprint is placement/reconstruction metadata. It may:
+
+- concatenate fragments;
+- reorder fragments;
+- insert fragments into arbitrary disjoint logical regions;
+- describe complete or partial logical coverage;
+- expose a logical tensor view when materialization is required.
+
+Blueprint must not silently perform arithmetic such as addition, multiplication,
+maximum, or reduction. Such transformations must be represented by following
+`spat.graph_compute` or `spat.graph_compute_batch` operations.
+
+A Blueprint consuming a physical fragment batch must explicitly identify the
+physical source slot for every logical fragment. It must not derive that slot
+from operand order unless that convention is explicitly represented and
+verified.
+
+## Multiple results
+
+A `spat.graph_compute_batch` may have several results.
+
+For each result `r` independently:
+
+- every lane produces one fragment of type `F_r`;
+- the graph result type is `N × F_r`;
+- its physical publication and logical reconstruction descriptor are verified
+  independently.
+
+Different results may use different fragment shapes.
+
+## Graph consumers
+
+A graph consumer of a batch result may:
+
+1. consume fragments directly as physical fragments;
+2. select one or more physical slots in a `spat.deferred_communication` body;
+3. use a Blueprint to obtain or describe a logical reconstruction;
+4. feed fragments to following graph computes or graph compute batches.
+
+A consumer must not treat the leading physical slot dimension as a logical
+model dimension unless an explicit graph operation intentionally performs such
+an interpretation.
+
+All constant selection, slicing, reshaping, concatenation, and other
+compile-time shaping needed for a scheduled consumer must be encoded inside the
+corresponding `spat.deferred_communication` body. Phase 2 must not recover
+missing graph semantics by inspecting consumers after the deferred operation.
+
+## Graph lane, scheduled lane, and physical core are different identities
+
+These concepts must never be conflated:
+
+- **graph lane**: the lane of the original `spat.graph_compute_batch`;
+- **physical fragment slot**: the slot in the graph batch result;
+- **scheduled lane**: one lane of a `spat.scheduled_compute_batch` equivalence
+  class;
+- **physical core**: the core selected by PEFT.
+
+The graph batch body or its Blueprint defines graph-lane-to-fragment-slot and
+fragment-slot-to-logical-region mappings.
+
+PEFT defines graph-instance-to-core placement.
+
+Scheduled communication defines how values move between cores.
+
+## Scheduled IR exclusion
+
+Do not add a verifier requiring `spat.scheduled_compute_batch` results to have
+`laneCount` as their first dimension.
+
+Do not rewrite scheduled values merely to resemble graph physical fragment
+collections.
+
+When lowering graph IR into scheduled IR:
+
+- resolve graph fragments and reconstruction metadata before erasing their
+  graph owners;
+- create local forwarding or `spat.channel_send`/`spat.channel_receive` for
+  cross-core dependencies;
+- allow scheduled result representation to follow the scheduled IR contract;
+- preserve numerical and deadlock correctness.
+
+The graph invariant is an input contract for scheduling, not a scheduled-value
+layout contract.
+
+## Required verifier properties
+
+`spat.graph_compute_batch` verification must establish, for every result:
+
+1. the result is a static or otherwise supported ranked tensor;
+2. result rank is exactly `fragment rank + 1`;
+3. result dimension 0 equals `laneCount`;
+4. every lane publication source has the same exact fragment type;
+5. the physical insertion targets the corresponding result block argument;
+6. physical insertion offsets have the fragment slot in dimension 0;
+7. all remaining physical offsets are zero;
+8. physical sizes are `[1] + fragment shape`;
+9. physical strides are unit;
+10. exactly one publication is defined for each graph result in the per-lane
+    body.
+
+These checks apply only to `spat.graph_compute_batch`, not to
+`spat.scheduled_compute_batch`.
+
+Blueprint verification must establish that every logical reconstruction entry:
+
+- references an existing physical batch operand;
+- references a valid physical fragment slot;
+- maps a fragment compatible with the declared logical slice;
+- stays within logical bounds;
+- follows the declared conflict and coverage policies.
+
+## Invalid representations
+
+The following are invariant violations.
+
+### Logical aggregate returned directly by graph batch
+
+```text
+laneCount = 16
+result     = tensor<1x4x16x16>
+```
+
+with each lane inserting into logical dimension 2.
+
+This is a logical assembly masquerading as a graph batch result. The graph
+result must instead be `16 × per-row-fragment`, and a Blueprint must describe
+placement into `tensor<1x4x16x16>`.
+
+### Physical storage derived from logical destination shape
+
+Code equivalent to:
+
+```cpp
+shape = logicalDestinationType.getShape();
+shape[logicalInsertionDimension] = laneCount;
+```
+
+is invalid.
+
+Physical graph storage must be derived from the per-lane fragment type:
+
+```cpp
+physicalShape = [laneCount] + fragmentType.getShape();
+```
+
+### Reconstruction inferred from result type
+
+It is invalid to assume that physical slot `i` belongs at logical offset `i`.
+The Blueprint or another explicit reconstruction descriptor must state the
+mapping.
+
+### Blueprint used for arithmetic
+
+It is invalid to encode `fragment0 + fragment1` as Blueprint reconstruction.
+Create a following graph compute or graph compute batch for the addition.
+
+## Ownership
+
+- ONNX-to-Spatial lowering owns creation of valid graph fragment batches.
+- Graph canonicalization owns normalization of temporary logical-assembly forms
+  into physical graph batches plus Blueprints.
+- `spat.graph_compute_batch` verifier rejects invalid physical publications.
+- `spat.blueprint` owns persistent logical reconstruction metadata.
+- Deferred communication Phase 1 owns complete consumer-side constant shaping.
+- Merge scheduling consumes this graph contract and introduces explicit
+  communication.
+- Scheduled IR verifiers validate scheduled execution and communication, not
+  the graph fragment representation.
+
+## No repair downstream
+
+If graph IR violates this invariant, fix the graph producer or graph
+canonicalization.
+
+Do not repair an invalid graph batch by:
+
+- guessing a lane dimension in `MergeComputeNodes`;
+- deriving physical storage from a logical destination tensor;
+- inspecting deferred-result users;
+- reconstructing omitted Blueprint data after graph erasure;
+- weakening graph verifiers;
+- imposing the graph representation on scheduled operations.

src/PIM/CMakeLists.txt

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

src/PIM/Common/CMakeLists.txt

@@ -8,6 +8,8 @@ add_pim_library(OMPimCommon
   IR/IndexingUtils.cpp
   IR/LoopUtils.cpp
   IR/ShapeUtils.cpp
+  IR/ShapingUtils.cpp
+  IR/StaticIntSequence.cpp
   IR/SubviewUtils.cpp
   IR/TensorSliceUtils.cpp
   IR/WeightUtils.cpp

src/PIM/Common/IR/AffineUtils.cpp

@@ -31,7 +31,7 @@ static FailureOr<int64_t> ceilDivSigned(int64_t lhs, int64_t rhs) {
 }
 
 Value createOrFoldAffineApply(
-  RewriterBase& rewriter, Location loc, AffineMap map, ValueRange operands, Operation* constantAnchor) {
+  OpBuilder& builder, Location loc, AffineMap map, ValueRange operands, Operation* constantAnchor) {
   assert(constantAnchor && "expected a valid constant anchor");
   assert(map.getNumResults() == 1 && "affine.apply expects a single-result affine map");
 
@@ -40,91 +40,91 @@ Value createOrFoldAffineApply(
   for (Value operand : operands) {
     std::optional<int64_t> constantValue = matchConstantIndexValue(operand);
     if (!constantValue)
-      return affine::AffineApplyOp::create(rewriter, loc, map, operands).getResult();
-    operandConstants.push_back(rewriter.getIndexAttr(*constantValue));
+      return affine::AffineApplyOp::create(builder, loc, map, operands).getResult();
+    operandConstants.push_back(builder.getIndexAttr(*constantValue));
   }
 
   SmallVector<Attribute> foldedResults;
   if (succeeded(map.constantFold(operandConstants, foldedResults)) && foldedResults.size() == 1)
     if (auto constantResult = dyn_cast<IntegerAttr>(foldedResults.front()))
-      return getOrCreateIndexConstant(rewriter, constantAnchor, constantResult.getInt());
+      return getOrCreateIndexConstant(builder, constantAnchor, constantResult.getInt());
 
-  return affine::AffineApplyOp::create(rewriter, loc, map, operands).getResult();
+  return affine::AffineApplyOp::create(builder, loc, map, operands).getResult();
 }
 
 Value createOrFoldAffineApply(
-  RewriterBase& rewriter, Location loc, AffineExpr expr, ValueRange dims, Operation* constantAnchor) {
+  OpBuilder& builder, Location loc, AffineExpr expr, ValueRange dims, Operation* constantAnchor) {
   AffineMap map = AffineMap::get(/*dimCount=*/dims.size(), /*symbolCount=*/0, expr);
-  return createOrFoldAffineApply(rewriter, loc, map, dims, constantAnchor);
+  return createOrFoldAffineApply(builder, loc, map, dims, constantAnchor);
 }
 
-Value affineMulConst(RewriterBase& rewriter, Location loc, Value value, int64_t multiplier, Operation* constantAnchor) {
+Value affineMulConst(OpBuilder& builder, Location loc, Value value, int64_t multiplier, Operation* constantAnchor) {
   assert(constantAnchor && "expected a valid constant anchor");
   if (multiplier == 0)
-    return getOrCreateIndexConstant(rewriter, constantAnchor, 0);
+    return getOrCreateIndexConstant(builder, constantAnchor, 0);
   if (multiplier == 1)
     return value;
 
-  AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext());
-  return createOrFoldAffineApply(rewriter, loc, d0 * multiplier, ValueRange {value}, constantAnchor);
+  AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
+  return createOrFoldAffineApply(builder, loc, d0 * multiplier, ValueRange {value}, constantAnchor);
 }
 
-Value affineAddConst(RewriterBase& rewriter, Location loc, Value value, int64_t offset, Operation* constantAnchor) {
+Value affineAddConst(OpBuilder& builder, 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);
+  AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
+  return createOrFoldAffineApply(builder, loc, d0 + offset, ValueRange {value}, constantAnchor);
 }
 
-Value affineModConst(RewriterBase& rewriter, Location loc, Value value, int64_t divisor, Operation* constantAnchor) {
+Value affineModConst(OpBuilder& builder, 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");
   if (divisor == 1)
-    return getOrCreateIndexConstant(rewriter, constantAnchor, 0);
+    return getOrCreateIndexConstant(builder, constantAnchor, 0);
 
-  AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext());
-  return createOrFoldAffineApply(rewriter, loc, d0 % divisor, ValueRange {value}, constantAnchor);
+  AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
+  return createOrFoldAffineApply(builder, loc, d0 % divisor, ValueRange {value}, constantAnchor);
 }
 
 Value affineFloorDivConst(
-  RewriterBase& rewriter, Location loc, Value value, int64_t divisor, Operation* constantAnchor) {
+  OpBuilder& builder, Location loc, Value value, 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 value;
 
-  AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext());
-  return createOrFoldAffineApply(rewriter, loc, d0.floorDiv(divisor), ValueRange {value}, constantAnchor);
+  AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
+  return createOrFoldAffineApply(builder, loc, d0.floorDiv(divisor), ValueRange {value}, constantAnchor);
 }
 
 Value affineAddModConst(
-  RewriterBase& rewriter, Location loc, Value value, int64_t offset, int64_t divisor, Operation* constantAnchor) {
+  OpBuilder& builder, 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);
+    return getOrCreateIndexConstant(builder, constantAnchor, 0);
 
-  AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext());
+  AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
   AffineExpr expr = d0;
   if (offset != 0)
     expr = expr + offset;
-  return createOrFoldAffineApply(rewriter, loc, expr % divisor, ValueRange {value}, constantAnchor);
+  return createOrFoldAffineApply(builder, loc, expr % divisor, ValueRange {value}, constantAnchor);
 }
 
 Value affineAddFloorDivConst(
-  RewriterBase& rewriter, Location loc, Value value, int64_t offset, int64_t divisor, Operation* constantAnchor) {
+  OpBuilder& builder, 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);
+    return offset == 0 ? value : affineAddConst(builder, loc, value, offset, constantAnchor);
 
-  AffineExpr d0 = getAffineDimExpr(0, rewriter.getContext());
+  AffineExpr d0 = getAffineDimExpr(0, builder.getContext());
   AffineExpr expr = d0;
   if (offset != 0)
     expr = expr + offset;
-  return createOrFoldAffineApply(rewriter, loc, expr.floorDiv(divisor), ValueRange {value}, constantAnchor);
+  return createOrFoldAffineApply(builder, loc, expr.floorDiv(divisor), ValueRange {value}, constantAnchor);
 }
 
 FailureOr<int64_t> evaluateAffineExpr(AffineExpr expr, ArrayRef<int64_t> dims, ArrayRef<int64_t> symbols) {

src/PIM/Common/IR/AffineUtils.hpp

@@ -11,50 +11,50 @@ namespace onnx_mlir {
 
 using IndexValueResolver = llvm::function_ref<llvm::FailureOr<int64_t>(mlir::Value)>;
 
-mlir::Value createOrFoldAffineApply(mlir::RewriterBase& rewriter,
+mlir::Value createOrFoldAffineApply(mlir::OpBuilder& builder,
                                     mlir::Location loc,
                                     mlir::AffineMap map,
                                     mlir::ValueRange operands,
                                     mlir::Operation* constantAnchor);
 
-mlir::Value createOrFoldAffineApply(mlir::RewriterBase& rewriter,
+mlir::Value createOrFoldAffineApply(mlir::OpBuilder& builder,
                                     mlir::Location loc,
                                     mlir::AffineExpr expr,
                                     mlir::ValueRange dims,
                                     mlir::Operation* constantAnchor);
 
-mlir::Value affineMulConst(mlir::RewriterBase& rewriter,
+mlir::Value affineMulConst(mlir::OpBuilder& builder,
                            mlir::Location loc,
                            mlir::Value value,
                            int64_t multiplier,
                            mlir::Operation* constantAnchor);
 
-mlir::Value affineAddConst(mlir::RewriterBase& rewriter,
+mlir::Value affineAddConst(mlir::OpBuilder& builder,
                            mlir::Location loc,
                            mlir::Value value,
                            int64_t offset,
                            mlir::Operation* constantAnchor);
 
-mlir::Value affineModConst(mlir::RewriterBase& rewriter,
+mlir::Value affineModConst(mlir::OpBuilder& builder,
                            mlir::Location loc,
                            mlir::Value value,
                            int64_t divisor,
                            mlir::Operation* constantAnchor);
 
-mlir::Value affineFloorDivConst(mlir::RewriterBase& rewriter,
+mlir::Value affineFloorDivConst(mlir::OpBuilder& builder,
                                 mlir::Location loc,
                                 mlir::Value value,
                                 int64_t divisor,
                                 mlir::Operation* constantAnchor);
 
-mlir::Value affineAddModConst(mlir::RewriterBase& rewriter,
+mlir::Value affineAddModConst(mlir::OpBuilder& builder,
                               mlir::Location loc,
                               mlir::Value value,
                               int64_t offset,
                               int64_t divisor,
                               mlir::Operation* constantAnchor);
 
-mlir::Value affineAddFloorDivConst(mlir::RewriterBase& rewriter,
+mlir::Value affineAddFloorDivConst(mlir::OpBuilder& builder,
                                    mlir::Location loc,
                                    mlir::Value value,
                                    int64_t offset,

src/PIM/Common/IR/ConstantUtils.cpp

@@ -10,6 +10,32 @@ using namespace mlir;
 
 namespace onnx_mlir {
 
+ConstantPool::ConstantPool(Operation *constantAnchor, OpBuilder &builder)
+  : anchor(constantAnchor), block(getConstantInsertionBlock(constantAnchor)),
+    builder(builder) {
+  for (Operation &op : *block)
+    if (auto constant = dyn_cast<arith::ConstantOp>(&op))
+      cache.try_emplace(
+        std::make_pair(constant.getType(), constant.getValue()),
+        constant.getResult());
+}
+
+Value ConstantPool::getIndex(int64_t value) {
+  return get(builder.getIndexType(), builder.getIndexAttr(value));
+}
+
+Value ConstantPool::get(Type type, Attribute value) {
+  auto key = std::make_pair(type, value);
+  if (Value existing = cache.lookup(key))
+    return existing;
+  OpBuilder::InsertionGuard guard(builder);
+  builder.setInsertionPointToStart(block);
+  Value constant = arith::ConstantOp::create(
+    builder, anchor->getLoc(), type, cast<TypedAttr>(value)).getResult();
+  cache.try_emplace(key, constant);
+  return constant;
+}
+
 static std::optional<int64_t> getIndexConstantValue(arith::ConstantOp constantOp) {
   if (!constantOp.getType().isIndex())
     return std::nullopt;
@@ -49,7 +75,7 @@ Value getOrCreateConstant(OperationFolder& folder, Operation* anchorOp, Attribut
   return folder.getOrCreateConstant(hostBlock, arithDialect, value, type);
 }
 
-Value getOrCreateConstant(RewriterBase& rewriter, Operation* anchorOp, Attribute value, Type type) {
+Value getOrCreateConstant(OpBuilder& builder, Operation* anchorOp, Attribute value, Type type) {
   assert(anchorOp && "expected a valid anchor operation");
   Block* hostBlock = getConstantInsertionBlock(anchorOp);
   for (Operation& op : *hostBlock) {
@@ -59,9 +85,16 @@ Value getOrCreateConstant(RewriterBase& rewriter, Operation* anchorOp, Attribute
     return constantOp.getResult();
   }
 
-  OpBuilder::InsertionGuard guard(rewriter);
-  rewriter.setInsertionPointToStart(hostBlock);
-  return arith::ConstantOp::create(rewriter, anchorOp->getLoc(), type, cast<TypedAttr>(value)).getResult();
+  OpBuilder::InsertionGuard guard(builder);
+  builder.setInsertionPointToStart(hostBlock);
+  return arith::ConstantOp::create(builder, anchorOp->getLoc(), type, cast<TypedAttr>(value)).getResult();
+}
+
+Value createConstantAtHostBlockStart(OpBuilder& builder, Operation* anchorOp, TypedAttr value) {
+  assert(anchorOp && "expected a valid anchor operation");
+  OpBuilder::InsertionGuard guard(builder);
+  builder.setInsertionPointToStart(getConstantInsertionBlock(anchorOp));
+  return arith::ConstantOp::create(builder, anchorOp->getLoc(), value).getResult();
 }
 
 Value getOrCreateConstantLike(OperationFolder& folder, arith::ConstantOp constantOp) {
@@ -73,9 +106,8 @@ Value getOrCreateIndexConstant(OperationFolder& folder, Operation* anchorOp, int
   return getOrCreateConstant(folder, anchorOp, builder.getIndexAttr(value), builder.getIndexType());
 }
 
-Value getOrCreateIndexConstant(RewriterBase& rewriter, Operation* anchorOp, int64_t value) {
-  Builder builder(anchorOp->getContext());
-  return getOrCreateConstant(rewriter, anchorOp, builder.getIndexAttr(value), builder.getIndexType());
+Value getOrCreateIndexConstant(OpBuilder& builder, Operation* anchorOp, int64_t value) {
+  return getOrCreateConstant(builder, anchorOp, builder.getIndexAttr(value), builder.getIndexType());
 }
 
 void hoistAndUniquifyIndexConstants(func::FuncOp funcOp, RewriterBase& rewriter) {

src/PIM/Common/IR/ConstantUtils.hpp

@@ -6,23 +6,42 @@
 #include "mlir/IR/Value.h"
 #include "mlir/Transforms/FoldUtils.h"
 
+#include "llvm/ADT/DenseMap.h"
+
 #include <optional>
 
 namespace onnx_mlir {
 
+class ConstantPool {
+public:
+  ConstantPool(mlir::Operation *constantAnchor, mlir::OpBuilder &builder);
+
+  mlir::Value getIndex(int64_t value);
+  mlir::Value get(mlir::Type type, mlir::Attribute value);
+
+private:
+  mlir::Operation *anchor;
+  mlir::Block *block;
+  mlir::OpBuilder &builder;
+  llvm::DenseMap<std::pair<mlir::Type, mlir::Attribute>, mlir::Value> cache;
+};
+
 mlir::Block* getConstantInsertionBlock(mlir::Operation* anchorOp);
 
 mlir::Value
 getOrCreateConstant(mlir::OperationFolder& folder, mlir::Operation* anchorOp, mlir::Attribute value, mlir::Type type);
 
 mlir::Value
-getOrCreateConstant(mlir::RewriterBase& rewriter, mlir::Operation* anchorOp, mlir::Attribute value, mlir::Type type);
+getOrCreateConstant(mlir::OpBuilder& builder, mlir::Operation* anchorOp, mlir::Attribute value, mlir::Type type);
+
+mlir::Value
+createConstantAtHostBlockStart(mlir::OpBuilder& builder, mlir::Operation* anchorOp, mlir::TypedAttr value);
 
 mlir::Value getOrCreateConstantLike(mlir::OperationFolder& folder, mlir::arith::ConstantOp constantOp);
 
 mlir::Value getOrCreateIndexConstant(mlir::OperationFolder& folder, mlir::Operation* anchorOp, int64_t value);
 
-mlir::Value getOrCreateIndexConstant(mlir::RewriterBase& rewriter, mlir::Operation* anchorOp, int64_t value);
+mlir::Value getOrCreateIndexConstant(mlir::OpBuilder& builder, mlir::Operation* anchorOp, int64_t value);
 
 void hoistAndUniquifyIndexConstants(mlir::func::FuncOp funcOp, mlir::RewriterBase& rewriter);
 

src/PIM/Common/IR/CoreBlockUtils.cpp

@@ -36,9 +36,10 @@ bool isCoreStaticAddressOp(mlir::Operation* op) {
 
 mlir::LogicalResult
 walkPimCoreBlock(mlir::Block& block,
-                 const StaticValueKnowledge& knowledge,
+                 const StaticValueKnowledge& initialKnowledge,
                  llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback) {
   bool hasFailure = false;
+  StaticValueKnowledge knowledge = initialKnowledge;
   for (mlir::Operation& op : block) {
     if (mlir::isa<pim::PimHaltOp, mlir::scf::YieldOp>(op) || isCoreStaticAddressOp(&op))
       continue;
@@ -74,6 +75,42 @@ 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 (auto switchOp = mlir::dyn_cast<mlir::scf::IndexSwitchOp>(op)) {
+      auto selector = resolveIndexValue(switchOp.getArg(), knowledge);
+      if (failed(selector)) {
+        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen");
+        hasFailure = true;
+        continue;
+      }
+      mlir::Region* selected = &switchOp.getDefaultRegion();
+      for (auto [caseValue, caseRegion] : llvm::zip(switchOp.getCases(), switchOp.getCaseRegions()))
+        if (caseValue == *selector) {
+          selected = &caseRegion;
+          break;
+        }
+      if (failed(walkPimCoreBlock(selected->front(), knowledge, callback)))
+        hasFailure = true;
+      auto yield = mlir::cast<mlir::scf::YieldOp>(selected->front().getTerminator());
+      for (auto [result, yielded] : llvm::zip(switchOp.getResults(), yield.getOperands()))
+        knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge);
+      continue;
+    }
+
     if (failed(callback(op, knowledge)))
       hasFailure = true;
   }
@@ -82,9 +119,10 @@ walkPimCoreBlock(mlir::Block& block,
 
 mlir::LogicalResult walkPimCoreBlockStructurally(
   mlir::Block& block,
-  const StaticValueKnowledge& knowledge,
+  const StaticValueKnowledge& initialKnowledge,
   llvm::function_ref<mlir::LogicalResult(mlir::Operation&, const StaticValueKnowledge&)> callback) {
   bool hasFailure = false;
+  StaticValueKnowledge knowledge = initialKnowledge;
   for (mlir::Operation& op : block) {
     if (mlir::isa<pim::PimHaltOp, mlir::scf::YieldOp>(op) || isCoreStaticAddressOp(&op))
       continue;
@@ -128,6 +166,44 @@ 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 (auto switchOp = mlir::dyn_cast<mlir::scf::IndexSwitchOp>(op)) {
+      auto selector = resolveIndexValue(switchOp.getArg(), knowledge);
+      if (failed(selector)) {
+        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM verification");
+        hasFailure = true;
+        continue;
+      }
+      mlir::Region* selected = &switchOp.getDefaultRegion();
+      for (auto [caseValue, caseRegion] : llvm::zip(switchOp.getCases(), switchOp.getCaseRegions()))
+        if (caseValue == *selector) {
+          selected = &caseRegion;
+          break;
+        }
+      for (mlir::Region& region : switchOp->getRegions())
+        if (failed(walkPimCoreBlockStructurally(region.front(), knowledge, callback)))
+          hasFailure = true;
+      auto yield = mlir::cast<mlir::scf::YieldOp>(selected->front().getTerminator());
+      for (auto [result, yielded] : llvm::zip(switchOp.getResults(), yield.getOperands()))
+        knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge);
+      continue;
+    }
+
     if (failed(callback(op, knowledge)))
       hasFailure = true;
   }

src/PIM/Common/IR/ShapingUtils.cpp

@@ -0,0 +1,39 @@
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Interfaces/SideEffectInterfaces.h"
+
+#include "ShapingUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+bool isShapingOnlyOp(Operation *op) {
+  return isa<tensor::CastOp,
+             tensor::CollapseShapeOp,
+             tensor::ExpandShapeOp,
+             tensor::ExtractSliceOp,
+             tensor::InsertSliceOp,
+             tensor::ConcatOp,
+             tensor::EmptyOp,
+             tensor::ExtractOp,
+             tensor::InsertOp,
+             tensor::SplatOp,
+             linalg::TransposeOp,
+             ONNXTransposeOp,
+             spatial::SpatConcatOp,
+             spatial::SpatExtractRowsOp>(op);
+}
+
+bool isPureIndexComputationOp(Operation *op) {
+  if (op->getNumRegions() != 0 || op->getNumResults() == 0 || op->hasTrait<OpTrait::IsTerminator>()
+      || !isMemoryEffectFree(op))
+    return false;
+  auto isIndexOrInteger = [](Type type) { return type.isIndex() || isa<IntegerType>(type); };
+  return llvm::all_of(op->getOperandTypes(), isIndexOrInteger)
+      && llvm::all_of(op->getResultTypes(), isIndexOrInteger);
+}
+
+} // namespace onnx_mlir

src/PIM/Common/IR/ShapingUtils.hpp

@@ -0,0 +1,13 @@
+#pragma once
+
+namespace mlir {
+class Operation;
+}
+
+namespace onnx_mlir {
+
+bool isShapingOnlyOp(mlir::Operation *op);
+
+bool isPureIndexComputationOp(mlir::Operation *op);
+
+} // namespace onnx_mlir

src/PIM/Common/IR/StaticIntSequence.cpp

@@ -0,0 +1,539 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/MathExtras.h"
+
+#include <limits>
+
+#include "AffineUtils.hpp"
+#include "ConstantUtils.hpp"
+#include "StaticIntSequence.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+static bool getAffineValue(int64_t base, int64_t step, size_t index,
+                           int64_t &value) {
+  if (index > static_cast<size_t>(std::numeric_limits<int64_t>::max()))
+    return false;
+  int64_t scaled;
+  return !llvm::MulOverflow(static_cast<int64_t>(index), step, scaled)
+      && !llvm::AddOverflow(base, scaled, value);
+}
+
+static FailureOr<SmallVector<int64_t>> getI64Values(Operation *op,
+                                                     StringRef name) {
+  Attribute attr = op->getAttr(name);
+  if (!attr)
+    return op->emitOpError() << "is missing " << name << " metadata",
+           failure();
+  if (auto scalar = dyn_cast<IntegerAttr>(attr))
+    return SmallVector<int64_t> {scalar.getInt()};
+  if (auto array = dyn_cast<DenseI64ArrayAttr>(attr))
+    return SmallVector<int64_t>(array.asArrayRef());
+  auto elements = dyn_cast<DenseIntElementsAttr>(attr);
+  auto type = elements ? dyn_cast<RankedTensorType>(elements.getType())
+                       : RankedTensorType();
+  if (!elements || !type || type.getRank() != 1
+      || !type.getElementType().isInteger(64))
+    return op->emitOpError() << "has invalid " << name << " metadata",
+           failure();
+  SmallVector<int64_t> values;
+  values.reserve(elements.getNumElements());
+  for (APInt value : elements.getValues<APInt>())
+    values.push_back(value.getSExtValue());
+  return values;
+}
+
+} // namespace
+
+StaticIntSequence StaticIntSequence::uniform(int64_t value, size_t count) {
+  assert(count != 0 && "empty static integer sequence");
+  StaticIntSequence result;
+  result.kind = StaticIntSequenceKind::Uniform;
+  result.count = count;
+  result.base = value;
+  return result;
+}
+
+StaticIntSequence StaticIntSequence::affine(int64_t base, int64_t step,
+                                             size_t count) {
+  assert(count != 0 && "empty static integer sequence");
+  int64_t last;
+  assert(getAffineValue(base, step, count - 1, last)
+         && "overflowing static affine sequence");
+  if (count == 1 || step == 0)
+    return uniform(base, count);
+  StaticIntSequence result;
+  result.kind = StaticIntSequenceKind::Affine;
+  result.count = count;
+  result.base = base;
+  result.step = step;
+  return result;
+}
+
+StaticIntSequence StaticIntSequence::runLengthEncoded(
+  ArrayRef<int64_t> runs, size_t count) {
+  assert(count != 0 && runs.size() % 2 == 0
+         && "invalid run-length encoded sequence");
+  StaticIntSequence result;
+  result.kind = StaticIntSequenceKind::RunLengthEncoded;
+  result.count = count;
+  result.data.assign(runs);
+  return result;
+}
+
+StaticIntSequence StaticIntSequence::fromValues(ArrayRef<int64_t> values) {
+  assert(!values.empty() && "empty static integer sequence");
+  if (llvm::all_equal(values))
+    return uniform(values.front(), values.size());
+  int64_t step;
+  bool isAffine = !llvm::SubOverflow(values[1], values[0], step);
+  for (size_t index = 1; isAffine && index < values.size(); ++index) {
+    int64_t difference;
+    isAffine = !llvm::SubOverflow(values[index], values[index - 1],
+                                  difference)
+        && difference == step;
+  }
+  if (isAffine)
+    return affine(values.front(), step, values.size());
+
+  SmallVector<int64_t> runs;
+  for (int64_t value : values) {
+    if (!runs.empty() && runs[runs.size() - 2] == value) {
+      ++runs.back();
+      continue;
+    }
+    runs.push_back(value);
+    runs.push_back(1);
+  }
+  if (runs.size() < values.size())
+    return runLengthEncoded(runs, values.size());
+  StaticIntSequence result;
+  result.kind = StaticIntSequenceKind::Dense;
+  result.count = values.size();
+  result.data.assign(values);
+  return result;
+}
+
+int64_t StaticIntSequence::valueAt(size_t index) const {
+  assert(index < count && "static integer sequence index out of range");
+  if (kind == StaticIntSequenceKind::Uniform)
+    return base;
+  if (kind == StaticIntSequenceKind::Affine) {
+    int64_t value;
+    bool valid = getAffineValue(base, step, index, value);
+    assert(valid && "overflowing static affine sequence");
+    return value;
+  }
+  if (kind == StaticIntSequenceKind::Dense)
+    return data[index];
+  for (size_t run = 0; run < data.size(); run += 2) {
+    size_t length = static_cast<size_t>(data[run + 1]);
+    if (index < length)
+      return data[run];
+    index -= length;
+  }
+  llvm_unreachable("malformed run-length encoded sequence");
+}
+
+std::optional<size_t> StaticIntSequence::find(int64_t value, size_t begin,
+                                              size_t length) const {
+  assert(begin <= count && length <= count - begin
+         && "invalid static integer sequence search");
+  if (length == 0)
+    return std::nullopt;
+  size_t end = begin + length;
+  if (kind == StaticIntSequenceKind::Uniform)
+    return value == base ? std::optional<size_t>(begin) : std::nullopt;
+  if (kind == StaticIntSequenceKind::Affine) {
+    int64_t delta;
+    if (llvm::SubOverflow(value, base, delta) || delta % step != 0)
+      return std::nullopt;
+    int64_t index = delta / step;
+    return index >= static_cast<int64_t>(begin)
+            && index < static_cast<int64_t>(end)
+        ? std::optional<size_t>(index)
+        : std::nullopt;
+  }
+  if (kind == StaticIntSequenceKind::Dense) {
+    ArrayRef<int64_t> selected = ArrayRef(data).slice(begin, length);
+    auto found = llvm::find(selected, value);
+    return found == selected.end()
+        ? std::nullopt
+        : std::optional<size_t>(begin + (found - selected.begin()));
+  }
+  size_t runBegin = 0;
+  for (size_t run = 0; run < data.size(); run += 2) {
+    size_t runEnd = runBegin + static_cast<size_t>(data[run + 1]);
+    if (runEnd > begin && runBegin < end && data[run] == value)
+      return std::max(begin, runBegin);
+    if (runBegin >= end)
+      break;
+    runBegin = runEnd;
+  }
+  return std::nullopt;
+}
+
+StaticIntSequence StaticIntSequence::slice(size_t begin, size_t length) const {
+  assert(length != 0 && begin <= count - length && "invalid sequence slice");
+  if (kind == StaticIntSequenceKind::Uniform)
+    return uniform(base, length);
+  if (kind == StaticIntSequenceKind::Affine)
+    return affine(valueAt(begin), step, length);
+  if (kind == StaticIntSequenceKind::Dense)
+    return fromValues(ArrayRef(data).slice(begin, length));
+  SmallVector<int64_t> runs;
+  size_t end = begin + length;
+  forEachEqualRun([&](int64_t value, size_t runBegin, size_t runCount) {
+    size_t selectedBegin = std::max(begin, runBegin);
+    size_t selectedEnd = std::min(end, runBegin + runCount);
+    if (selectedBegin >= selectedEnd)
+      return;
+    if (!runs.empty() && runs[runs.size() - 2] == value)
+      runs.back() += selectedEnd - selectedBegin;
+    else {
+      runs.push_back(value);
+      runs.push_back(selectedEnd - selectedBegin);
+    }
+  });
+  if (runs.size() == 2)
+    return uniform(runs.front(), length);
+  if (runs.size() < length)
+    return runLengthEncoded(runs, length);
+  SmallVector<int64_t> values;
+  for (size_t run = 0; run < runs.size(); run += 2)
+    values.append(runs[run + 1], runs[run]);
+  return fromValues(values);
+}
+
+StaticIntSequence StaticIntSequence::remap(ArrayRef<unsigned> indices) const {
+  assert(!indices.empty() && "empty static integer sequence remap");
+  SmallVector<int64_t> values;
+  values.reserve(indices.size());
+  for (unsigned index : indices)
+    values.push_back(valueAt(index));
+  return fromValues(values);
+}
+
+bool StaticIntSequence::operator==(const StaticIntSequence& other) const {
+  return kind == other.kind && count == other.count && base == other.base
+      && step == other.step && data == other.data;
+}
+
+llvm::hash_code StaticIntSequence::hash() const {
+  return llvm::hash_combine(kind, count, base, step,
+                            llvm::hash_combine_range(data.begin(), data.end()));
+}
+
+void StaticIntSequence::forEachEqualRun(
+  llvm::function_ref<void(int64_t, size_t, size_t)> callback) const {
+  if (kind == StaticIntSequenceKind::Uniform) {
+    callback(base, 0, count);
+    return;
+  }
+  if (kind == StaticIntSequenceKind::RunLengthEncoded) {
+    size_t begin = 0;
+    for (size_t run = 0; run < data.size(); run += 2) {
+      size_t runCount = static_cast<size_t>(data[run + 1]);
+      callback(data[run], begin, runCount);
+      begin += runCount;
+    }
+    return;
+  }
+  size_t begin = 0;
+  while (begin < count) {
+    int64_t value = valueAt(begin);
+    size_t end = begin + 1;
+    while (end < count && valueAt(end) == value)
+      ++end;
+    callback(value, begin, end - begin);
+    begin = end;
+  }
+}
+
+void StaticIntSequenceChain::append(const StaticIntSequence &sequence,
+                                    size_t begin, size_t length) {
+  assert(length != 0 && begin <= sequence.size() - length
+         && "invalid static integer sequence chain slice");
+  if (!slices.empty()) {
+    StaticIntSequenceSlice &last = slices.back();
+    if (last.sequence == &sequence && last.begin + last.count == begin) {
+      last.count += length;
+      count += length;
+      return;
+    }
+    auto affinePart = [](const StaticIntSequenceSlice &slice,
+                         int64_t &base, int64_t &step) {
+      base = slice.sequence->valueAt(slice.begin);
+      if (slice.count == 1) {
+        step = 0;
+        return true;
+      }
+      return !llvm::SubOverflow(slice.sequence->valueAt(slice.begin + 1),
+                                base, step)
+          && (slice.sequence->kind == StaticIntSequenceKind::Uniform
+              || slice.sequence->kind == StaticIntSequenceKind::Affine);
+    };
+    StaticIntSequenceSlice next {&sequence, begin, length};
+    int64_t leftBase, leftStep, rightBase, rightStep, expected;
+    if (affinePart(last, leftBase, leftStep)
+        && affinePart(next, rightBase, rightStep)
+        && (last.count == 1 || length == 1 || leftStep == rightStep)) {
+      int64_t step = last.count == 1 ? rightStep : leftStep;
+      if (getAffineValue(leftBase, step, last.count, expected)
+          && expected == rightBase) {
+        owned.push_back(std::make_unique<StaticIntSequence>(
+          StaticIntSequence::affine(leftBase, step, last.count + length)));
+        last = {owned.back().get(), 0, last.count + length};
+        count += length;
+        return;
+      }
+    }
+  }
+  slices.push_back({&sequence, begin, length});
+  count += length;
+}
+
+void StaticIntSequenceChain::append(StaticIntSequence sequence) {
+  size_t length = sequence.size();
+  owned.push_back(std::make_unique<StaticIntSequence>(std::move(sequence)));
+  append(*owned.back(), 0, length);
+}
+
+int64_t StaticIntSequenceChain::valueAt(size_t index) const {
+  assert(index < count && "static integer sequence chain index out of range");
+  for (const StaticIntSequenceSlice &slice : slices) {
+    if (index < slice.count)
+      return slice.sequence->valueAt(slice.begin + index);
+    index -= slice.count;
+  }
+  llvm_unreachable("malformed static integer sequence chain");
+}
+
+void StaticIntSequenceChain::forEachSegment(llvm::function_ref<void(
+  const StaticIntSequence &, size_t, size_t)> callback) const {
+  for (const StaticIntSequenceSlice &slice : slices)
+    callback(*slice.sequence, slice.begin, slice.count);
+}
+
+void StaticIntSequenceChain::forEachEqualRun(
+  llvm::function_ref<void(int64_t, size_t, size_t)> callback) const {
+  std::optional<int64_t> pendingValue;
+  size_t pendingBegin = 0, pendingCount = 0, chainBegin = 0;
+  auto flush = [&] {
+    if (pendingValue)
+      callback(*pendingValue, pendingBegin, pendingCount);
+  };
+  for (const StaticIntSequenceSlice &slice : slices) {
+    size_t sliceEnd = slice.begin + slice.count;
+    slice.sequence->forEachEqualRun(
+      [&](int64_t value, size_t runBegin, size_t runCount) {
+        size_t begin = std::max(slice.begin, runBegin);
+        size_t end = std::min(sliceEnd, runBegin + runCount);
+        if (begin >= end)
+          return;
+        size_t selectedCount = end - begin;
+        size_t globalBegin = chainBegin + begin - slice.begin;
+        if (pendingValue && *pendingValue == value
+            && pendingBegin + pendingCount == globalBegin) {
+          pendingCount += selectedCount;
+          return;
+        }
+        flush();
+        pendingValue = value;
+        pendingBegin = globalBegin;
+        pendingCount = selectedCount;
+      });
+    chainBegin += slice.count;
+  }
+  flush();
+}
+
+StaticIntSequence StaticIntSequenceChain::canonicalize() const {
+  assert(count != 0 && "empty static integer sequence chain");
+  int64_t first = valueAt(0);
+  bool uniform = true;
+  forEachEqualRun([&](int64_t value, size_t, size_t) {
+    uniform &= value == first;
+  });
+  if (uniform)
+    return StaticIntSequence::uniform(first, count);
+
+  int64_t step = 0, previous = first;
+  bool affine = true, haveStep = false;
+  size_t position = 0;
+  forEachSegment([&](const StaticIntSequence &sequence, size_t begin,
+                     size_t length) {
+    if (!affine)
+      return;
+    for (size_t index = 0; index < length; ++index) {
+      int64_t value = sequence.valueAt(begin + index);
+      if (position++ == 0) {
+        previous = value;
+        continue;
+      }
+      if (!haveStep) {
+        affine = !llvm::SubOverflow(value, previous, step);
+        haveStep = true;
+      } else if (haveStep) {
+        int64_t difference;
+        affine = !llvm::SubOverflow(value, previous, difference)
+            && difference == step;
+      }
+      previous = value;
+      if (!affine)
+        break;
+    }
+  });
+  if (affine && haveStep)
+    return StaticIntSequence::affine(first, step, count);
+
+  SmallVector<int64_t> runs;
+  forEachEqualRun([&](int64_t value, size_t, size_t runCount) {
+    runs.push_back(value);
+    runs.push_back(runCount);
+  });
+  if (runs.size() < count)
+    return StaticIntSequence::runLengthEncoded(runs, count);
+  SmallVector<int64_t> values;
+  values.reserve(count);
+  for (size_t run = 0; run < runs.size(); run += 2)
+    values.append(runs[run + 1], runs[run]);
+  return StaticIntSequence::fromValues(values);
+}
+
+int64_t StaticIntSequenceChainCursor::value() const {
+  assert(!done() && "static integer sequence chain cursor is done");
+  const StaticIntSequenceSlice &current = chain.slices[slice];
+  return current.sequence->valueAt(current.begin + offset);
+}
+
+void StaticIntSequenceChainCursor::advance() {
+  assert(!done() && "static integer sequence chain cursor is done");
+  if (++offset != chain.slices[slice].count)
+    return;
+  offset = 0;
+  ++slice;
+}
+
+void setStaticIntSequenceAttr(Operation *op, StringRef name,
+                              const StaticIntSequence &sequence,
+                              size_t logicalCount) {
+  assert(sequence.size() == logicalCount && logicalCount != 0
+         && "invalid static integer metadata count");
+  SmallVector<int64_t> values;
+  StringRef encoding;
+  switch (sequence.kind) {
+  case StaticIntSequenceKind::Uniform:
+    encoding = "uniform";
+    values.push_back(sequence.base);
+    break;
+  case StaticIntSequenceKind::Affine:
+    encoding = "affine";
+    values = {sequence.base, sequence.step};
+    break;
+  case StaticIntSequenceKind::RunLengthEncoded:
+    encoding = "rle";
+    values = sequence.data;
+    break;
+  case StaticIntSequenceKind::Dense:
+    encoding = "dense";
+    values = sequence.data;
+    break;
+  }
+  OpBuilder builder(op);
+  auto type = RankedTensorType::get(
+    {static_cast<int64_t>(values.size())}, builder.getI64Type());
+  op->setAttr(name, DenseIntElementsAttr::get(type, values));
+  if (sequence.kind != StaticIntSequenceKind::Dense)
+    op->setAttr((name + "_encoding").str(), builder.getStringAttr(encoding));
+}
+
+FailureOr<StaticIntSequence> getStaticIntSequenceAttr(
+  Operation *op, StringRef name, size_t logicalCount) {
+  if (logicalCount == 0)
+    return op->emitOpError() << "has zero logical count for " << name,
+           failure();
+  auto values = getI64Values(op, name);
+  if (failed(values))
+    return failure();
+  auto encoding = op->getAttrOfType<StringAttr>((name + "_encoding").str());
+  if (!encoding) {
+    if (values->size() != logicalCount)
+      return op->emitOpError() << "has invalid dense " << name << " count",
+             failure();
+    return StaticIntSequence::fromValues(*values);
+  }
+  if (encoding.getValue() == "uniform") {
+    if (values->size() != 1)
+      return op->emitOpError() << "has invalid uniform " << name,
+             failure();
+    return StaticIntSequence::uniform(values->front(), logicalCount);
+  }
+  if (encoding.getValue() == "affine") {
+    int64_t last;
+    if (values->size() != 2
+        || !getAffineValue((*values)[0], (*values)[1], logicalCount - 1, last))
+      return op->emitOpError() << "has invalid affine " << name,
+             failure();
+    return StaticIntSequence::affine((*values)[0], (*values)[1], logicalCount);
+  }
+  if (encoding.getValue() == "rle") {
+    size_t count = 0;
+    if (values->empty() || values->size() % 2 != 0)
+      return op->emitOpError() << "has invalid RLE " << name, failure();
+    for (size_t index = 1; index < values->size(); index += 2) {
+      if ((*values)[index] <= 0
+          || static_cast<uint64_t>((*values)[index]) > logicalCount - count)
+        return op->emitOpError() << "has invalid RLE " << name, failure();
+      count += (*values)[index];
+    }
+    if (count != logicalCount)
+      return op->emitOpError() << "has mismatched RLE " << name << " count",
+             failure();
+    return StaticIntSequence::runLengthEncoded(*values, count);
+  }
+  if (encoding.getValue() == "dense") {
+    if (values->size() != logicalCount)
+      return op->emitOpError() << "has invalid dense " << name << " count",
+             failure();
+    return StaticIntSequence::fromValues(*values);
+  }
+  return op->emitOpError() << "has unknown " << name << " encoding",
+         failure();
+}
+
+Value emitStaticIntLookup(const StaticIntSequence& sequence, Value position,
+                          Operation* constantAnchor,
+                          ConstantPool& constants, OpBuilder& builder,
+                          Location loc) {
+  if (sequence.getKind() == StaticIntSequenceKind::Uniform)
+    return constants.getIndex(sequence.valueAt(0));
+  if (sequence.getKind() == StaticIntSequenceKind::Affine) {
+    Value scaled = affineMulConst(builder, loc, position,
+                                  sequence.valueAt(1) - sequence.valueAt(0),
+                                  constantAnchor);
+    return affineAddConst(builder, loc, scaled, sequence.valueAt(0),
+                          constantAnchor);
+  }
+  SmallVector<int64_t> values;
+  values.reserve(sequence.size());
+  sequence.forEachEqualRun([&](int64_t value, size_t, size_t count) {
+    values.append(count, value);
+  });
+  auto type = RankedTensorType::get(
+    {static_cast<int64_t>(values.size())}, builder.getI64Type());
+  Value table = constants.get(type,
+    DenseElementsAttr::get(type, ArrayRef<int64_t>(values)));
+  Value selected = tensor::ExtractOp::create(
+    builder, loc, table, ValueRange {position});
+  return arith::IndexCastOp::create(
+    builder, loc, builder.getIndexType(), selected);
+}
+
+} // namespace onnx_mlir

src/PIM/Common/IR/StaticIntSequence.hpp

@@ -0,0 +1,117 @@
+#pragma once
+
+#include "mlir/IR/Builders.h"
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/FunctionExtras.h"
+#include "llvm/ADT/Hashing.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <memory>
+#include <optional>
+
+namespace onnx_mlir {
+
+class ConstantPool;
+
+enum class StaticIntSequenceKind {
+  Uniform,
+  Affine,
+  RunLengthEncoded,
+  Dense
+};
+
+class StaticIntSequence {
+public:
+  static StaticIntSequence fromValues(llvm::ArrayRef<int64_t> values);
+  static StaticIntSequence uniform(int64_t value, size_t count);
+  static StaticIntSequence affine(int64_t base, int64_t step, size_t count);
+
+  size_t size() const { return count; }
+  int64_t valueAt(size_t index) const;
+  std::optional<size_t> find(int64_t value, size_t begin, size_t length) const;
+  StaticIntSequence slice(size_t begin, size_t count) const;
+  StaticIntSequence remap(llvm::ArrayRef<unsigned> indices) const;
+  StaticIntSequenceKind getKind() const { return kind; }
+
+  bool operator==(const StaticIntSequence& other) const;
+  llvm::hash_code hash() const;
+
+  void forEachEqualRun(
+    llvm::function_ref<void(int64_t, size_t, size_t)> callback) const;
+
+private:
+  friend class StaticIntSequenceChain;
+  friend void setStaticIntSequenceAttr(mlir::Operation *, llvm::StringRef,
+                                       const StaticIntSequence &, size_t);
+  friend mlir::FailureOr<StaticIntSequence>
+  getStaticIntSequenceAttr(mlir::Operation *, llvm::StringRef, size_t);
+
+  static StaticIntSequence runLengthEncoded(
+    llvm::ArrayRef<int64_t> runs, size_t count);
+  StaticIntSequenceKind kind = StaticIntSequenceKind::Dense;
+  size_t count = 0;
+  int64_t base = 0;
+  int64_t step = 0;
+  llvm::SmallVector<int64_t> data;
+};
+
+struct StaticIntSequenceSlice {
+  const StaticIntSequence *sequence = nullptr;
+  size_t begin = 0;
+  size_t count = 0;
+};
+
+class StaticIntSequenceChain {
+public:
+  void append(const StaticIntSequence &sequence, size_t begin, size_t count);
+  void append(StaticIntSequence sequence);
+  size_t size() const { return count; }
+  int64_t valueAt(size_t index) const;
+  void forEachSegment(llvm::function_ref<void(
+    const StaticIntSequence &, size_t, size_t)> callback) const;
+  void forEachEqualRun(
+    llvm::function_ref<void(int64_t, size_t, size_t)> callback) const;
+  StaticIntSequence canonicalize() const;
+
+private:
+  friend class StaticIntSequenceChainCursor;
+  llvm::SmallVector<StaticIntSequenceSlice> slices;
+  llvm::SmallVector<std::unique_ptr<StaticIntSequence>> owned;
+  size_t count = 0;
+};
+
+class StaticIntSequenceChainCursor {
+public:
+  explicit StaticIntSequenceChainCursor(const StaticIntSequenceChain &chain)
+    : chain(chain) {}
+
+  bool done() const { return slice == chain.slices.size(); }
+  int64_t value() const;
+  void advance();
+
+private:
+  const StaticIntSequenceChain &chain;
+  size_t slice = 0;
+  size_t offset = 0;
+};
+
+void setStaticIntSequenceAttr(mlir::Operation *op, llvm::StringRef name,
+                              const StaticIntSequence &sequence,
+                              size_t logicalCount);
+
+mlir::FailureOr<StaticIntSequence>
+getStaticIntSequenceAttr(mlir::Operation *op, llvm::StringRef name,
+                         size_t logicalCount);
+
+mlir::Value emitStaticIntLookup(const StaticIntSequence& sequence,
+                                mlir::Value position,
+                                mlir::Operation* constantAnchor,
+                                ConstantPool& constants,
+                                mlir::OpBuilder& builder,
+                                mlir::Location loc);
+
+} // namespace onnx_mlir

src/PIM/Common/IR/TensorSliceUtils.cpp

@@ -68,4 +68,56 @@ Value insertStaticSlice(
     .getResult();
 }
 
+Value extractMixedSliceOrIdentity(RewriterBase &rewriter,
+                                  Location loc,
+                                  Value source,
+                                  RankedTensorType resultType,
+                                  const MixedSliceGeometry &geometry) {
+  return extractStaticSliceOrIdentity(rewriter, loc, source, resultType,
+                                      geometry.offsets, geometry.sizes,
+                                      geometry.strides);
+}
+
+Value insertMixedSlice(OpBuilder &builder, Location loc, Value source,
+                       Value dest, const MixedSliceGeometry &geometry) {
+  return tensor::InsertSliceOp::create(builder, loc, source, dest,
+                                       geometry.offsets, geometry.sizes,
+                                       geometry.strides);
+}
+
+FailureOr<Value> addLeadingUnitTensorDimension(OpBuilder& builder, Location loc, Value value) {
+  auto type = dyn_cast<RankedTensorType>(value.getType());
+  if (!type || !type.hasStaticShape())
+    return failure();
+  SmallVector<int64_t> shape {1};
+  llvm::append_range(shape, type.getShape());
+  auto resultType = RankedTensorType::get(shape, type.getElementType(), type.getEncoding());
+  SmallVector<ReassociationIndices> reassociation;
+  if (type.getRank() != 0) {
+    reassociation.push_back({0, 1});
+    for (int64_t dim = 1; dim < type.getRank(); ++dim)
+      reassociation.push_back({dim + 1});
+  }
+  return tensor::ExpandShapeOp::create(builder, loc, resultType, value, reassociation).getResult();
+}
+
+FailureOr<Value> removeLeadingUnitTensorDimension(
+  OpBuilder& builder, Location loc, Value value, RankedTensorType resultType) {
+  if (value.getType() == resultType)
+    return value;
+  auto type = dyn_cast<RankedTensorType>(value.getType());
+  if (!type || !resultType || !type.hasStaticShape() || !resultType.hasStaticShape()
+      || type.getRank() != resultType.getRank() + 1 || type.getDimSize(0) != 1
+      || type.getElementType() != resultType.getElementType()
+      || !llvm::equal(type.getShape().drop_front(), resultType.getShape()))
+    return failure();
+  SmallVector<ReassociationIndices> reassociation;
+  if (resultType.getRank() != 0) {
+    reassociation.push_back({0, 1});
+    for (int64_t dim = 1; dim < resultType.getRank(); ++dim)
+      reassociation.push_back({dim + 1});
+  }
+  return tensor::CollapseShapeOp::create(builder, loc, resultType, value, reassociation).getResult();
+}
+
 } // namespace onnx_mlir

src/PIM/Common/IR/TensorSliceUtils.hpp

@@ -8,6 +8,12 @@
 
 namespace onnx_mlir {
 
+struct MixedSliceGeometry {
+  llvm::SmallVector<mlir::OpFoldResult> offsets;
+  llvm::SmallVector<mlir::OpFoldResult> sizes;
+  llvm::SmallVector<mlir::OpFoldResult> strides;
+};
+
 mlir::Value extractAxisSlice(
   mlir::PatternRewriter& rewriter, mlir::Location loc, mlir::Value source, int64_t axis, int64_t offset, int64_t size);
 
@@ -25,4 +31,22 @@ mlir::Value insertStaticSlice(mlir::PatternRewriter& rewriter,
                               mlir::Value dest,
                               llvm::ArrayRef<mlir::OpFoldResult> offsets);
 
+mlir::Value extractMixedSliceOrIdentity(mlir::RewriterBase &rewriter,
+                                        mlir::Location loc,
+                                        mlir::Value source,
+                                        mlir::RankedTensorType resultType,
+                                        const MixedSliceGeometry &geometry);
+
+mlir::Value insertMixedSlice(mlir::OpBuilder &builder,
+                             mlir::Location loc,
+                             mlir::Value source,
+                             mlir::Value dest,
+                             const MixedSliceGeometry &geometry);
+
+mlir::FailureOr<mlir::Value>
+addLeadingUnitTensorDimension(mlir::OpBuilder& builder, mlir::Location loc, mlir::Value value);
+
+mlir::FailureOr<mlir::Value> removeLeadingUnitTensorDimension(
+  mlir::OpBuilder& builder, mlir::Location loc, mlir::Value value, mlir::RankedTensorType resultType);
+
 } // namespace onnx_mlir

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/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);
@@ -1114,7 +1118,9 @@ enum class CompiledCoreOpKind : uint8_t {
 struct CompiledCoreNode {
   enum class Kind : uint8_t {
     Op,
-    Loop
+    Loop,
+    If,
+    IndexSwitch
   };
 
   Kind kind = Kind::Op;
@@ -1123,7 +1129,13 @@ 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;
+  llvm::SmallVector<int64_t> caseValues;
+  llvm::SmallVector<std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>>> caseBodies;
+  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> defaultBody;
 };
 
 static FailureOr<CompiledCoreOpKind> classifyCompiledCoreOpKind(Operation& op) {
@@ -1201,6 +1213,53 @@ 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;
+    }
+
+    if (auto switchOp = dyn_cast<mlir::scf::IndexSwitchOp>(op)) {
+      auto selector = compileIndexExpr(switchOp.getArg());
+      if (failed(selector)) {
+        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen");
+        return failure();
+      }
+      CompiledCoreNode switchNode;
+      switchNode.kind = CompiledCoreNode::Kind::IndexSwitch;
+      switchNode.op = switchOp.getOperation();
+      switchNode.condition = *selector;
+      llvm::append_range(switchNode.caseValues, switchOp.getCases());
+      for (mlir::Region& region : switchOp.getCaseRegions()) {
+        auto body = std::make_unique<llvm::SmallVector<CompiledCoreNode, 8>>();
+        if (failed(compileCoreEmissionPlan(region.front(), weightOwner, *body)))
+          return failure();
+        switchNode.caseBodies.push_back(std::move(body));
+      }
+      switchNode.defaultBody = std::make_unique<llvm::SmallVector<CompiledCoreNode, 8>>();
+      if (failed(compileCoreEmissionPlan(
+            switchOp.getDefaultRegion().front(), weightOwner, *switchNode.defaultBody)))
+        return failure();
+      plan.push_back(std::move(switchNode));
+      continue;
+    }
+
     auto opKind = classifyCompiledCoreOpKind(op);
     if (failed(opKind)) {
       InFlightDiagnostic diag = op.emitError() << "unsupported codegen for op '" << op.getName().getStringRef() << "'";
@@ -1263,6 +1322,51 @@ 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;
+    }
+
+    if (node.kind == CompiledCoreNode::Kind::IndexSwitch) {
+      auto selector = node.condition.evaluate(knowledge);
+      auto switchOp = cast<mlir::scf::IndexSwitchOp>(node.op);
+      if (failed(selector)) {
+        switchOp.emitOpError("requires a statically evaluable scf.index_switch selector for PIM codegen");
+        return failure();
+      }
+      const llvm::SmallVectorImpl<CompiledCoreNode>* selectedBody = node.defaultBody.get();
+      mlir::Region* selectedRegion = &switchOp.getDefaultRegion();
+      for (auto [index, caseValue] : llvm::enumerate(node.caseValues))
+        if (caseValue == *selector) {
+          selectedBody = node.caseBodies[index].get();
+          selectedRegion = &switchOp.getCaseRegions()[index];
+          break;
+        }
+      if (failed(executeCompiledCorePlan(*selectedBody, coreCodeGen, knowledge,
+                                         resolveWeightSlot, processedOperations,
+                                         batchLane, batchLaneCount)))
+        return failure();
+      auto yield = cast<mlir::scf::YieldOp>(selectedRegion->front().getTerminator());
+      for (auto [result, yielded] : llvm::zip(switchOp.getResults(), yield.getOperands()))
+        knowledge.aliases[result] = resolveLoopCarriedAlias(yielded, knowledge);
+      continue;
+    }
+
     switch (node.opKind) {
     case CompiledCoreOpKind::Load:
       coreCodeGen.codeGenLoadOp(cast<pim::PimMemCopyHostToDevOp>(node.op), knowledge);
@@ -1363,6 +1467,36 @@ static int64_t codeGenCoreOps(
   return failed(result) ? -1 : static_cast<int64_t>(processedOperations);
 }
 
+static OnnxMlirCompilerErrorCodes emitEmptyCoreArtifacts(StringRef outputDirPath, size_t emittedCoreId) {
+  std::string outputCorePath =
+    (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".pim").str();
+  std::error_code errorCode;
+  raw_fd_ostream coreBinaryStream(outputCorePath, errorCode, sys::fs::OF_None);
+  if (errorCode) {
+    errs() << "Error while opening core file `" << outputCorePath << "`: " << errorCode.message() << '\n';
+    return InvalidOutputFileAccess;
+  }
+
+  pim_binary::writeHeader(coreBinaryStream);
+  pim_binary::patchInstructionCount(coreBinaryStream, 0);
+  coreBinaryStream.close();
+
+  if (!pimEmitJson.getValue())
+    return CompilerSuccess;
+
+  std::string outputCoreJsonPath =
+    (outputDirPath + "/core_" + std::to_string(emittedCoreId) + ".json").str();
+  errorCode = std::error_code();
+  raw_fd_ostream coreJsonStream(outputCoreJsonPath, errorCode);
+  if (errorCode) {
+    errs() << "Error while opening core json file `" << outputCoreJsonPath << "`: " << errorCode.message() << '\n';
+    return InvalidOutputFileAccess;
+  }
+  coreJsonStream << "[]";
+  coreJsonStream.close();
+  return CompilerSuccess;
+}
+
 OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::string& outputDirPath) {
   if (!outputDirPath.empty()) {
     if (auto error = sys::fs::create_directory(outputDirPath)) {
@@ -1607,6 +1741,13 @@ OnnxMlirCompilerErrorCodes onnx_mlir::compileToPimCode(ModuleOp& moduleOp, std::
     if (jobResults[jobIndex].status != CompilerSuccess)
       return jobResults[jobIndex].status;
 
+  if (jobs.empty()) {
+    if (auto err = emitEmptyCoreArtifacts(outputDirPath, 0))
+      return err;
+    xbarsPerArrayGroup["core0"] = json::Array {};
+    memory.recordCoreReport(0, MemoryReportRow {});
+  }
+
   llvm::SmallVector<WeightFileRequest, 8> weightRequests;
   weightRequests.reserve(jobs.size());
   for (size_t jobIndex = 0; jobIndex < jobs.size(); ++jobIndex) {

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -57,6 +57,20 @@ llvm::cl::opt<PimConvLoweringType> pimConvLowering(
   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 for Spatial pipeline snapshots"),
+  llvm::cl::values(clEnumValN(SpatialDataflowExportNone, "none", "Do not emit Spatial dataflow CSV reports")),
+  llvm::cl::values(
+    clEnumValN(SpatialDataflowExportSpatial1, "spatial1", "Emit spatial1 graph dataflow CSV reports")),
+  llvm::cl::values(
+    clEnumValN(SpatialDataflowExportSpatial2, "spatial2", "Emit spatial2 scheduled dataflow CSV reports")),
+  llvm::cl::values(
+    clEnumValN(SpatialDataflowExportSpatial3, "spatial3", "Emit spatial3 realized dataflow CSV reports")),
+  llvm::cl::values(clEnumValN(SpatialDataflowExportAll, "all", "Emit all 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)"),

src/PIM/Compiler/PimCompilerOptions.hpp

@@ -42,11 +42,20 @@ typedef enum {
   PimConvLoweringTiled2D = 8,
 } PimConvLoweringType;
 
+typedef enum {
+  SpatialDataflowExportNone = 0,
+  SpatialDataflowExportSpatial1 = 1,
+  SpatialDataflowExportSpatial2 = 2,
+  SpatialDataflowExportSpatial3 = 3,
+  SpatialDataflowExportAll = 4,
+} 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;

src/PIM/Compiler/PimMemoryLiveness.cpp

@@ -291,7 +291,26 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ord
 
       if (auto yieldOp = dyn_cast<scf::YieldOp>(user)) {
         auto forOp = dyn_cast<scf::ForOp>(yieldOp->getParentOp());
-        if (!forOp) {
+        auto ifOp = dyn_cast<scf::IfOp>(yieldOp->getParentOp());
+        auto indexSwitch = dyn_cast<scf::IndexSwitchOp>(yieldOp->getParentOp());
+        if (ifOp) {
+          for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) {
+            if (operand != value)
+              continue;
+            pendingValues.push_back(ifOp.getResult(index));
+            appendAliasDescription(interval.aliasesFollowed, ifOp.getResult(index));
+          }
+        }
+        else if (indexSwitch) {
+          for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) {
+            if (operand != value)
+              continue;
+            pendingValues.push_back(indexSwitch.getResult(index));
+            appendAliasDescription(interval.aliasesFollowed,
+                                   indexSwitch.getResult(index));
+          }
+        }
+        else if (!forOp) {
           addFallbackReason(interval.fallbackReason, "yield without scf.for parent");
         }
         else {

src/PIM/Conversion/CMakeLists.txt

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

src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt

@@ -20,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
@@ -30,8 +31,10 @@ add_pim_library(OMONNXToSpatial
   SpatialLayoutPlanningPass.cpp
   LowerSpatialPlansPass.cpp
   Common/AttributeUtils.cpp
+  Common/BiasAddUtils.cpp
   Common/ComputeRegionBuilder.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/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,6 +137,31 @@ inline mlir::Value createSpatConcat(RewriterT& rewriter, mlir::Location loc, int
   return spatial::SpatConcatOp::create(rewriter, loc, outputType, rewriter.getI64IntegerAttr(axis), inputs).getOutput();
 }
 
+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>
@@ -97,16 +172,8 @@ auto createSpatGraphCompute(RewriterT& rewriter,
                             mlir::ValueRange inputs,
                             BodyFn&& body) {
   assert(inputs.size() == NumInputs && "NumInputs must match the number of input values");
-  auto computeOp = spatial::SpatGraphCompute::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>) {
@@ -140,16 +207,8 @@ auto createSpatGraphCompute(RewriterT& rewriter,
                             mlir::ValueRange weights,
                             mlir::ValueRange inputs,
                             BodyFn&& body) {
-  auto computeOp = spatial::SpatGraphCompute::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>) {
@@ -170,14 +229,15 @@ auto createSpatGraphCompute(RewriterT& rewriter,
   }
 }
 
-template <typename RewriterT, typename BodyFn>
-auto createSpatGraphComputeBatch(RewriterT& rewriter,
+template <typename RewriterT>
+auto createEmptySpatGraphComputeBatch(RewriterT& rewriter,
                                       mlir::Location loc,
                                       mlir::TypeRange resultTypes,
                                       int64_t laneCount,
                                       mlir::ValueRange weights,
                                       mlir::ValueRange inputs,
-                                 BodyFn&& body) {
+                                      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());
 
@@ -186,27 +246,36 @@ auto createSpatGraphComputeBatch(RewriterT& rewriter,
     return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
 
   auto batchOp = spatial::SpatGraphComputeBatch::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);
+  rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), blockArgTypes, blockArgLocs);
+  rewriter.setInsertionPointToStart(&batchOp.getBody().front());
+  return mlir::FailureOr<spatial::SpatGraphComputeBatch>(batchOp);
 }
 
-  auto* block =
-    rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), mlir::TypeRange(blockArgTypes), blockArgLocs);
-  rewriter.setInsertionPointToStart(block);
+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 createSpatGraphComputeBatch(RewriterT& rewriter,
+                                 mlir::Location loc,
+                                 mlir::TypeRange resultTypes,
+                                 int64_t laneCount,
+                                 mlir::ValueRange weights,
+                                 mlir::ValueRange inputs,
+                                 BodyFn&& body) {
+  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,18 +286,18 @@ auto createSpatGraphComputeBatch(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::SpatGraphComputeBatch>(batchOp);
+    rewriter.setInsertionPointAfter(*batchOp);
+    return batchOp;
   }
   else {
     auto bodyResult = std::forward<BodyFn>(body)(args);
     if (mlir::failed(bodyResult)) {
-      rewriter.setInsertionPointAfter(batchOp);
-      rewriter.eraseOp(batchOp);
+      rewriter.setInsertionPointAfter(*batchOp);
+      rewriter.eraseOp(*batchOp);
       return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
     }
-    rewriter.setInsertionPointAfter(batchOp);
-    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(batchOp);
+    rewriter.setInsertionPointAfter(*batchOp);
+    return batchOp;
   }
 }
 
@@ -277,6 +346,52 @@ inline void createParallelInsertSliceIntoBatchOutput(mlir::PatternRewriter& rewr
   mlir::tensor::ParallelInsertSliceOp::create(rewriter, loc, source, dest, offsets, sizes, strides);
 }
 
+inline void publishGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
+                                              mlir::Location loc,
+                                              mlir::Value fragment,
+                                              mlir::Value output,
+                                              mlir::Value physicalSlot) {
+  auto fragmentType = mlir::cast<mlir::RankedTensorType>(fragment.getType());
+  mlir::SmallVector<mlir::OpFoldResult> offsets {physicalSlot};
+  mlir::SmallVector<mlir::OpFoldResult> sizes {rewriter.getIndexAttr(1)};
+  mlir::SmallVector<mlir::OpFoldResult> strides {rewriter.getIndexAttr(1)};
+  for (int64_t dim : fragmentType.getShape()) {
+    offsets.push_back(rewriter.getIndexAttr(0));
+    sizes.push_back(rewriter.getIndexAttr(dim));
+    strides.push_back(rewriter.getIndexAttr(1));
+  }
+  createParallelInsertSliceIntoBatchOutput(rewriter, loc, fragment, output, offsets, sizes, strides);
+}
+
+inline mlir::FailureOr<mlir::Value>
+extractGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
+                                  mlir::Location loc,
+                                  mlir::Value physicalBatch,
+                                  mlir::OpFoldResult slot,
+                                  mlir::RankedTensorType fragmentType) {
+  if (fragmentType.getRank() == 0)
+    return mlir::failure();
+  auto physicalType = mlir::dyn_cast<mlir::RankedTensorType>(physicalBatch.getType());
+  if (!physicalType || physicalType.getRank() != fragmentType.getRank() + 1)
+    return mlir::failure();
+  mlir::SmallVector<int64_t> selectedShape {1};
+  llvm::append_range(selectedShape, fragmentType.getShape());
+  auto selectedType = mlir::RankedTensorType::get(selectedShape, fragmentType.getElementType(), fragmentType.getEncoding());
+  mlir::SmallVector<mlir::OpFoldResult> offsets {slot};
+  mlir::SmallVector<mlir::OpFoldResult> sizes {rewriter.getIndexAttr(1)};
+  mlir::SmallVector<mlir::OpFoldResult> strides {rewriter.getIndexAttr(1)};
+  for (int64_t dim : fragmentType.getShape()) {
+    offsets.push_back(rewriter.getIndexAttr(0));
+    sizes.push_back(rewriter.getIndexAttr(dim));
+    strides.push_back(rewriter.getIndexAttr(1));
+  }
+  mlir::Value selected = mlir::tensor::ExtractSliceOp::create(rewriter, loc, selectedType, physicalBatch, offsets, sizes, strides);
+  mlir::SmallVector<mlir::ReassociationIndices> reassociation {{0, 1}};
+  for (int64_t dim = 2; dim <= fragmentType.getRank(); ++dim)
+    reassociation.push_back({dim});
+  return mlir::tensor::CollapseShapeOp::create(rewriter, loc, fragmentType, selected, reassociation).getResult();
+}
+
 template <typename BodyFn>
 mlir::Value materializeOrComputeUnary(mlir::Value input,
                                       mlir::RankedTensorType resultType,

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

@@ -0,0 +1,211 @@
+#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 spatial::getGraphBatchPhysicalResultType(logicalType.getDimSize(2), getRowStripFragmentType(logicalType));
+}
+
+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};
+}
+
+Value extractRowStripFragment(Value storage,
+                              RankedTensorType logicalType,
+                              OpFoldResult row,
+                              PatternRewriter& rewriter,
+                              Location loc) {
+  return *extractGraphBatchPhysicalFragment(rewriter, loc, storage, row, getRowStripFragmentType(logicalType));
+}
+
+void insertRowStripFragment(Value fragment,
+                            Value output,
+                            RankedTensorType logicalType,
+                            OpFoldResult row,
+                            PatternRewriter& rewriter,
+                            Location loc) {
+  assert(fragment.getType() == getRowStripFragmentType(logicalType));
+  assert(output.getType() == getRowStripStorageType(logicalType));
+  auto slot = dyn_cast<Value>(row);
+  assert(slot && "row-strip graph publication requires a dynamic physical slot");
+  publishGraphBatchPhysicalFragment(rewriter, loc, fragment, output, slot);
+}
+
+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>
+createRowStripAssemblyBlueprint(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
+  auto storageType = dyn_cast<RankedTensorType>(storage.getType());
+  if (!storageType || storageType != getRowStripStorageType(logicalType))
+    return failure();
+
+  auto [offsets, sizes] = buildRowStripMetadata(logicalType);
+  int64_t height = logicalType.getDimSize(2);
+  SmallVector<int64_t> operandIndices(height, 0), sourceSlots, sourceOffsets(height, 0), strides(height * 4, 1);
+  for (int64_t row = 0; row < height; ++row)
+    sourceSlots.push_back(row);
+  return spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, storage, ValueRange {},
+    rewriter.getStringAttr("nchw"), rewriter.getStringAttr("nchw_row_strip"),
+    rewriter.getDenseI64ArrayAttr(offsets), rewriter.getDenseI64ArrayAttr(sizes),
+    rewriter.getStringAttr("nchw_row_strip_fragments"), rewriter.getStringAttr("fragment_assembly"),
+    rewriter.getDenseI64ArrayAttr(operandIndices), rewriter.getDenseI64ArrayAttr(sourceSlots),
+    rewriter.getDenseI64ArrayAttr(sourceOffsets), rewriter.getDenseI64ArrayAttr(strides),
+    rewriter.getStringAttr("disjoint"), rewriter.getStringAttr("complete")).getOutput();
+}
+
+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> createRowStripAssemblyBlueprint(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/CompileTime.cpp

@@ -9,10 +9,12 @@
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 
+#include <cstring>
 #include <utility>
 
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -21,24 +23,7 @@ using namespace mlir;
 
 namespace onnx_mlir {
 
-namespace {
-
-static bool hasStaticUnitStrides(tensor::ExtractSliceOp extractSliceOp) {
-  return llvm::all_of(extractSliceOp.getStaticStrides(), [](int64_t stride) { return stride == 1; });
-}
-
-static bool hasConstantIndices(tensor::ExtractOp extractOp) {
-  return llvm::all_of(extractOp.getIndices(), [](Value index) { return matchConstantIndexValue(index).has_value(); });
-}
-
-static bool isStaticTensorResult(Operation* op) {
-  return llvm::all_of(op->getResultTypes(), [](Type type) {
-    auto shapedType = dyn_cast<ShapedType>(type);
-    return shapedType && shapedType.hasStaticShape();
-  });
-}
-
-static FailureOr<DenseElementsAttr> transposeDenseElements(DenseElementsAttr denseAttr, ArrayRef<int64_t> perms) {
+FailureOr<DenseElementsAttr> transposeDenseElementsAttr(DenseElementsAttr denseAttr, ArrayRef<int64_t> perms) {
   auto tensorType = dyn_cast<RankedTensorType>(denseAttr.getType());
   if (!tensorType)
     return failure();
@@ -59,7 +44,45 @@ static FailureOr<DenseElementsAttr> transposeDenseElements(DenseElementsAttr den
 
   auto transposedType = RankedTensorType::get(transposedShape, tensorType.getElementType(), tensorType.getEncoding());
   if (denseAttr.isSplat())
-    return DenseElementsAttr::get(transposedType, denseAttr.getSplatValue<Attribute>());
+    return DenseElementsAttr::getFromRawBuffer(transposedType, denseAttr.getRawData());
+
+  const unsigned elementBitWidth = tensorType.getElementTypeBitWidth();
+  const ArrayRef<char> inputData = denseAttr.getRawData();
+  if (elementBitWidth % 8 == 0) {
+    const size_t elementBytes = elementBitWidth / 8;
+    const size_t expectedBytes = denseAttr.getNumElements() * elementBytes;
+    if (inputData.size() == expectedBytes) {
+      SmallVector<char> transposedData(expectedBytes);
+      if (rank == 2 && perms[0] == 1 && perms[1] == 0) {
+        const int64_t rows = tensorType.getDimSize(0);
+        const int64_t columns = tensorType.getDimSize(1);
+        for (int64_t row = 0; row < rows; ++row)
+          for (int64_t column = 0; column < columns; ++column)
+            std::memcpy(transposedData.data() + (column * rows + row) * elementBytes,
+                        inputData.data() + (row * columns + column) * elementBytes,
+                        elementBytes);
+        return DenseElementsAttr::getFromRawBuffer(transposedType, transposedData);
+      }
+
+      SmallVector<int64_t> originalStrides = computeRowMajorStrides(tensorType.getShape());
+      SmallVector<int64_t> transposedStrides = computeRowMajorStrides(transposedShape);
+      SmallVector<int64_t> originalIndices(rank);
+      for (int64_t linearIndex = 0; linearIndex < tensorType.getNumElements(); ++linearIndex) {
+        int64_t remaining = linearIndex;
+        for (int64_t dim = 0; dim < rank; ++dim) {
+          originalIndices[dim] = originalStrides.empty() ? 0 : remaining / originalStrides[dim];
+          remaining = originalStrides.empty() ? 0 : remaining % originalStrides[dim];
+        }
+        int64_t transposedLinearIndex = 0;
+        for (int64_t dim = 0; dim < rank; ++dim)
+          transposedLinearIndex += originalIndices[perms[dim]] * transposedStrides[dim];
+        std::memcpy(transposedData.data() + transposedLinearIndex * elementBytes,
+                    inputData.data() + linearIndex * elementBytes,
+                    elementBytes);
+      }
+      return DenseElementsAttr::getFromRawBuffer(transposedType, transposedData);
+    }
+  }
 
   SmallVector<Attribute> originalValues(denseAttr.getValues<Attribute>());
   SmallVector<Attribute> transposedValues(originalValues.size());
@@ -84,16 +107,30 @@ static FailureOr<DenseElementsAttr> transposeDenseElements(DenseElementsAttr den
   return DenseElementsAttr::get(transposedType, transposedValues);
 }
 
+namespace {
+
+static bool hasStaticUnitStrides(tensor::ExtractSliceOp extractSliceOp) {
+  return llvm::all_of(extractSliceOp.getStaticStrides(), [](int64_t stride) { return stride == 1; });
+}
+
+static bool hasConstantIndices(tensor::ExtractOp extractOp) {
+  return llvm::all_of(extractOp.getIndices(), [](Value index) { return matchConstantIndexValue(index).has_value(); });
+}
+
+static bool isStaticTensorResult(Operation* op) {
+  return llvm::all_of(op->getResultTypes(), [](Type type) {
+    auto shapedType = dyn_cast<ShapedType>(type);
+    return shapedType && shapedType.hasStaticShape();
+  });
+}
+
 static FailureOr<DenseElementsAttr> reshapeDenseElements(DenseElementsAttr denseAttr, RankedTensorType resultType) {
   auto sourceType = dyn_cast<RankedTensorType>(denseAttr.getType());
-  if (!sourceType || !resultType || sourceType.getNumElements() != resultType.getNumElements())
+  if (!sourceType || !resultType || sourceType.getNumElements() != resultType.getNumElements()
+      || sourceType.getElementType() != resultType.getElementType())
     return failure();
 
-  if (denseAttr.isSplat())
-    return DenseElementsAttr::get(resultType, denseAttr.getSplatValue<Attribute>());
-
-  SmallVector<Attribute> values(denseAttr.getValues<Attribute>());
-  return DenseElementsAttr::get(resultType, values);
+  return DenseElementsAttr::getFromRawBuffer(resultType, denseAttr.getRawData());
 }
 
 static FailureOr<DenseElementsAttr> extractSliceDenseElements(DenseElementsAttr denseAttr,
@@ -161,7 +198,7 @@ static DenseElementsAttr getHostConstantDenseElementsAttrImpl(Value value, llvm:
     perm.reserve(transposeOp.getPermAttr().size());
     for (IntegerAttr attr : transposeOp.getPermAttr().getAsRange<IntegerAttr>())
       perm.push_back(attr.getInt());
-    auto transposedAttr = transposeDenseElements(inputAttr, perm);
+    auto transposedAttr = transposeDenseElementsAttr(inputAttr, perm);
     return succeeded(transposedAttr) ? *transposedAttr : nullptr;
   }
 
@@ -171,7 +208,7 @@ static DenseElementsAttr getHostConstantDenseElementsAttrImpl(Value value, llvm:
       return nullptr;
 
     SmallVector<int64_t> perm(transposeOp.getPermutation().begin(), transposeOp.getPermutation().end());
-    auto transposedAttr = transposeDenseElements(inputAttr, perm);
+    auto transposedAttr = transposeDenseElementsAttr(inputAttr, perm);
     return succeeded(transposedAttr) ? *transposedAttr : nullptr;
   }
 
@@ -219,6 +256,9 @@ getCompileTimeSourceImpl(Operation* op, llvm::SmallPtrSetImpl<Operation*>& visit
 
   chainLength += 1;
 
+  if (!isShapingOnlyOp(op))
+    return std::nullopt;
+
   if (auto extractOp = dyn_cast<tensor::ExtractOp>(op))
     return hasConstantIndices(extractOp)
            ? getCompileTimeSourceImpl(extractOp.getTensor().getDefiningOp(), visited, chainLength)

src/PIM/Conversion/ONNXToSpatial/CompileTime.hpp

@@ -4,6 +4,8 @@
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/Value.h"
 
+#include "llvm/ADT/ArrayRef.h"
+
 namespace onnx_mlir {
 
 struct CompileTimeSource {
@@ -19,4 +21,7 @@ bool isCompileTimeOp(mlir::Operation* op);
 
 mlir::DenseElementsAttr getHostConstDenseElementsAttr(mlir::Value value);
 
+mlir::FailureOr<mlir::DenseElementsAttr> transposeDenseElementsAttr(
+  mlir::DenseElementsAttr denseAttr, llvm::ArrayRef<int64_t> permutation);
+
 } // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -11,12 +11,16 @@
 #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/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
 #include "src/Accelerators/PIM/Pass/PIMPasses.h"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
 
@@ -28,14 +32,6 @@ namespace {
 static constexpr StringLiteral kDenseLayout = "dense_nchw";
 static constexpr StringLiteral kRowStripLayout = "nchw_row_strip";
 
-struct RowStripPhysicalValue {
-  Value physicalValue;
-  RankedTensorType logicalType;
-  SmallVector<int64_t, 16> fragmentOffsets;
-  SmallVector<int64_t, 16> fragmentSizes;
-  std::string indexMap;
-};
-
 static FailureOr<RowStripPhysicalValue> getRowStripValue(llvm::DenseMap<Value, RowStripPhysicalValue>& rowStripValues,
                                                          Value value) {
   auto it = rowStripValues.find(value);
@@ -45,112 +41,42 @@ static FailureOr<RowStripPhysicalValue> getRowStripValue(llvm::DenseMap<Value, R
 }
 
 static FailureOr<RowStripPhysicalValue> buildRowStripValue(spatial::SpatBlueprintOp blueprint,
-                                                           Value physicalValue) {
+                                                           Value storage) {
   auto logicalType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
   if (!logicalType)
     return blueprint.emitOpError("requires ranked logical output type"), failure();
   RowStripPhysicalValue value;
-  value.physicalValue = physicalValue;
+  value.storage = storage;
   value.logicalType = logicalType;
   value.fragmentOffsets.append(blueprint.getFragmentOffsets().begin(), blueprint.getFragmentOffsets().end());
   value.fragmentSizes.append(blueprint.getFragmentSizes().begin(), blueprint.getFragmentSizes().end());
-  value.indexMap = blueprint.getIndexMap().str();
+  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) {
-  auto packedType = cast<RankedTensorType>(input.physicalValue.getType());
-  auto computeOp =
-    createSpatCompute<1>(rewriter, planOp.getLoc(), TypeRange {packedType}, {}, input.physicalValue, [&](Value x) {
-      auto relu = spatial::SpatReluOp::create(rewriter, planOp.getLoc(), packedType, x);
-      spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), relu.getResult());
-    });
-  return computeOp.getResult(0);
+  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) {
-  auto packedType = dyn_cast<RankedTensorType>(rowStripValue.physicalValue.getType());
-  if (!packedType || packedType.getRank() != 3 || !packedType.hasStaticShape())
-    return failure();
   if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape())
     return failure();
-  if (rowStripValue.indexMap != "packed_hwc_rows_to_nchw")
+  auto [expectedOffsets, expectedSizes] = buildRowStripMetadata(rowStripValue.logicalType);
+  if (!llvm::equal(rowStripValue.fragmentOffsets, expectedOffsets) || !llvm::equal(rowStripValue.fragmentSizes, expectedSizes))
     return failure();
-
-  const int64_t rank = rowStripValue.logicalType.getRank();
-  const int64_t fragmentCount = rowStripValue.fragmentOffsets.size() / rank;
-  const int64_t packedWidth = packedType.getDimSize(1);
-  const int64_t packedChannels = packedType.getDimSize(2);
-  if (fragmentCount != packedType.getDimSize(0))
-    return failure();
-  for (int64_t fragmentIndex = 0; fragmentIndex < fragmentCount; ++fragmentIndex) {
-    if (rowStripValue.fragmentOffsets[fragmentIndex * rank + 0] != 0
-        || rowStripValue.fragmentOffsets[fragmentIndex * rank + 1] != 0
-        || rowStripValue.fragmentOffsets[fragmentIndex * rank + 2] != fragmentIndex
-        || rowStripValue.fragmentOffsets[fragmentIndex * rank + 3] != 0)
-      return failure();
-    if (rowStripValue.fragmentSizes[fragmentIndex * rank + 0] != 1
-        || rowStripValue.fragmentSizes[fragmentIndex * rank + 1] != packedChannels
-        || rowStripValue.fragmentSizes[fragmentIndex * rank + 2] != 1
-        || rowStripValue.fragmentSizes[fragmentIndex * rank + 3] != packedWidth)
-      return failure();
-  }
-
-  auto packedSliceType =
-    RankedTensorType::get({1, packedWidth, packedChannels}, packedType.getElementType(), packedType.getEncoding());
-  auto expandedType =
-    RankedTensorType::get({1, 1, packedWidth, packedChannels}, packedType.getElementType(), packedType.getEncoding());
-  auto logicalFragmentType =
-    RankedTensorType::get({1, packedChannels, 1, packedWidth}, packedType.getElementType(), packedType.getEncoding());
-  auto batchOp = createSpatComputeBatch(
-    rewriter,
-    loc,
-    TypeRange {rowStripValue.logicalType},
-    fragmentCount,
-    {},
-    ValueRange {rowStripValue.physicalValue},
-    [&](detail::SpatComputeBatchBodyArgs args) {
-      SmallVector<OpFoldResult> packedOffsets {args.lane, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> packedSizes {
-        rewriter.getIndexAttr(1), rewriter.getIndexAttr(packedWidth), rewriter.getIndexAttr(packedChannels)};
-      Value packedSlice = tensor::ExtractSliceOp::create(
-        rewriter, loc, packedSliceType, args.inputs.front(), packedOffsets, packedSizes, getUnitStrides(rewriter, 3));
-
-      Value expanded = tensor::ExpandShapeOp::create(rewriter,
-                                                     loc,
-                                                     expandedType,
-                                                     packedSlice,
-                                                     SmallVector<ReassociationIndices> {
-                                                       {0, 1},
-                                                       {2},
-                                                       {3}
-      });
-      Value transposeInit =
-        tensor::EmptyOp::create(rewriter, loc, logicalFragmentType.getShape(), logicalFragmentType.getElementType());
-      Value logicalFragment =
-        linalg::TransposeOp::create(rewriter, loc, expanded, transposeInit, SmallVector<int64_t> {0, 3, 1, 2})
-          .getResult()[0];
-
-      SmallVector<OpFoldResult> logicalOffsets {
-        rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), args.lane, rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> logicalSizes {rewriter.getIndexAttr(1),
-                                              rewriter.getIndexAttr(packedChannels),
-                                              rewriter.getIndexAttr(1),
-                                              rewriter.getIndexAttr(packedWidth)};
-      createParallelInsertSliceIntoBatchOutput(rewriter,
-                                               loc,
-                                               logicalFragment,
-                                               args.outputs.front(),
-                                               logicalOffsets,
-                                               logicalSizes,
-                                               getUnitStrides(rewriter, 4));
-      return success();
-    });
-  if (failed(batchOp))
-    return failure();
-  return batchOp->getResult(0);
+  return createRowStripAssemblyBlueprint(rowStripValue.storage, rowStripValue.logicalType, rewriter, loc);
 }
 
 struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
@@ -193,7 +119,7 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
           rewriter.setInsertionPoint(planOp);
           FailureOr<Value> lowered = lowerSelectedConv2DPlan(
             planOp,
-            succeeded(rowStripInput) ? std::optional<Value> {rowStripInput->physicalValue} : std::nullopt,
+            succeeded(rowStripInput) ? std::optional<Value> {rowStripInput->storage} : std::nullopt,
             /*emitRowStripLayout=*/true,
             rewriter);
           if (failed(lowered)) {
@@ -265,6 +191,64 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
         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) {
@@ -294,6 +278,8 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
         continue;
       }
       if (auto blueprintOp = dyn_cast<spatial::SpatBlueprintOp>(&op)) {
+        if (std::optional<StringRef> mode = blueprintOp.getMode(); mode && *mode == "fragment_assembly")
+          continue;
         if (blueprintOp.getPhysicalLayout() == kDenseLayout) {
           rewriter.replaceOp(blueprintOp, blueprintOp.getInput());
           continue;
@@ -345,17 +331,25 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
     RewritePatternSet helperPatterns(ctx);
     populateGemmPatterns(helperPatterns, ctx);
     populateTransposePatterns(helperPatterns, ctx);
-    if (failed(applyPartialConversion(moduleOp, helperTarget, std::move(helperPatterns)))) {
+    FrozenRewritePatternSet frozenHelperPatterns(
+      std::move(helperPatterns));
+    SmallVector<Operation*> topLevelHelperOps;
+    funcOp.walk([&](Operation* op) {
+      if (isa<spatial::SpatGraphCompute,
+              spatial::SpatGraphComputeBatch>(op))
+        return WalkResult::skip();
+      if (isa<ONNXGemmOp, ONNXTransposeOp>(op))
+        topLevelHelperOps.push_back(op);
+      return WalkResult::advance();
+    });
+    for (Operation *helper : topLevelHelperOps) {
+      if (failed(applyPartialConversion(
+            helper, helperTarget, frozenHelperPatterns))) {
         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,
@@ -371,7 +365,8 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
         computeLikeOps.push_back(op);
     });
     for (Operation* op : computeLikeOps) {
-      if (failed(applyFullConversion(op, nestedHelperTarget, nestedHelperPatterns))) {
+      if (failed(applyFullConversion(
+            op, nestedHelperTarget, frozenHelperPatterns))) {
         op->emitOpError("failed to lower nested helper ONNX ops emitted by selected Spatial plan lowering");
         signalPassFailure();
         return;
@@ -383,19 +378,37 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
     moduleOp.walk([&](Operation* op) {
       if (isa<ONNXEntryPointOp>(op))
         return;
-      if (isa<spatial::SpatConv2DPlanOp,
+      if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
+        if (std::optional<StringRef> mode = blueprint.getMode(); mode && *mode == "fragment_assembly")
+          return;
+        op->emitOpError("planning blueprint must not remain after LowerSpatialPlans");
+        hasIllegalOps = true;
+      } else 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;
       }
     });
-    if (hasIllegalOps)
+
+    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_premerge");
+    } else {
+      dumpModule(moduleOp, "spatial1_graph");
+      spatial::SpatialDataflowExportStage exportMode = spatial::getSpatialDataflowExportStage();
+      if (spatial::shouldExportSpatialDataflowStage(exportMode, spatial::SpatialDataflowExportStage::Spatial1)
+          && failed(spatial::exportSpatialDataflowCsvGraph(funcOp, "spatial1_graph"))) {
+        signalPassFailure();
+        return;
+      }
+    }
 
     if (!verifyLogicalPhase("at the end of LowerSpatialPlans"))
       return;

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -13,6 +13,7 @@
 
 #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"
@@ -45,11 +46,12 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
   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() || !reluPlans.empty() || !blueprints.empty()
-      || !materializers.empty()) {
+  if (!computes.empty() || !computeBatches.empty() || !convPlans.empty() || !biasAddPlans.empty() || !reluPlans.empty()
+      || !blueprints.empty() || !materializers.empty()) {
     return;
   }
 
@@ -65,9 +67,9 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
     sourceLocs.push_back(source.getLoc());
   }
 
-  auto newCompute = spatial::SpatGraphCompute::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())});
@@ -103,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);
@@ -142,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>();
@@ -173,11 +176,6 @@ 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))) {
@@ -213,13 +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;
     }
-  dumpModule(moduleOp, "spatial0");
-
   if (failed(verifyONNXToSpatial(*entryFunc))) {
     moduleOp.emitError("ONNX-to-Spatial host legality verification failed");
     signalPassFailure();

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp

@@ -56,13 +56,18 @@ bool isLegalExternalCapture(Value value, Region& region) {
   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))
+      if (isLegalExternalCapture(value, body) || isRecordedDeferredCommunicationSource(nestedOp, value))
         continue;
 
       Operation* definingOp = value.getDefiningOp();
@@ -90,21 +95,29 @@ bool isLegalHostBackedValue(Value value) {
   return definingOp->getDialect()->getNamespace() != "spat";
 }
 
+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(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(compute.getOperation(), [&](Operation* illegalOp) {
       InFlightDiagnostic diag = illegalOp->emitOpError()
-                                << kPhaseMarker << " " << kind << " input #" << inputIndex
+                                << kPhaseMarker << " " << kind << " input #" << currentInputIndex
                                 << (allowChannelReceiveInputs ? " must come from the host or explicit spat.channel_receive"
                                                               : " must come from the host");
       if (definingOp)
@@ -132,6 +145,7 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
             spatial::SpatGraphCompute,
             spatial::SpatGraphComputeBatch,
             spatial::SpatConv2DPlanOp,
+            spatial::SpatBiasAddPlanOp,
             spatial::SpatReluPlanOp,
             spatial::SpatBlueprintOp,
             spatial::SpatMaterializeLayoutOp>(&op)) {
@@ -162,9 +176,9 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
 
 void verifyScheduledTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter& diagnostics) {
   for (Operation& op : funcOp.getOps()) {
-    if (isa<spatial::SpatGraphCompute, spatial::SpatGraphComputeBatch>(&op)) {
+    if (isa<spatial::SpatChannelSendOp, spatial::SpatChannelReceiveOp>(&op)) {
       diagnostics.report(&op, [&](Operation* illegalOp) {
-        illegalOp->emitOpError() << kPhaseMarker << " graph Spatial compute op remained after merge materialization";
+        illegalOp->emitOpError() << kPhaseMarker << " real channel communication is not allowed in scheduled phase 1";
       });
     }
   }

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);

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

@@ -251,10 +251,7 @@ static FailureOr<spatial::SpatComputeBatch> createVmmBatch(Value a,
         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)};
-      SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
-      createParallelInsertSliceIntoBatchOutput(
-        rewriter, loc, piece, args.outputs.front(), pieceOffsets, pieceSizes, unitStrides);
+      publishGraphBatchPhysicalFragment(rewriter, loc, piece, args.outputs.front(), args.lane);
     });
   if (failed(batchOp))
     return failure();
@@ -401,11 +398,7 @@ static FailureOr<spatial::SpatComputeBatch> createVvdmulBatch(Value a,
       Value bVector = extractDynamicGemmBColumn(args.inputs[1], column, vectorType, rewriter, loc);
       Value scalar = spatial::SpatVVDMulOp::create(rewriter, loc, scalarType, aVector, bVector).getResult();
 
-      SmallVector<OpFoldResult> outputOffsets {args.lane, rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-      SmallVector<OpFoldResult> unitStrides = getUnitStrides(rewriter, 2);
-      createParallelInsertSliceIntoBatchOutput(
-        rewriter, loc, scalar, args.outputs.front(), outputOffsets, scalarSizes, unitStrides);
+      publishGraphBatchPhysicalFragment(rewriter, loc, scalar, args.outputs.front(), args.lane);
     });
   if (failed(batchOp))
     return failure();
@@ -447,15 +440,14 @@ static FailureOr<spatial::SpatCompute> createDynamicGemmOutputCompute(Value scal
         Value row = createDynamicGemmBatchRow(lane, numOutCols, rewriter, nestedLoc);
         Value column =
           onnx_mlir::affineModConst(rewriter, nestedLoc, lane, numOutCols, rewriter.getInsertionBlock()->getParentOp());
-        SmallVector<OpFoldResult> scalarOffsets {lane, rewriter.getIndexAttr(0)};
         SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
         SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-        Value scalar = tensor::ExtractSliceOp::create(
-                         rewriter, nestedLoc, scalarType, pieces, scalarOffsets, scalarSizes, unitStrides)
-                         .getResult();
+        FailureOr<Value> scalar = extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, scalarType);
+        if (failed(scalar))
+          return failure();
         if (alpha != 1.0f) {
           Value alphaTensor = createScalarTensorConstant(scalarType, alpha, rewriter, nestedLoc);
-          scalar = spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, scalar, alphaTensor).getResult();
+          *scalar = spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, *scalar, alphaTensor).getResult();
         }
         if (biasArg) {
           Value biasScalar =
@@ -465,11 +457,11 @@ static FailureOr<spatial::SpatCompute> createDynamicGemmOutputCompute(Value scal
             biasScalar =
               spatial::SpatVMulOp::create(rewriter, nestedLoc, scalarType, biasScalar, betaTensor).getResult();
           }
-          scalar = spatial::SpatVAddOp::create(rewriter, nestedLoc, scalarType, scalar, biasScalar).getResult();
+          *scalar = spatial::SpatVAddOp::create(rewriter, nestedLoc, scalarType, *scalar, biasScalar).getResult();
         }
         SmallVector<OpFoldResult> outputOffsets {row, column};
         Value outputNext =
-          tensor::InsertSliceOp::create(rewriter, nestedLoc, scalar, outputAcc, outputOffsets, scalarSizes, unitStrides)
+          tensor::InsertSliceOp::create(rewriter, nestedLoc, *scalar, outputAcc, outputOffsets, scalarSizes, unitStrides)
             .getResult();
         yielded.push_back(outputNext);
         return success();
@@ -505,14 +497,13 @@ static Value extractReductionPiece(Value partialPiecesArg,
                                    int64_t numOutRows,
                                    ConversionPatternRewriter& rewriter,
                                    Location loc) {
-  SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-  SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(numOutRows),
-                                        rewriter.getIndexAttr(crossbarSize.getValue())};
+  SmallVector<OpFoldResult> unitStrides {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
+  SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
   SmallVector<OpFoldResult> pieceOffsets {
-    createPartialGroupOffset(hSlice, kSlice, numKSlices, numOutRows, rewriter, loc), rewriter.getIndexAttr(0)};
-  return tensor::ExtractSliceOp::create(
-           rewriter, loc, pieceType, partialPiecesArg, pieceOffsets, pieceSizes, unitStrides)
-    .getResult();
+    createPartialGroupOffset(hSlice, kSlice, numKSlices, numOutRows, rewriter, loc), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+  auto selectedType = RankedTensorType::get({numOutRows, 1, static_cast<int64_t>(crossbarSize.getValue())}, pieceType.getElementType());
+  Value selected = tensor::ExtractSliceOp::create(rewriter, loc, selectedType, partialPiecesArg, pieceOffsets, pieceSizes, unitStrides);
+  return tensor::CollapseShapeOp::create(rewriter, loc, pieceType, selected, SmallVector<ReassociationIndices> {{0, 1}, {2}});
 }
 
 static Value reducePartialPiecesForHSlice(Value partialPiecesArg,
@@ -730,7 +721,7 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
       return failure();
     }
 
-    auto scalarPiecesType = RankedTensorType::get({laneCount64, 1}, outType.getElementType());
+    auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(laneCount64, RankedTensorType::get({1, 1}, outType.getElementType()));
     auto batchOp = createVvdmulBatch(a, b, aType, bType, scalarPiecesType, outType, rewriter, loc);
     if (failed(batchOp))
       return failure();
@@ -802,8 +793,8 @@ LogicalResult GemmToSpatialComputes::matchAndRewrite(ONNXGemmOp gemmOp,
     return failure();
   }
 
-  auto partialPiecesType =
-    RankedTensorType::get({laneCount64, static_cast<int64_t>(crossbarSize.getValue())}, outType.getElementType());
+  auto partialPiecesType = spatial::getGraphBatchPhysicalResultType(
+    laneCount64, RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, outType.getElementType()));
   auto batchOp =
     createVmmBatch(a, b, aType, paddedBType, partialPiecesType, numOutRows, numKSlices, numOutHSlices, rewriter, loc);
   if (failed(batchOp))

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

@@ -398,10 +398,7 @@ static FailureOr<spatial::SpatComputeBatch> createBatchedVmmBatch(Value a,
         args.weights.front(), bBatchShape, outputBatchShape, batch, kOffset, hOffset, bTileType, rewriter, loc);
       Value piece = spatial::SpatVMMOp::create(rewriter, loc, pieceType, bTile, aTile).getResult();
 
-      SmallVector<OpFoldResult> pieceOffsets {args.lane, rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
-      createParallelInsertSliceIntoBatchOutput(
-        rewriter, loc, piece, args.outputs.front(), pieceOffsets, pieceSizes, getUnitStrides(rewriter, 2));
+      publishGraphBatchPhysicalFragment(rewriter, loc, piece, args.outputs.front(), args.lane);
     });
   if (failed(batchOp))
     return failure();
@@ -506,10 +503,7 @@ static FailureOr<spatial::SpatComputeBatch> createBatchedVvdmulBatch(Value a,
       Value bVector = extractDynamicBatchedBColumn(
         args.inputs[1], bBatchShape, outputBatchShape, batch, column, vectorType, rewriter, loc);
       Value scalar = spatial::SpatVVDMulOp::create(rewriter, loc, scalarType, aVector, bVector).getResult();
-      SmallVector<OpFoldResult> outputOffsets {args.lane, rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-      createParallelInsertSliceIntoBatchOutput(
-        rewriter, loc, scalar, args.outputs.front(), outputOffsets, scalarSizes, getUnitStrides(rewriter, 2));
+      publishGraphBatchPhysicalFragment(rewriter, loc, scalar, args.outputs.front(), args.lane);
     });
   if (failed(batchOp))
     return failure();
@@ -548,14 +542,13 @@ static FailureOr<Value> createBatchedDynamicOutputCompute(Value scalarPieces,
           Value batchLane = affineModConst(rewriter, nestedLoc, lane, numOutRows * numOutCols, anchorOp);
           Value row = affineFloorDivConst(rewriter, nestedLoc, batchLane, numOutCols, anchorOp);
           Value column = affineModConst(rewriter, nestedLoc, batchLane, numOutCols, anchorOp);
-          SmallVector<OpFoldResult> scalarOffsets {lane, rewriter.getIndexAttr(0)};
-          SmallVector<OpFoldResult> scalarSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
-          Value scalar = tensor::ExtractSliceOp::create(
-            rewriter, nestedLoc, scalarType, pieces, scalarOffsets, scalarSizes, getUnitStrides(rewriter, 2));
+          FailureOr<Value> scalar = extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, scalarType);
+          if (failed(scalar))
+            return failure();
           Value expanded = tensor::ExpandShapeOp::create(rewriter,
                                                          nestedLoc,
                                                          outputScalarType,
-                                                         scalar,
+                                                         *scalar,
                                                          SmallVector<ReassociationIndices> {
                                                            {0},
                                                            {1, 2}
@@ -596,10 +589,11 @@ static Value extractBatchedReductionPiece(Value partialPiecesArg,
   Value kOffset = getOrCreateIndexConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), kSlice * numOutRows);
   Value batchAndHSlice = arith::AddIOp::create(rewriter, loc, batchOffset, hOffset);
   Value pieceOffset = arith::AddIOp::create(rewriter, loc, batchAndHSlice, kOffset);
-  SmallVector<OpFoldResult> offsets {pieceOffset, rewriter.getIndexAttr(0)};
-  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(crossbarSize.getValue())};
-  return tensor::ExtractSliceOp::create(
-    rewriter, loc, pieceType, partialPiecesArg, offsets, sizes, getUnitStrides(rewriter, 2));
+  SmallVector<OpFoldResult> offsets {pieceOffset, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
+  SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
+  auto selectedType = RankedTensorType::get({numOutRows, 1, static_cast<int64_t>(crossbarSize.getValue())}, pieceType.getElementType());
+  Value selected = tensor::ExtractSliceOp::create(rewriter, loc, selectedType, partialPiecesArg, offsets, sizes, getUnitStrides(rewriter, 3));
+  return tensor::CollapseShapeOp::create(rewriter, loc, pieceType, selected, SmallVector<ReassociationIndices> {{0, 1}, {2}});
 }
 
 static Value reduceBatchedPartialPiecesForHSlice(Value partialPiecesArg,
@@ -917,9 +911,7 @@ struct MatMulToGemm : OpRewritePattern<ONNXMatMulOp> {
     if (failed(shapeInfo) || shapeInfo->lhsWasVector || shapeInfo->rhsWasVector)
       return failure();
 
-    const bool hasNonSingletonOutputBatch =
-      !shapeInfo->outputBatchShape.empty() && getStaticShapeElementCount(shapeInfo->outputBatchShape) != 1;
-    if (hasNonSingletonOutputBatch)
+    if (!shapeInfo->outputBatchShape.empty())
       return failure();
 
     Location loc = matmulOp.getLoc();
@@ -1021,8 +1013,8 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
       if (succeeded(paddedRhs)) {
         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());
+        auto partialPiecesType = spatial::getGraphBatchPhysicalResultType(
+          laneCount, RankedTensorType::get({1, static_cast<int64_t>(crossbarSize.getValue())}, shapeInfo->outType.getElementType()));
         auto batchOp = createBatchedVmmBatch(paddedLhs,
                                              *paddedRhs,
                                              paddedLhsType,
@@ -1063,7 +1055,8 @@ struct MatMulBatchedToSpatialComputes : OpRewritePattern<ONNXMatMulOp> {
       }
     }
     const int64_t laneCount = plan.batch * plan.m * plan.n;
-    auto scalarPiecesType = RankedTensorType::get({laneCount, 1}, shapeInfo->outType.getElementType());
+    auto scalarPiecesType = spatial::getGraphBatchPhysicalResultType(
+      laneCount, RankedTensorType::get({1, 1}, shapeInfo->outType.getElementType()));
     auto batchOp = createBatchedVvdmulBatch(plan.lhs,
                                             plan.lhsBatchShape,
                                             plan.rhs,

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

@@ -5,7 +5,6 @@
 #include "llvm/ADT/SmallVector.h"
 
 #include <algorithm>
-#include <numeric>
 #include <optional>
 #include <type_traits>
 
@@ -122,14 +121,6 @@ static RankedTensorType getKeepdimsType(RankedTensorType inputType, Type element
   return RankedTensorType::get(shape, elementType, inputType.getEncoding());
 }
 
-static RankedTensorType getCompactKeptType(RankedTensorType inputType, Type elementType, ArrayRef<bool> reducedAxes) {
-  SmallVector<int64_t> shape;
-  for (auto [dim, isReduced] : llvm::zip_equal(inputType.getShape(), reducedAxes))
-    if (!isReduced)
-      shape.push_back(dim);
-  return RankedTensorType::get(shape, elementType, inputType.getEncoding());
-}
-
 static RankedTensorType getReducedSliceType(RankedTensorType inputType, ArrayRef<bool> reducedAxes) {
   SmallVector<int64_t> shape;
   shape.reserve(inputType.getRank());
@@ -139,9 +130,7 @@ static RankedTensorType getReducedSliceType(RankedTensorType inputType, ArrayRef
 }
 
 static RankedTensorType getLanePackedKeepdimsType(int64_t laneCount, RankedTensorType leafType) {
-  SmallVector<int64_t> shape(leafType.getShape().begin(), leafType.getShape().end());
-  shape.front() = laneCount;
-  return RankedTensorType::get(shape, leafType.getElementType(), leafType.getEncoding());
+  return spatial::getGraphBatchPhysicalResultType(laneCount, leafType);
 }
 
 static SmallVector<int64_t> getKeptAxes(ArrayRef<bool> reducedAxes) {
@@ -191,12 +180,9 @@ static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
 
   SmallVector<OpFoldResult> sliceOffsets;
   SmallVector<OpFoldResult> sliceSizes;
-  SmallVector<OpFoldResult> insertOffsets;
-  SmallVector<OpFoldResult> insertSizes(inputType.getRank(), rewriter.getIndexAttr(1));
   SmallVector<OpFoldResult> unitStrides = getUnitStrides(rewriter, inputType.getRank());
   sliceOffsets.reserve(inputType.getRank());
   sliceSizes.reserve(inputType.getRank());
-  insertOffsets.reserve(inputType.getRank());
 
   auto batchOp =
     createSpatComputeBatch(rewriter,
@@ -209,7 +195,6 @@ static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
                              size_t keptAxisIndex = 0;
                              sliceOffsets.clear();
                              sliceSizes.clear();
-                             insertOffsets.clear();
                              for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
                                if (isReduced) {
                                  sliceOffsets.push_back(rewriter.getIndexAttr(0));
@@ -224,72 +209,90 @@ static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
                                sliceSizes.push_back(rewriter.getIndexAttr(1));
                              }
 
-                             insertOffsets.push_back(args.lane);
-                             insertOffsets.append(inputType.getRank() - 1, rewriter.getIndexAttr(0));
-
                              Value slice = tensor::ExtractSliceOp::create(
                                rewriter, loc, sliceType, args.inputs.front(), sliceOffsets, sliceSizes, unitStrides);
                              Value reduced = spatial::SpatVAvgOp::create(rewriter, loc, leafType, slice).getResult();
-                             createParallelInsertSliceIntoBatchOutput(
-                               rewriter, loc, reduced, args.outputs.front(), insertOffsets, insertSizes, unitStrides);
+                             publishGraphBatchPhysicalFragment(rewriter, loc, reduced, args.outputs.front(), args.lane);
                            });
   if (failed(batchOp))
     return failure();
   return (*batchOp).getResult(0);
 }
 
-static Value buildKeepdimsFromLanePackedBatch(Value batchValue,
-                                              RankedTensorType keepdimsType,
-                                              RankedTensorType compactKeptType,
-                                              ArrayRef<bool> reducedAxes,
-                                              ConversionPatternRewriter& rewriter,
+static FailureOr<Value> buildReduceMeanKeepdimsBlueprint(
+    Value batchValue, RankedTensorType keepdimsType,
+    ArrayRef<bool> reducedAxes, ConversionPatternRewriter& rewriter,
     Location loc) {
-  auto batchType = cast<RankedTensorType>(batchValue.getType());
-  if (batchType == keepdimsType)
-    return batchValue;
+  auto batchType = dyn_cast<RankedTensorType>(batchValue.getType());
+  int64_t rank = keepdimsType.getRank();
+  if (!batchType || !batchType.hasStaticShape()
+      || !keepdimsType.hasStaticShape()
+      || static_cast<int64_t>(reducedAxes.size()) != rank
+      || batchType.getRank() != rank + 1
+      || batchType.getElementType() != keepdimsType.getElementType())
+    return failure();
 
-  SmallVector<ReassociationIndices> collapseToFlat {{}};
-  for (int64_t axis = 0; axis < batchType.getRank(); ++axis)
-    collapseToFlat.front().push_back(axis);
+  int64_t laneCount = 1;
+  SmallVector<int64_t> keptAxes;
+  SmallVector<int64_t> keptAxisStrides;
+  for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
+    int64_t dim = keepdimsType.getDimSize(axis);
+    if (dim <= 0 || (isReduced && dim != 1))
+      return failure();
+    if (!isReduced)
+      keptAxes.push_back(axis);
+  }
+  keptAxisStrides.resize(keptAxes.size(), 1);
+  for (int64_t index = static_cast<int64_t>(keptAxes.size()) - 1;
+       index >= 0; --index) {
+    keptAxisStrides[index] = laneCount;
+    int64_t dim = keepdimsType.getDimSize(keptAxes[index]);
+    if (laneCount > std::numeric_limits<int64_t>::max() / dim)
+      return failure();
+    laneCount *= dim;
+  }
+  if (batchType.getDimSize(0) != laneCount
+      || llvm::any_of(batchType.getShape().drop_front(),
+                      [](int64_t dim) { return dim != 1; }))
+    return failure();
 
-  SmallVector<ReassociationIndices> expandFlatToCompact(1);
-  for (int64_t axis = 0; axis < compactKeptType.getRank(); ++axis)
-    expandFlatToCompact.front().push_back(axis);
-
-  SmallVector<ReassociationIndices> expandCompactToKeepdims;
-  ReassociationIndices pendingLeadingReducedAxes;
+  SmallVector<int64_t> operandIndices(laneCount, 0);
+  SmallVector<int64_t> sourceSlots;
+  SmallVector<int64_t> sourceOffsets(laneCount, 0);
+  SmallVector<int64_t> fragmentOffsets;
+  sourceSlots.reserve(laneCount);
+  fragmentOffsets.reserve(laneCount * rank);
+  for (int64_t lane = 0; lane < laneCount; ++lane) {
+    sourceSlots.push_back(lane);
+    size_t keptAxisIndex = 0;
     for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
       if (isReduced) {
-      if (expandCompactToKeepdims.empty())
-        pendingLeadingReducedAxes.push_back(axis);
-      else
-        expandCompactToKeepdims.back().push_back(axis);
+        fragmentOffsets.push_back(0);
         continue;
       }
-
-    expandCompactToKeepdims.emplace_back();
-    auto& group = expandCompactToKeepdims.back();
-    group.append(pendingLeadingReducedAxes.begin(), pendingLeadingReducedAxes.end());
-    pendingLeadingReducedAxes.clear();
-    group.push_back(axis);
+      int64_t dim = keepdimsType.getDimSize(axis);
+      fragmentOffsets.push_back(
+        (lane / keptAxisStrides[keptAxisIndex]) % dim);
+      ++keptAxisIndex;
     }
-  if (!pendingLeadingReducedAxes.empty())
-    expandCompactToKeepdims.back().append(pendingLeadingReducedAxes.begin(), pendingLeadingReducedAxes.end());
-
-  auto reshapeCompute =
-    createSpatCompute<1>(rewriter, loc, TypeRange {keepdimsType}, {}, ValueRange {batchValue}, [&](Value input) {
-      auto flatType =
-        RankedTensorType::get({batchType.getDimSize(0)}, batchType.getElementType(), batchType.getEncoding());
-      Value flat = tensor::CollapseShapeOp::create(rewriter, loc, flatType, input, collapseToFlat);
-      Value compact = flat;
-      if (compactKeptType != flatType)
-        compact = tensor::ExpandShapeOp::create(rewriter, loc, compactKeptType, flat, expandFlatToCompact);
-      Value keepdims = compact;
-      if (keepdimsType != compactKeptType)
-        keepdims = tensor::ExpandShapeOp::create(rewriter, loc, keepdimsType, compact, expandCompactToKeepdims);
-      spatial::SpatYieldOp::create(rewriter, loc, keepdims);
-    });
-  return reshapeCompute.getResult(0);
+  }
+  SmallVector<int64_t> fragmentSizes(fragmentOffsets.size(), 1);
+  SmallVector<int64_t> fragmentStrides(fragmentOffsets.size(), 1);
+  return spatial::SpatBlueprintOp::create(
+           rewriter, loc, keepdimsType, batchValue, ValueRange {},
+           rewriter.getStringAttr("nchw"),
+           rewriter.getStringAttr("fragmented"),
+           rewriter.getDenseI64ArrayAttr(fragmentOffsets),
+           rewriter.getDenseI64ArrayAttr(fragmentSizes),
+           rewriter.getStringAttr("reduce_mean_keepdims_fragments"),
+           rewriter.getStringAttr("fragment_assembly"),
+           rewriter.getDenseI64ArrayAttr(operandIndices),
+           rewriter.getDenseI64ArrayAttr(sourceSlots),
+           rewriter.getDenseI64ArrayAttr(sourceOffsets),
+           rewriter.getDenseI64ArrayAttr(fragmentStrides),
+           rewriter.getStringAttr("disjoint"),
+           rewriter.getStringAttr("complete"))
+    .getOutput();
 }
 
 static SmallVector<ReassociationIndices> buildCollapseReassociation(ArrayRef<bool> reducedAxes) {
@@ -366,26 +369,36 @@ struct ReduceMeanToSpatialCompute : OpConversionPattern<ReduceMeanOp> {
 
     Location loc = reduceMeanOp.getLoc();
     RankedTensorType leafType = getAllOnesType(inputType, resultType.getElementType());
-    RankedTensorType compactKeptType = getCompactKeptType(inputType, resultType.getElementType(), reducedAxes);
     RankedTensorType keepdimsType = getKeepdimsType(inputType, resultType.getElementType(), reducedAxes);
     int64_t laneCount = 1;
-    for (int64_t dim : compactKeptType.getShape())
+    for (auto [dim, isReduced] : llvm::zip_equal(keepdimsType.getShape(), reducedAxes)) {
+      if (isReduced)
+        continue;
+      if (dim <= 0 || laneCount > std::numeric_limits<int32_t>::max() / dim)
+        return rewriter.notifyMatchFailure(
+          reduceMeanOp, "ReduceMean physical lane count is not representable");
       laneCount *= dim;
+    }
     RankedTensorType batchType = getLanePackedKeepdimsType(laneCount, leafType);
 
     auto lanePackedKeepdims =
       buildReduceMeanKeepdimsBatch(adaptor.getData(), reducedAxes, batchType, leafType, rewriter, loc);
     if (failed(lanePackedKeepdims))
       return failure();
-    Value reducedKeepdims =
-      buildKeepdimsFromLanePackedBatch(*lanePackedKeepdims, keepdimsType, compactKeptType, reducedAxes, rewriter, loc);
+    auto reducedKeepdims = buildReduceMeanKeepdimsBlueprint(
+      *lanePackedKeepdims, keepdimsType, reducedAxes, rewriter, loc);
+    if (failed(reducedKeepdims))
+      return rewriter.notifyMatchFailure(
+        reduceMeanOp,
+        "cannot build physical-fragment ReduceMean keepdims reconstruction");
 
     if (semantics->keepdims != 0) {
-      rewriter.replaceOp(reduceMeanOp, reducedKeepdims);
+      rewriter.replaceOp(reduceMeanOp, *reducedKeepdims);
       return success();
     }
 
-    Value reduced = squeezeReducedAxes(reducedKeepdims, resultType, reducedAxes, rewriter, loc);
+    Value reduced = squeezeReducedAxes(
+      *reducedKeepdims, resultType, reducedAxes, rewriter, loc);
     rewriter.replaceOp(reduceMeanOp, reduced);
     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"
@@ -128,8 +129,6 @@ struct PromoteWeightLikeComputeInputsPattern : OpRewritePattern<spatial::SpatGra
     Block& oldBlock = compute.getBody().front();
 
     rewriter.setInsertionPointAfter(compute);
-    auto newCompute = spatial::SpatGraphCompute::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::SpatGra
     }
     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());
@@ -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::SpatGraphComputeBatch::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();
   }
 };

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/Patterns/Tensor/Transpose.cpp

@@ -5,7 +5,7 @@
 
 #include "llvm/ADT/SmallVector.h"
 
-#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Dialect/ONNX/ONNXOps.hpp"
@@ -52,35 +52,12 @@ static FailureOr<Value> materializeTransposedConstant(Value input,
     return failure();
   }
 
-  if (denseAttr.isSplat())
+  auto transposedAttr = transposeDenseElementsAttr(denseAttr, permutation);
+  if (failed(transposedAttr) || transposedAttr->getType() != resultType)
+    return failure();
   return getOrCreateConstant(rewriter,
                              rewriter.getInsertionBlock()->getParentOp(),
-                               DenseElementsAttr::get(resultType, denseAttr.getSplatValue<Attribute>()),
-                               resultType);
-
-  SmallVector<Attribute> inputValues(denseAttr.getValues<Attribute>());
-  SmallVector<Attribute> resultValues(inputValues.size());
-  SmallVector<int64_t> inputStrides = computeRowMajorStrides(inputType.getShape());
-  SmallVector<int64_t> resultStrides = computeRowMajorStrides(resultType.getShape());
-  SmallVector<int64_t> inputIndices(inputType.getRank(), 0);
-
-  for (auto [linearIndex, value] : llvm::enumerate(inputValues)) {
-    int64_t remaining = static_cast<int64_t>(linearIndex);
-    for (int64_t dim = 0; dim < inputType.getRank(); ++dim) {
-      inputIndices[dim] = inputStrides.empty() ? 0 : remaining / inputStrides[dim];
-      remaining = inputStrides.empty() ? 0 : remaining % inputStrides[dim];
-    }
-
-    int64_t resultLinearIndex = 0;
-    for (int64_t dim = 0; dim < resultType.getRank(); ++dim)
-      resultLinearIndex += inputIndices[permutation[dim]] * resultStrides[dim];
-
-    resultValues[resultLinearIndex] = value;
-  }
-
-  return getOrCreateConstant(rewriter,
-                             rewriter.getInsertionBlock()->getParentOp(),
-                             DenseElementsAttr::get(resultType, resultValues),
+                             *transposedAttr,
                              resultType);
 }
 

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -6,10 +6,11 @@
 
 #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"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
 
 using namespace mlir;
 
@@ -19,7 +20,6 @@ namespace {
 static constexpr StringLiteral kLogicalLayout = "nchw";
 static constexpr StringLiteral kDenseLayout = "dense_nchw";
 static constexpr StringLiteral kRowStripLayout = "nchw_row_strip";
-static constexpr StringLiteral kRowStripIndexMap = "packed_hwc_rows_to_nchw";
 
 enum class SelectedLayout {
   DenseNchw,
@@ -34,6 +34,8 @@ static SelectedLayout getSelectedLayout(llvm::DenseMap<Value, SelectedLayout>& l
 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;
@@ -49,21 +51,26 @@ static bool allUsersCanHandleRowStrip(Value value, llvm::DenseMap<Value, Selecte
   return true;
 }
 
-static 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});
+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);
   }
-  return {offsets, sizes};
+  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());
@@ -76,6 +83,9 @@ 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;
@@ -85,11 +95,27 @@ 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);
@@ -108,6 +134,7 @@ static spatial::SpatBlueprintOp insertRowStripBlueprint(IRRewriter& rewriter, Va
                                               nullptr,
                                               nullptr,
                                               nullptr,
+                                              nullptr,
                                               nullptr);
 }
 
@@ -173,6 +200,14 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
           }
           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;
+        }
       }
     }
 
@@ -180,6 +215,8 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
       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

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

@@ -141,7 +141,8 @@ collectTopLevelFragmentAssemblyCopies(OpResult result, RankedTensorType packedRe
     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)
+    std::optional<ArrayRef<int64_t>> sourceSlotsAttr = blueprint.getFragmentSourceSlots();
+    if (!mode || *mode != "fragment_assembly" || !operandIndicesAttr || !sourceOffsetsAttr || !sourceSlotsAttr)
       return failure();
     if (!blueprint.getOutput().hasOneUse() || !isa<func::ReturnOp>(*blueprint.getOutput().getUsers().begin()))
       return failure();
@@ -153,6 +154,9 @@ collectTopLevelFragmentAssemblyCopies(OpResult result, RankedTensorType packedRe
 
     ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
     ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
+    ArrayRef<int64_t> sourceSlots = *sourceSlotsAttr;
+    if (sourceSlots.size() != operandIndices.size())
+      return failure();
     ArrayRef<int64_t> flatOffsets = blueprint.getFragmentOffsets();
     ArrayRef<int64_t> flatSizes = blueprint.getFragmentSizes();
     ArrayRef<int64_t> flatStrides = *stridesAttr;
@@ -174,7 +178,8 @@ collectTopLevelFragmentAssemblyCopies(OpResult result, RankedTensorType packedRe
       if (operandIndices[fragmentIndex] != static_cast<int64_t>(use.getOperandNumber()))
         continue;
 
-      int64_t sourceElementOffset = sourceOffsets[fragmentIndex];
+      int64_t sourceElementOffset =
+        sourceSlots[fragmentIndex] * payloadElementCount + sourceOffsets[fragmentIndex];
       int64_t lane = sourceElementOffset / payloadElementCount;
       if (lane < 0 || lane >= static_cast<int64_t>(laneCount))
         return failure();
@@ -395,6 +400,11 @@ LogicalResult raptor::SpatialToPimPass::lowerComputeBatchOp(spatial::SpatSchedul
     if (isa<spatial::SpatYieldOp>(op))
       continue;
 
+    // Cloning a region-bearing operation may leave the rewriter inside that
+    // region. Every old-block operation is lowered at the core-batch body
+    // boundary.
+    rewriter.setInsertionPointToEnd(newBlock);
+
     if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
       std::optional<StringRef> modeAttr = blueprint.getMode();
       if (modeAttr && *modeAttr == "fragment_assembly") {

src/PIM/Conversion/SpatialToPim/Common.cpp

@@ -8,6 +8,8 @@
 #include <limits>
 
 #include "Common.hpp"
+#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/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
@@ -192,21 +194,23 @@ forEachContiguousDestinationChunk(ArrayRef<int64_t> destShape,
 }
 
 static mlir::Value
-createSteppedOffset(OpBuilder& builder, Location loc, mlir::Value start, mlir::Value index, int64_t stepBytes) {
+createSteppedOffset(OpBuilder& builder, Location loc, mlir::Value start, mlir::Value index,
+                    int64_t stepBytes, Operation *constantAnchor) {
   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();
+  return createOrFoldAffineApply(
+    builder, loc, builder.getAffineDimExpr(0) + builder.getAffineDimExpr(1) * stepBytes,
+    ValueRange {start, index}, constantAnchor);
 }
 
 static mlir::Value createIndexedOffset(OpBuilder& builder,
                                        Location loc,
                                        mlir::Value indexArg,
-                                       ArrayRef<int64_t> values) {
+                                       ArrayRef<int64_t> values,
+                                       Operation *constantAnchor) {
   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());
+    return getOrCreateIndexConstant(builder, constantAnchor, values.front());
 
   if (values.size() >= 2) {
     int64_t step = values[1] - values[0];
@@ -214,21 +218,18 @@ static mlir::Value createIndexedOffset(OpBuilder& builder,
       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();
+      return createOrFoldAffineApply(
+        builder, loc, builder.getAffineDimExpr(0) * step + values.front(),
+        ValueRange {indexArg}, constantAnchor);
     }
   }
 
-  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;
+  RankedTensorType tableType = RankedTensorType::get(
+    {static_cast<int64_t>(values.size())}, builder.getI64Type());
+  DenseElementsAttr tableAttr = DenseElementsAttr::get(tableType, values);
+  mlir::Value table = getOrCreateConstant(builder, constantAnchor, tableAttr, tableType);
+  mlir::Value selected = tensor::ExtractOp::create(builder, loc, table, ValueRange {indexArg});
+  return arith::IndexCastOp::create(builder, loc, builder.getIndexType(), selected).getResult();
 }
 
 struct FragmentAssemblyCopyRunFamily {
@@ -433,11 +434,11 @@ static FailureOr<mlir::Value> emitFragmentAssemblyCopyRun(OpBuilder& builder,
   mlir::Value hostStart;
   mlir::Value sourceStart;
   if (laneArg) {
-    hostStart = createIndexedOffset(builder, loc, *laneArg, run.hostStartBytesByLane);
-    sourceStart = createIndexedOffset(builder, loc, *laneArg, run.sourceStartBytesByLane);
+    hostStart = createIndexedOffset(builder, loc, *laneArg, run.hostStartBytesByLane, anchor);
+    sourceStart = createIndexedOffset(builder, loc, *laneArg, run.sourceStartBytesByLane, anchor);
   } else {
-    hostStart = arith::ConstantIndexOp::create(builder, loc, run.hostStartBytesByLane.front());
-    sourceStart = arith::ConstantIndexOp::create(builder, loc, run.sourceStartBytesByLane.front());
+    hostStart = getOrCreateIndexConstant(builder, anchor, run.hostStartBytesByLane.front());
+    sourceStart = getOrCreateIndexConstant(builder, anchor, run.sourceStartBytesByLane.front());
   }
 
   if (hostRunStartDelta)
@@ -459,9 +460,9 @@ static FailureOr<mlir::Value> emitFragmentAssemblyCopyRun(OpBuilder& builder,
       .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);
+  mlir::Value lowerBound = getOrCreateIndexConstant(builder, anchor, 0);
+  mlir::Value upperBound = getOrCreateIndexConstant(builder, anchor, run.count);
+  mlir::Value step = getOrCreateIndexConstant(builder, anchor, 1);
   FailureOr<NormalizedLoopResult> loop = buildNormalizedScfFor(
     builder,
     loc,
@@ -474,9 +475,10 @@ static FailureOr<mlir::Value> emitFragmentAssemblyCopyRun(OpBuilder& builder,
         mlir::Value flatIndex,
         ValueRange iterArgs,
         SmallVectorImpl<mlir::Value>& yielded) {
-      mlir::Value hostOffset = createSteppedOffset(loopBuilder, bodyLoc, hostStart, flatIndex, run.hostStepBytes);
+      mlir::Value hostOffset = createSteppedOffset(
+        loopBuilder, bodyLoc, hostStart, flatIndex, run.hostStepBytes, anchor);
       mlir::Value sourceOffset =
-        createSteppedOffset(loopBuilder, bodyLoc, sourceStart, flatIndex, run.sourceStepBytes);
+        createSteppedOffset(loopBuilder, bodyLoc, sourceStart, flatIndex, run.sourceStepBytes, anchor);
       mlir::Value copied =
         pim::PimMemCopyDevToHostOp::create(loopBuilder,
                                            bodyLoc,
@@ -506,9 +508,9 @@ static FailureOr<mlir::Value> emitFragmentAssemblyCopyRunFamily(OpBuilder& build
     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);
+  mlir::Value lowerBound = getOrCreateIndexConstant(builder, anchor, 0);
+  mlir::Value upperBound = getOrCreateIndexConstant(builder, anchor, family.sourceRunStartDeltas.size());
+  mlir::Value step = getOrCreateIndexConstant(builder, anchor, 1);
   FailureOr<NormalizedLoopResult> outerLoop = buildNormalizedScfFor(
     builder,
     loc,
@@ -522,9 +524,9 @@ static FailureOr<mlir::Value> emitFragmentAssemblyCopyRunFamily(OpBuilder& build
         ValueRange iterArgs,
         SmallVectorImpl<mlir::Value>& yielded) {
       mlir::Value sourceRunStartDelta =
-        createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.sourceRunStartDeltas);
+        createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.sourceRunStartDeltas, anchor);
       mlir::Value hostRunStartDelta =
-        createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.hostRunStartDeltas);
+        createIndexedOffset(loopBuilder, bodyLoc, runIndex, family.hostRunStartDeltas, anchor);
       FailureOr<mlir::Value> copied = emitFragmentAssemblyCopyRun(loopBuilder,
                                                                   bodyLoc,
                                                                   family.prototype,

src/PIM/Conversion/SpatialToPim/CoreLoweringPatterns.cpp

@@ -10,7 +10,9 @@
 #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/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
@@ -180,16 +182,79 @@ static LogicalResult collectHelperComputeChain(spatial::SpatScheduledCompute com
   return success();
 }
 
+static bool isHostMaterializableHelperOp(Operation* op) {
+  if (isa<spatial::SpatYieldOp>(op))
+    return true;
+  if (isa<arith::ConstantOp>(op) || op->hasTrait<OpTrait::ConstantLike>())
+    return true;
+  if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
+    std::optional<StringRef> mode = blueprint.getMode();
+    return mode && *mode == "fragment_assembly";
+  }
+  return isShapingOnlyOp(op) || isPureIndexComputationOp(op);
+}
+
+static FailureOr<DenseMap<Value, Attribute>>
+analyzeHostMaterializableHelper(spatial::SpatScheduledCompute computeOp) {
+  DenseMap<Value, Attribute> folded;
+  for (auto [weightIndex, weight] : llvm::enumerate(computeOp.getWeights())) {
+    auto argument = computeOp.getWeightArgument(weightIndex);
+    if (!argument)
+      return failure();
+    Attribute constant;
+    if (matchPattern(weight, m_Constant(&constant)))
+      folded[*argument] = constant;
+  }
+  Block& block = computeOp.getBody().front();
+  for (Operation& op : block) {
+    if (!isHostMaterializableHelperOp(&op))
+      return failure();
+    if (isa<spatial::SpatYieldOp, spatial::SpatBlueprintOp>(op)
+        || (isShapingOnlyOp(&op) && !isPureIndexComputationOp(&op)))
+      continue;
+    if (isa<arith::ConstantOp>(op) || op.hasTrait<OpTrait::ConstantLike>()) {
+      for (Value result : op.getResults()) {
+        Attribute constant;
+        if (!matchPattern(result, m_Constant(&constant)))
+          return failure();
+        folded[result] = constant;
+      }
+      continue;
+    }
+    if (!isPureIndexComputationOp(&op) || op.getNumRegions() != 0)
+      return failure();
+    SmallVector<Attribute> operands;
+    for (Value operand : op.getOperands()) {
+      auto it = folded.find(operand);
+      if (it == folded.end())
+        return failure();
+      operands.push_back(it->second);
+    }
+    SmallVector<OpFoldResult> results;
+    if (failed(op.fold(operands, results))
+        || results.size() != op.getNumResults())
+      return failure();
+    for (auto [result, foldResult] : llvm::zip(op.getResults(), results)) {
+      auto attribute = dyn_cast<Attribute>(foldResult);
+      if (!attribute)
+        return failure();
+      folded[result] = attribute;
+    }
+  }
+  return folded;
+}
+
 static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatScheduledCompute computeOp,
                                                            IRRewriter& rewriter,
                                                            OperationFolder& constantFolder) {
   if (!computeOp.getInputs().empty() || computeOp.getNumResults() != 1)
     return false;
+  if (computeOp.getResult(0).use_empty())
+    return false;
   if (!llvm::all_of(computeOp.getResult(0).getUsers(), [](Operation* user) {
         return isa<spatial::SpatScheduledCompute, spatial::SpatScheduledComputeBatch, pim::PimCoreOp, pim::PimCoreBatchOp>(user);
       }))
     return false;
-
   Block& block = computeOp.getBody().front();
   if (block.getNumArguments() != computeOp.getWeights().size())
     return false;
@@ -197,6 +262,9 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatSchedule
   auto yieldOp = dyn_cast<spatial::SpatYieldOp>(block.getTerminator());
   if (!yieldOp || yieldOp.getNumOperands() != 1)
     return false;
+  auto folded = analyzeHostMaterializableHelper(computeOp);
+  if (failed(folded))
+    return false;
 
   rewriter.setInsertionPoint(computeOp);
   IRMapping mapping;
@@ -218,6 +286,20 @@ static bool inlineInputlessHelperComputeForWeightLikeUsers(spatial::SpatSchedule
       }
     }
 
+    if (isa<arith::ConstantOp>(op) || op.hasTrait<OpTrait::ConstantLike>()
+        || isPureIndexComputationOp(&op)) {
+      for (Value result : op.getResults()) {
+        auto it = folded->find(result);
+        if (it == folded->end())
+          return false;
+        mapping.map(
+          result,
+          getOrCreateConstant(constantFolder, computeOp, it->second,
+                              result.getType()));
+      }
+      continue;
+    }
+
     cloneMappedHelperOperands(&op, mapping, rewriter, constantFolder);
     Operation* clonedOp = rewriter.clone(op, mapping);
     for (auto [originalResult, newResult] : llvm::zip(op.getResults(), clonedOp->getResults()))

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

@@ -1,10 +1,23 @@
 #include "Dialect/Pim/Transforms/Bufferization/Common.hpp"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
 using namespace mlir;
 
+static SmallVector<Region *> getSelectionRegions(OpResult result) {
+  SmallVector<Region *> regions;
+  if (auto selection = dyn_cast<scf::IndexSwitchOp>(result.getOwner()))
+    for (Region &region : selection->getRegions())
+      regions.push_back(&region);
+  else if (auto selection = dyn_cast<scf::IfOp>(result.getOwner())) {
+    regions.push_back(&selection.getThenRegion());
+    regions.push_back(&selection.getElseRegion());
+  }
+  return regions;
+}
+
 static bool isCoreBatchInputArgument(Value value) {
   auto blockArg = dyn_cast<BlockArgument>(value);
   if (!blockArg)
@@ -92,20 +105,46 @@ FailureOr<Value> onnx_mlir::pim::getPimAddressBase(Value value, const StaticValu
 }
 
 bool onnx_mlir::pim::isHostBackedPimAddress(Value value, const StaticValueKnowledge& knowledge) {
-  auto base = getPimStorageBase(value, knowledge);
-  if (failed(base))
+  llvm::SmallPtrSet<Value, 8> visited;
+  std::function<bool(Value)> isHost = [&](Value current) {
+    auto base = getPimStorageBase(current, knowledge);
+    if (failed(base) || !visited.insert(*base).second)
       return false;
-
-  if (isCoreBatchInputArgument(*base))
-    return true;
-
-  return isa_and_nonnull<memref::GetGlobalOp>(base->getDefiningOp());
+    bool resultIsHost = isCoreBatchInputArgument(*base)
+                        || isa_and_nonnull<memref::GetGlobalOp>(base->getDefiningOp());
+    auto result = dyn_cast<OpResult>(*base);
+    SmallVector<Region *> regions = result ? getSelectionRegions(result)
+                                          : SmallVector<Region *>();
+    if (!resultIsHost && !regions.empty())
+      resultIsHost = llvm::all_of(regions, [&](Region *region) {
+        auto yield = dyn_cast<scf::YieldOp>(region->front().getTerminator());
+        return yield && result.getResultNumber() < yield.getNumOperands()
+               && isHost(yield.getOperand(result.getResultNumber()));
+      });
+    visited.erase(*base);
+    return resultIsHost;
+  };
+  return isHost(value);
 }
 
 bool onnx_mlir::pim::isDeviceLocalPimAddress(Value value, const StaticValueKnowledge& knowledge) {
-  auto base = getPimStorageBase(value, knowledge);
-  if (failed(base))
+  llvm::SmallPtrSet<Value, 8> visited;
+  std::function<bool(Value)> isDevice = [&](Value current) {
+    auto base = getPimStorageBase(current, knowledge);
+    if (failed(base) || !visited.insert(*base).second)
       return false;
-
-  return isa_and_nonnull<memref::AllocOp>(base->getDefiningOp());
+    bool resultIsDevice = isa_and_nonnull<memref::AllocOp>(base->getDefiningOp());
+    auto result = dyn_cast<OpResult>(*base);
+    SmallVector<Region *> regions = result ? getSelectionRegions(result)
+                                          : SmallVector<Region *>();
+    if (!resultIsDevice && !regions.empty())
+      resultIsDevice = llvm::all_of(regions, [&](Region *region) {
+        auto yield = dyn_cast<scf::YieldOp>(region->front().getTerminator());
+        return yield && result.getResultNumber() < yield.getNumOperands()
+               && isDevice(yield.getOperand(result.getResultNumber()));
+      });
+    visited.erase(*base);
+    return resultIsDevice;
+  };
+  return isDevice(value);
 }

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

@@ -2,6 +2,8 @@
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 
+#include "llvm/Support/MathExtras.h"
+
 #include "ContiguityPatterns.hpp"
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
@@ -33,6 +35,7 @@ struct CopyEndpointPlan {
 
 struct CopyLoopPlan {
   SmallVector<int64_t> outerShape;
+  int64_t outerElements = 0;
   int64_t chunkBytes = 0;
   ByteOffsetExpr targetBaseOffset;
   ByteOffsetExpr sourceBaseOffset;
@@ -74,6 +77,24 @@ static void appendTerm(ByteOffsetExpr& expr, Value value, int64_t scale) {
     expr.terms.push_back(ByteOffsetTerm {value, scale});
 }
 
+static FailureOr<int64_t> checkedPositiveMul(int64_t lhs, int64_t rhs) {
+  int64_t result = 0;
+  if (lhs < 0 || rhs < 0 || llvm::MulOverflow(lhs, rhs, result))
+    return failure();
+  return result;
+}
+
+static FailureOr<int64_t> checkedPositiveProduct(ArrayRef<int64_t> values) {
+  int64_t result = 1;
+  for (int64_t value : values) {
+    auto product = checkedPositiveMul(result, value);
+    if (failed(product))
+      return failure();
+    result = *product;
+  }
+  return result;
+}
+
 static FailureOr<SmallVector<int64_t>> getStaticMemRefStrides(MemRefType type) {
   SmallVector<int64_t> strides;
   int64_t offset = 0;
@@ -84,6 +105,165 @@ static FailureOr<SmallVector<int64_t>> getStaticMemRefStrides(MemRefType type) {
   return strides;
 }
 
+static FailureOr<SmallVector<int64_t>> getProvenMemRefStrides(Value value) {
+  llvm::SmallPtrSet<Value, 8> visiting;
+  std::function<FailureOr<SmallVector<int64_t>>(Value)> prove =
+    [&](Value current) -> FailureOr<SmallVector<int64_t>> {
+      auto type = dyn_cast<MemRefType>(current.getType());
+      if (!type || !visiting.insert(current).second)
+        return failure();
+      if (auto strides = getStaticMemRefStrides(type); succeeded(strides)) {
+        visiting.erase(current);
+        return strides;
+      }
+      if (auto castOp = current.getDefiningOp<memref::CastOp>()) {
+        auto strides = prove(castOp.getSource());
+        visiting.erase(current);
+        return strides;
+      }
+      if (auto subview = current.getDefiningOp<memref::SubViewOp>()) {
+        auto sourceStrides = prove(subview.getSource());
+        if (failed(sourceStrides) || subview.getSourceType().getRank() != subview.getType().getRank()) {
+          visiting.erase(current);
+          return failure();
+        }
+        SmallVector<int64_t> strides;
+        for (auto [sourceStride, viewStride] :
+             llvm::zip_equal(*sourceStrides, subview.getStaticStrides())) {
+          if (ShapedType::isDynamic(viewStride) || viewStride < 0) {
+            visiting.erase(current);
+            return failure();
+          }
+          auto stride = checkedPositiveMul(sourceStride, viewStride);
+          if (failed(stride)) {
+            visiting.erase(current);
+            return failure();
+          }
+          strides.push_back(*stride);
+        }
+        visiting.erase(current);
+        return strides;
+      }
+      if (auto expand = current.getDefiningOp<memref::ExpandShapeOp>()) {
+        auto sourceStrides = prove(expand.getSrc());
+        auto resultType = dyn_cast<MemRefType>(expand.getResult().getType());
+        auto sourceType = dyn_cast<MemRefType>(expand.getSrc().getType());
+        if (failed(sourceStrides) || !sourceType || !resultType
+            || !resultType.hasStaticShape()
+            || sourceStrides->size() != static_cast<size_t>(sourceType.getRank())
+            || llvm::any_of(resultType.getShape(), [](int64_t dim) {
+                 return dim <= 0;
+               })) {
+          visiting.erase(current);
+          return failure();
+        }
+        SmallVector<int64_t> strides(resultType.getRank());
+        SmallVector<bool> assigned(resultType.getRank(), false);
+        for (auto [sourceDim, group] :
+             llvm::enumerate(expand.getReassociationIndices())) {
+          if (sourceDim >= sourceStrides->size() || group.empty()) {
+            visiting.erase(current);
+            return failure();
+          }
+          int64_t stride = (*sourceStrides)[sourceDim];
+          for (int64_t resultDim : llvm::reverse(group)) {
+            if (resultDim < 0 || resultDim >= resultType.getRank()
+                || assigned[resultDim]) {
+              visiting.erase(current);
+              return failure();
+            }
+            strides[resultDim] = stride;
+            assigned[resultDim] = true;
+            auto nextStride = checkedPositiveMul(
+              stride, resultType.getDimSize(resultDim));
+            if (failed(nextStride)) {
+              visiting.erase(current);
+              return failure();
+            }
+            stride = *nextStride;
+          }
+        }
+        if (llvm::is_contained(assigned, false)) {
+          visiting.erase(current);
+          return failure();
+        }
+        visiting.erase(current);
+        return strides;
+      }
+      if (auto collapse = current.getDefiningOp<memref::CollapseShapeOp>()) {
+        auto sourceStrides = prove(collapse.getSrc());
+        auto sourceType = dyn_cast<MemRefType>(collapse.getSrc().getType());
+        if (failed(sourceStrides) || !sourceType
+            || !sourceType.hasStaticShape()
+            || sourceStrides->size() != static_cast<size_t>(sourceType.getRank())) {
+          visiting.erase(current);
+          return failure();
+        }
+        SmallVector<int64_t> strides;
+        for (ArrayRef<int64_t> group : collapse.getReassociationIndices()) {
+          if (group.empty()) {
+            visiting.erase(current);
+            return failure();
+          }
+          for (int64_t dim : group)
+            if (dim < 0 || dim >= sourceType.getRank()
+                || sourceType.getDimSize(dim) <= 0
+                || (*sourceStrides)[dim] < 0) {
+              visiting.erase(current);
+              return failure();
+            }
+          for (auto pair : llvm::zip(group.drop_back(), group.drop_front())) {
+            int64_t outer = std::get<0>(pair);
+            int64_t inner = std::get<1>(pair);
+            auto expectedOuterStride = checkedPositiveMul(
+              (*sourceStrides)[inner], sourceType.getDimSize(inner));
+            if (failed(expectedOuterStride)
+                || (*sourceStrides)[outer] != *expectedOuterStride) {
+              visiting.erase(current);
+              return failure();
+            }
+          }
+          strides.push_back((*sourceStrides)[group.back()]);
+        }
+        visiting.erase(current);
+        return strides;
+      }
+      auto result = dyn_cast<OpResult>(current);
+      SmallVector<Region *> regions;
+      if (result) {
+        if (auto selection = dyn_cast<scf::IndexSwitchOp>(result.getOwner()))
+          for (Region &region : selection->getRegions())
+            regions.push_back(&region);
+        else if (auto selection = dyn_cast<scf::IfOp>(result.getOwner())) {
+          regions.push_back(&selection.getThenRegion());
+          regions.push_back(&selection.getElseRegion());
+        }
+      }
+      if (regions.empty()) {
+        visiting.erase(current);
+        return failure();
+      }
+      std::optional<SmallVector<int64_t>> common;
+      for (Region *region : regions) {
+        auto yield = dyn_cast<scf::YieldOp>(region->front().getTerminator());
+        if (!yield || result.getResultNumber() >= yield.getNumOperands()) {
+          visiting.erase(current);
+          return failure();
+        }
+        auto strides = prove(yield.getOperand(result.getResultNumber()));
+        if (failed(strides) || (common && *common != *strides)) {
+          visiting.erase(current);
+          return failure();
+        }
+        common = std::move(*strides);
+      }
+      visiting.erase(current);
+      return common ? FailureOr<SmallVector<int64_t>>(std::move(*common))
+                    : FailureOr<SmallVector<int64_t>>(failure());
+    };
+  return prove(value);
+}
+
 static FailureOr<int64_t> getShapedByteSize(MemRefType type) {
   if (!type.hasStaticShape() || !hasByteSizedElementType(type.getElementType()))
     return failure();
@@ -119,12 +299,15 @@ inferLogicalCopyShape(MemRefType targetType, MemRefType sourceType, int64_t size
   return failure();
 }
 
-static FailureOr<int64_t> getContiguousSuffixRank(MemRefType type, ArrayRef<int64_t> copyShape) {
-  if (!type.hasStaticShape() || !hasByteSizedElementType(type.getElementType())
+static FailureOr<int64_t> getContiguousSuffixRank(Value value, ArrayRef<int64_t> copyShape) {
+  auto type = dyn_cast<MemRefType>(value.getType());
+  if (!type || !type.hasStaticShape() || !hasByteSizedElementType(type.getElementType())
       || type.getRank() != static_cast<int64_t>(copyShape.size()))
     return failure();
+  if (llvm::any_of(copyShape, [](int64_t dim) { return dim <= 0; }))
+    return failure();
 
-  auto strides = getStaticMemRefStrides(type);
+  auto strides = getProvenMemRefStrides(value);
   if (failed(strides))
     return failure();
 
@@ -134,7 +317,10 @@ static FailureOr<int64_t> getContiguousSuffixRank(MemRefType type, ArrayRef<int6
     if ((*strides)[dim] != expectedStride)
       break;
     ++contiguousSuffixRank;
-    expectedStride *= copyShape[dim];
+    auto nextStride = checkedPositiveMul(expectedStride, copyShape[dim]);
+    if (failed(nextStride))
+      return failure();
+    expectedStride = *nextStride;
   }
   return contiguousSuffixRank;
 }
@@ -174,18 +360,25 @@ static FailureOr<CopyEndpointPlan> analyzeCopyEndpoint(Value value, Value initia
     if (!sourceType || !sourceType.hasStaticShape() || !hasByteSizedElementType(sourceType.getElementType()))
       return failure();
 
-    auto sourceStrides = getStaticMemRefStrides(sourceType);
+    auto sourceStrides = getProvenMemRefStrides(subviewOp.getSource());
     if (failed(sourceStrides))
       return failure();
 
     int64_t elementByteWidth = static_cast<int64_t>(getElementTypeSizeInBytes(sourceType.getElementType()));
     for (auto [offset, stride] : llvm::zip_equal(subviewOp.getMixedOffsets(), *sourceStrides)) {
-      int64_t byteScale = stride * elementByteWidth;
+      auto byteScale = checkedPositiveMul(stride, elementByteWidth);
+      if (failed(byteScale))
+        return failure();
       if (auto attr = dyn_cast<Attribute>(offset)) {
-        endpoint.offset.constant += cast<IntegerAttr>(attr).getInt() * byteScale;
+        auto constantOffset = checkedPositiveMul(
+          cast<IntegerAttr>(attr).getInt(), *byteScale);
+        if (failed(constantOffset)
+            || llvm::AddOverflow(endpoint.offset.constant, *constantOffset,
+                                 endpoint.offset.constant))
+          return failure();
         continue;
       }
-      appendTerm(endpoint.offset, cast<Value>(offset), byteScale);
+      appendTerm(endpoint.offset, cast<Value>(offset), *byteScale);
     }
 
     endpoint.base = subviewOp.getSource();
@@ -204,17 +397,34 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO
   if (!targetType || !sourceType || size <= 0)
     return failure();
 
-  auto logicalCopyShape = inferLogicalCopyShape(targetType, sourceType, size);
-  if (failed(logicalCopyShape))
-    return failure();
-
   auto targetPlan = analyzeCopyEndpoint(target, targetOffset, targetType);
   auto sourcePlan = analyzeCopyEndpoint(source, sourceOffset, sourceType);
   if (failed(targetPlan) || failed(sourcePlan))
     return failure();
 
-  auto targetSuffixRank = getContiguousSuffixRank(targetType, *logicalCopyShape);
-  auto sourceSuffixRank = getContiguousSuffixRank(sourceType, *logicalCopyShape);
+  auto targetBytes = getShapedByteSize(targetType);
+  auto sourceBytes = getShapedByteSize(sourceType);
+  if (targetType.getElementType() == sourceType.getElementType() && succeeded(targetBytes) && succeeded(sourceBytes)
+      && *targetBytes == size && *sourceBytes == size) {
+    auto targetSuffixRank = getContiguousSuffixRank(target, targetType.getShape());
+    auto sourceSuffixRank = getContiguousSuffixRank(source, sourceType.getShape());
+    if (succeeded(targetSuffixRank) && succeeded(sourceSuffixRank)
+        && *targetSuffixRank == targetType.getRank() && *sourceSuffixRank == sourceType.getRank()) {
+      CopyRewritePlan plan;
+      plan.kind = CopyRewritePlan::Kind::Direct;
+      plan.target = *targetPlan;
+      plan.source = *sourcePlan;
+      plan.directBytes = size;
+      return plan;
+    }
+  }
+
+  auto logicalCopyShape = inferLogicalCopyShape(targetType, sourceType, size);
+  if (failed(logicalCopyShape))
+    return failure();
+
+  auto targetSuffixRank = getContiguousSuffixRank(target, *logicalCopyShape);
+  auto sourceSuffixRank = getContiguousSuffixRank(source, *logicalCopyShape);
   if (failed(targetSuffixRank) || failed(sourceSuffixRank))
     return failure();
 
@@ -229,8 +439,8 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO
     return plan;
   }
 
-  auto targetStrides = getStaticMemRefStrides(targetType);
-  auto sourceStrides = getStaticMemRefStrides(sourceType);
+  auto targetStrides = getProvenMemRefStrides(target);
+  auto sourceStrides = getProvenMemRefStrides(source);
   if (failed(targetStrides) || failed(sourceStrides))
     return failure();
 
@@ -240,11 +450,27 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO
   plan.loop.sourceBaseOffset = plan.source.offset;
   plan.loop.outerShape.assign(logicalCopyShape->begin(), logicalCopyShape->end() - contiguousSuffixRank);
   SmallVector<int64_t> chunkShape(logicalCopyShape->end() - contiguousSuffixRank, logicalCopyShape->end());
-  plan.loop.chunkBytes = getNumElements(chunkShape) * elementByteWidth;
-  for (int64_t stride : ArrayRef<int64_t>(*targetStrides).take_front(plan.loop.outerShape.size()))
-    plan.loop.targetOuterByteStrides.push_back(stride * elementByteWidth);
-  for (int64_t stride : ArrayRef<int64_t>(*sourceStrides).take_front(plan.loop.outerShape.size()))
-    plan.loop.sourceOuterByteStrides.push_back(stride * elementByteWidth);
+  auto outerElements = checkedPositiveProduct(plan.loop.outerShape);
+  auto chunkElements = checkedPositiveProduct(chunkShape);
+  auto chunkBytes = failed(chunkElements)
+                      ? FailureOr<int64_t>(failure())
+                      : checkedPositiveMul(*chunkElements, elementByteWidth);
+  if (failed(outerElements) || failed(chunkBytes))
+    return failure();
+  plan.loop.outerElements = *outerElements;
+  plan.loop.chunkBytes = *chunkBytes;
+  for (int64_t stride : ArrayRef<int64_t>(*targetStrides).take_front(plan.loop.outerShape.size())) {
+    auto byteStride = checkedPositiveMul(stride, elementByteWidth);
+    if (failed(byteStride))
+      return failure();
+    plan.loop.targetOuterByteStrides.push_back(*byteStride);
+  }
+  for (int64_t stride : ArrayRef<int64_t>(*sourceStrides).take_front(plan.loop.outerShape.size())) {
+    auto byteStride = checkedPositiveMul(stride, elementByteWidth);
+    if (failed(byteStride))
+      return failure();
+    plan.loop.sourceOuterByteStrides.push_back(*byteStride);
+  }
   if (plan.loop.chunkBytes <= 0)
     return failure();
   return plan;
@@ -344,7 +570,7 @@ static LogicalResult rewriteCopyLikeOp(CopyOp copyOp,
   }
 
   Value c0 = createIndexConstant(rewriter, anchorOp, 0);
-  Value cUpper = createIndexConstant(rewriter, anchorOp, getNumElements(plan->loop.outerShape));
+  Value cUpper = createIndexConstant(rewriter, anchorOp, plan->loop.outerElements);
   Value cStep = createIndexConstant(rewriter, anchorOp, 1);
   auto loop = buildNormalizedScfFor(
     rewriter,

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

@@ -302,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,
@@ -365,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) {

src/PIM/Dialect/Spatial/CMakeLists.txt

@@ -7,12 +7,22 @@ add_pim_library(SpatialOps
   SpatialOpsVerify.cpp
   SpatialOpsCanonicalization.cpp
   ${PIM_SRC_ROOT}/Conversion/ONNXToSpatial/CompileTime.cpp
-  Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp
-  Transforms/MergeComputeNodes/HostOutputFinalization.cpp
-  Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp
-  Transforms/MergeComputeNodes/ProjectedFragments.cpp
   Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp
   Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp
+  Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp
+  Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp
+  Transforms/MergeComputeNodes/DeferredTransferPlanning.cpp
+  Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp
+  Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp
+  Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp
+  Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp
+  Transforms/MergeComputeNodes/DeferredResultRealization.cpp
+  Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp
+  Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp
+  Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp
+  Transforms/MergeComputeNodes/ScheduledComputeReport.cpp
+  Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp
+  Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.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"
 
@@ -27,7 +28,7 @@ def SpatTensor :
 //===----------------------------------------------------------------------===//
 
 class SpatComputeLikeBase<string mnemonic> : SpatOp<mnemonic,
-    [SingleBlock, AttrSizedOperandSegments,
+    [AttrSizedOperandSegments,
      DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmBlockArgumentNames"]>]> {
   let summary = "Compute region with attached constant weights";
 
@@ -40,7 +41,7 @@ class SpatComputeLikeBase<string mnemonic> : SpatOp<mnemonic,
     Variadic<SpatTensor>:$outputs
   );
 
-  let regions = (region SizedRegion<1>:$body);
+  let regions = (region MinSizedRegion<1>:$body);
 
   let hasVerifier = 1;
   let hasFolder = 1;
@@ -76,7 +77,7 @@ def SpatScheduledCompute : SpatComputeLikeBase<"scheduled_compute"> {
 }
 
 class SpatComputeBatchLikeBase<string mnemonic> : SpatOp<mnemonic,
-    [SingleBlock, AttrSizedOperandSegments,
+    [AttrSizedOperandSegments,
      DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmBlockArgumentNames"]>]> {
   let summary = "Tensor-native batch of equivalent compute lanes with shared weights and packed inputs";
 
@@ -90,13 +91,14 @@ class SpatComputeBatchLikeBase<string mnemonic> : SpatOp<mnemonic,
     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);
@@ -113,6 +115,7 @@ def SpatGraphComputeBatch : SpatComputeBatchLikeBase<"graph_compute_batch"> {
 }
 
 def SpatScheduledComputeBatch : SpatComputeBatchLikeBase<"scheduled_compute_batch"> {
+  let hasCanonicalizer = 1;
   let extraClassDeclaration = [{
     std::optional<::mlir::BlockArgument> getLaneArgument();
     std::optional<::mlir::BlockArgument> getWeightArgument(unsigned idx);
@@ -161,6 +164,41 @@ 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 payload derivation placeholder";
+
+  let arguments = (ins
+    Variadic<SpatTensor>:$sources
+  );
+
+  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";
 
@@ -232,6 +270,22 @@ def SpatReluPlanOp : SpatOp<"relu_plan", []> {
   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";
 
@@ -245,6 +299,7 @@ def SpatBlueprintOp : SpatOp<"blueprint", []> {
     StrAttr:$indexMap,
     OptionalAttr<StrAttr>:$mode,
     OptionalAttr<DenseI64ArrayAttr>:$fragmentOperandIndices,
+    OptionalAttr<DenseI64ArrayAttr>:$fragmentSourceSlots,
     OptionalAttr<DenseI64ArrayAttr>:$fragmentSourceOffsets,
     OptionalAttr<DenseI64ArrayAttr>:$fragmentStrides,
     OptionalAttr<StrAttr>:$conflictPolicy,

src/PIM/Dialect/Spatial/SpatialOps.cpp

@@ -10,6 +10,18 @@ using namespace mlir;
 
 namespace onnx_mlir {
 namespace spatial {
+
+RankedTensorType getGraphBatchPhysicalResultType(int64_t laneCount, RankedTensorType fragmentType) {
+  SmallVector<int64_t> shape {laneCount};
+  llvm::append_range(shape, fragmentType.getShape());
+  return RankedTensorType::get(shape, fragmentType.getElementType(), fragmentType.getEncoding());
+}
+
+FailureOr<RankedTensorType> getGraphBatchFragmentType(RankedTensorType physicalType, int64_t expectedLaneCount) {
+  if (!physicalType || physicalType.getRank() < 1 || physicalType.getDimSize(0) != expectedLaneCount)
+    return failure();
+  return RankedTensorType::get(physicalType.getShape().drop_front(), physicalType.getElementType(), physicalType.getEncoding());
+}
 namespace {
 
 std::optional<BlockArgument> getBlockArgument(Region& body, unsigned argIdx) {
@@ -238,6 +250,15 @@ void SpatScheduledCompute::getAsmBlockArgumentNames(Region& region, OpAsmSetValu
   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);

src/PIM/Dialect/Spatial/SpatialOps.hpp

@@ -30,6 +30,10 @@
 namespace onnx_mlir {
 namespace spatial {
 
+mlir::RankedTensorType getGraphBatchPhysicalResultType(int64_t laneCount, mlir::RankedTensorType fragmentType);
+mlir::FailureOr<mlir::RankedTensorType>
+getGraphBatchFragmentType(mlir::RankedTensorType physicalType, int64_t expectedLaneCount);
+
 bool isGraphComputeLike(mlir::Operation* op);
 bool isGraphBatchComputeLike(mlir::Operation* op);
 bool isScheduledComputeLike(mlir::Operation* op);

src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp

@@ -160,7 +160,7 @@ void printComputeLikeOp(ComputeOpTy op, OpAsmPrinter& printer) {
   printer << " -> ";
   printCompressedTypeSequence(printer, op.getResultTypes());
   printer << " ";
-  printer.printRegion(op.getBody(), /*printEntryBlockArgs=*/false);
+  printer.printRegion(op.getBody(), /*printEntryBlockArgs=*/!op.getBody().hasOneBlock());
 }
 
 template <typename ComputeOpTy>
@@ -290,7 +290,7 @@ void printComputeBatchLikeOp(ComputeBatchOpTy op, OpAsmPrinter& printer) {
   printer << " -> ";
   printCompressedTypeSequence(printer, op.getResultTypes());
   printer << " ";
-  printer.printRegion(op.getBody(), /*printEntryBlockArgs=*/false);
+  printer.printRegion(op.getBody(), /*printEntryBlockArgs=*/!op.getBody().hasOneBlock());
 }
 
 template <typename ComputeBatchOpTy>
@@ -407,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());
@@ -493,6 +576,10 @@ void SpatBlueprintOp::print(OpAsmPrinter& printer) {
     printer << " operandIndices ";
     printCompressedIntegerList(printer, *operandIndices);
   }
+  if (std::optional<ArrayRef<int64_t>> sourceSlots = getFragmentSourceSlots()) {
+    printer << " sourceSlots ";
+    printCompressedIntegerList(printer, *sourceSlots);
+  }
   if (std::optional<ArrayRef<int64_t>> sourceOffsets = getFragmentSourceOffsets()) {
     printer << " sourceOffsets ";
     printCompressedIntegerList(printer, *sourceOffsets);
@@ -514,6 +601,7 @@ void SpatBlueprintOp::print(OpAsmPrinter& printer) {
                                  getIndexMapAttrName().getValue(),
                                  getModeAttrName().getValue(),
                                  getFragmentOperandIndicesAttrName().getValue(),
+                                 getFragmentSourceSlotsAttrName().getValue(),
                                  getFragmentSourceOffsetsAttrName().getValue(),
                                  getFragmentStridesAttrName().getValue(),
                                  getConflictPolicyAttrName().getValue(),
@@ -537,6 +625,7 @@ ParseResult SpatBlueprintOp::parse(OpAsmParser& parser, OperationState& result)
   SmallVector<int64_t> fragmentOffsets;
   SmallVector<int64_t> fragmentSizes;
   SmallVector<int64_t> fragmentOperandIndices;
+  SmallVector<int64_t> fragmentSourceSlots;
   SmallVector<int64_t> fragmentSourceOffsets;
   SmallVector<int64_t> fragmentStrides;
 
@@ -554,6 +643,9 @@ ParseResult SpatBlueprintOp::parse(OpAsmParser& parser, OperationState& result)
   if (succeeded(parser.parseOptionalKeyword("operandIndices"))
       && parseCompressedIntegerList(parser, fragmentOperandIndices))
     return failure();
+  if (succeeded(parser.parseOptionalKeyword("sourceSlots"))
+      && parseCompressedIntegerList(parser, fragmentSourceSlots))
+    return failure();
   if (succeeded(parser.parseOptionalKeyword("sourceOffsets"))
       && parseCompressedIntegerList(parser, fragmentSourceOffsets))
     return failure();
@@ -584,6 +676,8 @@ ParseResult SpatBlueprintOp::parse(OpAsmParser& parser, OperationState& result)
     result.addAttribute("mode", mode);
   if (!fragmentOperandIndices.empty())
     result.addAttribute("fragmentOperandIndices", builder.getDenseI64ArrayAttr(fragmentOperandIndices));
+  if (!fragmentSourceSlots.empty())
+    result.addAttribute("fragmentSourceSlots", builder.getDenseI64ArrayAttr(fragmentSourceSlots));
   if (!fragmentSourceOffsets.empty())
     result.addAttribute("fragmentSourceOffsets", builder.getDenseI64ArrayAttr(fragmentSourceOffsets));
   if (!fragmentStrides.empty())

src/PIM/Dialect/Spatial/SpatialOpsCanonicalization.cpp

@@ -1,8 +1,14 @@
+#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/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
@@ -40,5 +46,177 @@ LogicalResult SpatScheduledCompute::fold(FoldAdaptor adaptor, ::llvm::SmallVecto
   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 = getOrCreateIndexConstant(rewriter, compute.getOperation(), 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"
@@ -59,6 +60,21 @@ 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;
@@ -80,7 +96,7 @@ static bool isStaticIndexExpr(Value value) {
 }
 
 static bool isSupportedLaneOffsetExpr(Value value, BlockArgument laneArg) {
-  if (value == laneArg || isStaticIndexExpr(value))
+  if (value == laneArg || isStaticIndexExpr(value) || isStaticScfForInductionVar(value))
     return true;
 
   auto affineApply = value.getDefiningOp<affine::AffineApplyOp>();
@@ -176,12 +192,18 @@ 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 =
@@ -203,8 +225,35 @@ static LogicalResult verifyOnlyConstantExternalValues(Operation* ownerOp, Region
   return success(!hasFailure);
 }
 
+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) {
+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)
@@ -219,6 +268,7 @@ static LogicalResult verifyBatchBody(ComputeBatchOpTy 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")))
@@ -436,6 +486,39 @@ LogicalResult SpatReluPlanOp::verify() {
   return success();
 }
 
+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 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");
+
+  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";
@@ -491,20 +574,26 @@ LogicalResult SpatBlueprintOp::verify() {
 
   auto stridesAttr = getFragmentStridesAttr();
   auto operandIndicesAttr = getFragmentOperandIndicesAttr();
+  auto sourceSlotsAttr = getFragmentSourceSlotsAttr();
   auto sourceOffsetsAttr = getFragmentSourceOffsetsAttr();
   if (!operandIndicesAttr)
     return emitError("fragment assembly blueprint requires fragment operand indices");
+  if (!sourceSlotsAttr)
+    return emitError("fragment assembly blueprint requires physical fragment source slots");
   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> sourceSlots = sourceSlotsAttr.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 (sourceSlots.size() != operandIndices.size())
+    return emitError("fragment source slot count must match fragment operand index count");
   if (!getConflictPolicyAttr() || !getCoveragePolicyAttr())
     return emitError("fragment assembly blueprint requires conflict and coverage policies");
   if (getConflictPolicy() != "disjoint")
@@ -539,14 +628,19 @@ LogicalResult SpatBlueprintOp::verify() {
     int64_t operandIndex = operandIndices[fragmentIndex];
     if (operandIndex < 0 || operandIndex >= operandCount)
       return emitError("fragment assembly operand index is out of range");
+    if (sourceSlots[fragmentIndex] < 0)
+      return emitError("fragment assembly physical source slot must be nonnegative");
     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");
+    if (operandType.getRank() != rank + 1)
+      return emitError("fragment assembly physical operand must have one leading source-slot dimension");
+    if (sourceSlots[fragmentIndex] >= operandType.getDimSize(0))
+      return emitError("fragment assembly physical source slot is out of range");
+    auto fragmentType = RankedTensorType::get(operandType.getShape().drop_front(), operandType.getElementType(), operandType.getEncoding());
 
     SmallVector<int64_t, 4> fragmentOffsets;
     SmallVector<int64_t, 4> fragmentSizes;
@@ -562,12 +656,12 @@ LogicalResult SpatBlueprintOp::verify() {
     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");
+    if (sourceOffsets[fragmentIndex] + fragmentElements > fragmentType.getNumElements())
+      return emitError("fragment assembly source offset exceeds the selected physical fragment bounds");
     SmallVector<int64_t, 4> sourceSliceOffsets =
-      expandFlatElementIndex(sourceOffsets[fragmentIndex], operandType.getShape());
+      expandFlatElementIndex(sourceOffsets[fragmentIndex], fragmentType.getShape());
     for (int64_t dim = 0; dim < rank; ++dim)
-      if (sourceSliceOffsets[dim] + fragmentSizes[dim] > operandType.getDimSize(dim))
+      if (sourceSliceOffsets[dim] + fragmentSizes[dim] > fragmentType.getDimSize(dim))
         return emitError("fragment assembly source offset must describe a valid unit-stride slice");
 
     for (const auto& [existingOffsets, existingSizes] : slices) {
@@ -630,7 +724,9 @@ 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(), [](Operation* op) {
+        return llvm::all_of(result.getUsers(), [result](Operation* op) {
+          if (isRecordedDeferredCommunicationSource(op, result))
+            return true;
           return !isAnySpatialComputeLike(op->getParentOp());
         });
       })) {
@@ -641,33 +737,44 @@ LogicalResult verifyComputeResultsUses(Operation* op) {
 
 template <typename ComputeOpTy>
 LogicalResult verifyComputeLikeOp(ComputeOpTy compute, StringRef opName) {
-  auto& block = compute.getBody().front();
   unsigned expectedArgCount = compute.getWeights().size() + compute.getInputs().size();
-  if (block.getNumArguments() != expectedArgCount)
+  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())) {
-    auto blockArg = compute.getWeightArgument(weightIndex);
-    if (!blockArg || blockArg->getType() != weight.getType())
+    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())) {
-    auto blockArg = compute.getInputArgument(inputIndex);
-    if (!blockArg || blockArg->getType() != input.getType())
+    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");
-  }
 
-  if (block.mightHaveTerminator()) {
-    auto yieldOp = dyn_cast_or_null<SpatYieldOp>(block.getTerminator());
-    if (!yieldOp)
+    Operation* terminator = block.getTerminator();
+    if (auto yieldOp = dyn_cast_or_null<SpatYieldOp>(terminator)) {
+      auto realized = compute->template getAttrOfType<BoolAttr>("scheduled.realized");
+      if (isScheduled && (!realized || !realized.getValue() || !compute.getBody().hasOneBlock()))
+        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();
-    auto yieldTypes = yieldOp->getOperandTypes();
-    if (resultTypes.size() != yieldTypes.size())
-      return compute.emitOpError("ComputeOp must have same number of results as yieldOp operands");
+  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, yieldTypes))) {
+  for (auto it : llvm::reverse(llvm::zip(resultTypes, yieldedTypes))) {
     auto resultType = std::get<0>(it);
     auto yieldType = std::get<1>(it);
 
@@ -677,18 +784,18 @@ LogicalResult verifyComputeLikeOp(ComputeOpTy compute, StringRef opName) {
     if (auto resultRankedType = dyn_cast<RankedTensorType>(resultType)) {
       if (auto yieldRankedType = dyn_cast<RankedTensorType>(yieldType)) {
         if (resultRankedType.getEncoding() != yieldRankedType.getEncoding())
-            return compute.emitOpError("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 compute.emitOpError("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 compute.emitOpError("ComputeOp output must not have an encoding if yieldOp operand has one");
     }
   }
-  }
 
+  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");
@@ -705,6 +812,41 @@ LogicalResult SpatGraphCompute::verify() { return verifyComputeLikeOp(*this, "sp
 
 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");
+  static constexpr StringLiteral staleAttributes[] = {
+    "exchangeId", "logicalProducer", "logicalConsumer", "sourceClass", "targetClass", "sourceCore",
+    "targetCore", "sourceLane", "targetLane", "transferKind", "resultIndex", "projectedTransfer",
+    "hostOutputOwner", "source_cpus", "source_classes", "source_lane_ranges", "target_cpus",
+    "target_classes", "target_lane_ranges", "batched", "source_operand_for_scheduled_lane",
+    "multi_source_payload"};
+  for (StringLiteral name : staleAttributes)
+    if (getOperation()->hasAttr(name))
+      return emitOpError() << "does not accept stale routing attribute '" << name
+                           << "'; source selection and shaping belong in the body and routing is derived in Phase 2";
+  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();
@@ -727,30 +869,33 @@ LogicalResult verifyComputeBatchLikeOp(ComputeBatchOpTy batch, StringRef opName)
         return batch.emitOpError("compute_batch coreIds values must be unique");
   }
 
-  Block& block = batch.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 batch.emitOpError("compute_batch body must have exactly one lane block argument");
-  unsigned expectedArgCount = 1 + batch.getWeights().size() + batch.getInputs().size() + batch.getNumResults();
     if (block.getNumArguments() != expectedArgCount)
-    return batch.emitOpError("compute_batch body block arguments must match lane, weight, input, and output operands/results");
-  auto laneArg = batch.getLaneArgument();
-  if (!laneArg || !laneArg->getType().isIndex())
+      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(batch.getWeights())) {
-    auto blockArg = batch.getWeightArgument(weightIndex);
-    if (!blockArg || blockArg->getType() != weight.getType())
+    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())) {
-    auto blockArg = batch.getInputArgument(inputIndex);
-    if (!blockArg || blockArg->getType() != input.getType())
+    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())) {
-    auto blockArg = batch.getOutputArgument(resultIndex);
-    if (!blockArg || blockArg->getType() != resultType)
+    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(batch.getOperation())))
@@ -759,7 +904,7 @@ LogicalResult verifyComputeBatchLikeOp(ComputeBatchOpTy batch, StringRef opName)
     return failure();
   if (failed(verifyOnlyConstantExternalValues(batch.getOperation(), batch.getBody(), opName)))
     return failure();
-  return verifyBatchBody(batch, block);
+  return success();
 }
 
 LogicalResult SpatGraphComputeBatch::verify() { return verifyComputeBatchLikeOp(*this, "spat.graph_compute_batch"); }

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

@@ -0,0 +1,658 @@
+#include "llvm/ADT/MapVector.h"
+
+#include "DeferredBoundaryPlanning.hpp"
+#include "DeferredCommunicationScheduling.hpp"
+#include "DeferredProjectionAnalysis.hpp"
+#include "DeferredTransferPlanning.hpp"
+
+namespace onnx_mlir::spatial {
+using namespace mlir;
+namespace {
+
+enum class BoundaryEventKind { Send, Receive };
+
+struct BoundaryEvent {
+  BoundaryEventKind kind = BoundaryEventKind::Send;
+  ScheduledTransferSlice slice;
+  LaneSet activeLanes;
+  TransferEmissionSignature emission;
+};
+
+struct SequenceNode {
+  unsigned previous = 0;
+  unsigned instruction = 0;
+};
+
+struct IntervalClass {
+  LaneSet lanes;
+  unsigned sequence = 0;
+};
+
+enum class SemanticKind {
+  Send,
+  Availability,
+  Result
+};
+
+struct SemanticKey {
+  SemanticKind kind = SemanticKind::Send;
+  TransferEmissionSignature emission;
+  DeferredExchangePlan* exchange = nullptr;
+  RequirementCoordinate coordinate;
+  Type fragmentType;
+
+  bool operator==(const SemanticKey& other) const {
+    if (kind != other.kind)
+      return false;
+    if (kind == SemanticKind::Send)
+      return emission == other.emission;
+    if (kind == SemanticKind::Result)
+      return exchange == other.exchange;
+    return exchange == other.exchange && coordinate == other.coordinate
+        && fragmentType == other.fragmentType;
+  }
+};
+
+struct PendingToken {
+  LaneSet lanes;
+  unsigned semantic = 0;
+  std::variant<unsigned, LocalAvailabilityFamily*, DeferredExchangePlan*> value;
+};
+
+struct SequenceCursor {
+  LaneSet lanes;
+  unsigned position = 0;
+};
+
+struct CanonicalAction {
+  SemanticKey key;
+  SmallVector<ScheduledTransferSlice> slices;
+  SmallVector<LocalAvailabilityFamily*> locals;
+  LaneSet instructionLanes;
+  LaneSet receiveLanes;
+  LaneSet localLanes;
+  DeferredExchangePlan* result = nullptr;
+};
+
+struct BoundaryWork {
+  BoundaryProgram program;
+  SmallVector<BoundaryEvent> events;
+};
+
+static unsigned getBoundaryIndex(SmallVectorImpl<BoundaryWork>& boundaries,
+                                 DenseMap<BoundaryKey, unsigned>& indices,
+                                 BoundaryKey key) {
+  if (auto it = indices.find(key); it != indices.end())
+    return it->second;
+  unsigned index = boundaries.size();
+  indices[key] = index;
+  BoundaryWork work;
+  work.program.key = key;
+  boundaries.push_back(std::move(work));
+  return index;
+}
+
+static unsigned getResultBoundary(DeferredExchangePlan& exchange,
+                                  const ScheduledCommunicationPlan& schedule) {
+  std::optional<unsigned> step;
+  for (const ScheduledTransferSlice& slice : schedule.slices)
+    if (slice.family->requirement->exchange == &exchange)
+      step = std::max(step.value_or(0), slice.targetInsertionStep);
+  return step.value_or(exchange.consumerStep);
+}
+
+static size_t hashSemanticKey(const SemanticKey& key) {
+  if (key.kind == SemanticKind::Send)
+    return llvm::hash_combine(key.kind,
+                              key.emission.scheduled,
+                              key.emission.payload.getAsOpaquePointer(),
+                              key.emission.fragmentType.getAsOpaquePointer(),
+                              key.emission.hasGraphLane,
+                              key.emission.sourceIsBatch);
+  if (key.kind == SemanticKind::Result)
+    return llvm::hash_combine(key.kind, key.exchange);
+  return llvm::hash_combine(key.kind, key.exchange,
+                            key.coordinate.leafIndex,
+                            key.coordinate.selectedPosition,
+                            key.fragmentType.getAsOpaquePointer());
+}
+
+static unsigned internSemantic(const SemanticKey& key,
+                               SmallVectorImpl<SemanticKey>& keys,
+                               DenseMap<size_t, SmallVector<unsigned>>& byHash) {
+  size_t hash = hashSemanticKey(key);
+  for (unsigned candidate : byHash.lookup(hash))
+    if (keys[candidate] == key)
+      return candidate;
+  unsigned id = keys.size();
+  keys.push_back(key);
+  byHash[hash].push_back(id);
+  return id;
+}
+
+static SemanticKey getEventKey(const BoundaryEvent& event) {
+  if (event.kind == BoundaryEventKind::Send) {
+    SemanticKey key;
+    key.kind = SemanticKind::Send;
+    key.emission = event.emission;
+    return key;
+  }
+  RequirementFamily& requirement = *event.slice.family->requirement;
+  SemanticKey key;
+  key.kind = SemanticKind::Availability;
+  key.exchange = requirement.exchange;
+  key.coordinate = requirement.coordinate;
+  key.fragmentType = requirement.publicationFragmentType;
+  return key;
+}
+
+static SemanticKey getLocalKey(LocalAvailabilityFamily& local) {
+  RequirementFamily& requirement = *local.requirement;
+  SemanticKey key;
+  key.kind = SemanticKind::Availability;
+  key.exchange = requirement.exchange;
+  key.coordinate = requirement.coordinate;
+  key.fragmentType = requirement.publicationFragmentType;
+  return key;
+}
+
+static SmallVector<ScheduledTransferSlice> intersectReceiveSlice(const ScheduledTransferSlice& slice,
+                                                                 const LaneSet& lanes) {
+  LaneInterval family = slice.family->targetLanes.intervals().front();
+  unsigned sliceBegin = family.begin + slice.familyOffset;
+  LaneSet active = LaneSet::range(sliceBegin, sliceBegin + slice.transferCount).intersect(lanes);
+  SmallVector<ScheduledTransferSlice> result;
+  for (LaneInterval selected : active.intervals()) {
+    ScheduledTransferSlice part = slice;
+    part.familyOffset += selected.begin - sliceBegin;
+    part.transferCount = selected.end - selected.begin;
+    result.push_back(part);
+  }
+  return result;
+}
+
+static void mergeSequenceCursors(SmallVectorImpl<SequenceCursor>& states) {
+  DenseMap<unsigned, unsigned> byPosition;
+  SmallVector<SequenceCursor> merged;
+  for (SequenceCursor& state : states) {
+    auto [it, inserted] = byPosition.try_emplace(state.position, merged.size());
+    if (inserted)
+      merged.push_back(std::move(state));
+    else
+      merged[it->second].lanes = merged[it->second].lanes.unite(state.lanes);
+  }
+  states = std::move(merged);
+}
+
+static LogicalResult appendReplayToken(const PendingToken& token,
+                                       const LaneSet& lanes,
+                                       CanonicalAction& action,
+                                       ArrayRef<BoundaryEvent> events) {
+  if (auto eventId = std::get_if<unsigned>(&token.value)) {
+    const BoundaryEvent& event = events[*eventId];
+    if (event.kind == BoundaryEventKind::Send) {
+      action.slices.push_back(event.slice);
+      action.instructionLanes =
+        action.instructionLanes.unite(event.activeLanes);
+    }
+    else {
+      llvm::append_range(action.slices, intersectReceiveSlice(event.slice, lanes));
+      action.receiveLanes = action.receiveLanes.unite(lanes);
+    }
+    return success();
+  }
+  if (auto local = std::get_if<LocalAvailabilityFamily*>(&token.value)) {
+    if (!llvm::is_contained(action.locals, *local))
+      action.locals.push_back(*local);
+    action.localLanes = action.localLanes.unite(lanes);
+    return success();
+  }
+  action.result = *std::get_if<DeferredExchangePlan*>(&token.value);
+  return success();
+}
+
+static FailureOr<SmallVector<CanonicalAction, 0>> collectCanonicalActions(
+  ArrayRef<unsigned> sequence, ArrayRef<PendingToken> tokens,
+  ArrayRef<SemanticKey> semantics, ArrayRef<BoundaryEvent> events,
+  const LaneSet& lanes) {
+  SmallVector<CanonicalAction, 0> actions;
+  for (unsigned semantic : sequence) {
+    CanonicalAction action;
+    action.key = semantics[semantic];
+    actions.push_back(std::move(action));
+  }
+  SmallVector<SequenceCursor> states {
+    {lanes, 0}
+  };
+  for (const PendingToken& token : tokens) {
+    SmallVector<SequenceCursor> next;
+    for (const SequenceCursor& state : states) {
+      LaneSet intersection = state.lanes.intersect(token.lanes);
+      LaneSet difference = state.lanes.subtract(token.lanes);
+      if (!difference.empty())
+        next.push_back({difference, state.position});
+      if (intersection.empty())
+        continue;
+      if (state.position == sequence.size()) {
+        next.push_back({intersection, state.position});
+        continue;
+      }
+      if (state.position >= sequence.size() || sequence[state.position] != token.semantic
+          || failed(appendReplayToken(token, intersection, actions[state.position], events)))
+        return failure();
+      next.push_back({intersection, state.position + 1});
+    }
+    mergeSequenceCursors(next);
+    states = std::move(next);
+  }
+  if (llvm::any_of(states, [&](const SequenceCursor& state) { return state.position != sequence.size(); }))
+    return failure();
+  return actions;
+}
+
+static bool matchesAssembly(ArrayRef<CanonicalAction> actions, size_t begin, SmallVectorImpl<unsigned>& entryOrder) {
+  if (begin >= actions.size())
+    return false;
+  const CanonicalAction& first = actions[begin];
+  DeferredExchangePlan* exchange = first.key.exchange;
+  auto& assembly = exchange->program.insertAssembly;
+  if (first.key.kind != SemanticKind::Availability || !first.locals.empty()
+      || !assembly || assembly->entries.empty()
+      || begin + assembly->entries.size() > actions.size())
+    return false;
+  SmallVector<bool> matched(assembly->entries.size());
+  for (size_t offset = 0; offset < assembly->entries.size(); ++offset) {
+    const CanonicalAction& action = actions[begin + offset];
+    if (action.key.kind != SemanticKind::Availability || !action.locals.empty()
+        || action.key.exchange != exchange || action.slices.empty())
+      return false;
+    std::optional<unsigned> entry;
+    for (auto [entryIndex, candidate] : llvm::enumerate(assembly->entries))
+      if (!matched[entryIndex] && action.key.coordinate == candidate.coordinate) {
+        entry = entryIndex;
+        break;
+      }
+    if (!entry)
+      return false;
+    matched[*entry] = true;
+    entryOrder.push_back(*entry);
+  }
+  if (assembly->entries.size() > 1
+      && llvm::any_of(ArrayRef(assembly->entries).drop_front(), [&](const DeferredInsertAssemblyEntryTemplate& entry) {
+           return entry.sourceTransform != assembly->entries.front().sourceTransform
+               || entry.sourceType != assembly->entries.front().sourceType;
+         }))
+    return false;
+  return true;
+}
+
+static bool matchesProjectionAssembly(ArrayRef<CanonicalAction> actions,
+                                      size_t begin,
+                                      const LaneSet& lanes,
+                                      unsigned& leafIndex,
+                                      SmallVectorImpl<unsigned>& positions) {
+  if (begin >= actions.size() || lanes.empty())
+    return false;
+  const CanonicalAction& first = actions[begin];
+  DeferredExchangePlan* exchange = first.key.exchange;
+  if (first.key.kind != SemanticKind::Availability || !first.locals.empty()
+      || !exchange || exchange->program.insertAssembly
+      || first.key.coordinate.leafIndex >= exchange->program.leaves.size())
+    return false;
+  leafIndex = first.key.coordinate.leafIndex;
+  const DeferredProjectionLeafTemplate& leaf = exchange->program.leaves[leafIndex];
+  if (leaf.form != DeferredLeafForm::DirectSource)
+    return false;
+  unsigned representative = lanes.intervals().front().begin;
+  unsigned positionCount = 0;
+  unsigned requirementCount = 0;
+  Type fragmentType;
+  for (RequirementFamily& requirement : exchange->requirements) {
+    if (requirement.coordinate.leafIndex != leafIndex || !requirement.targetLanes.contains(representative))
+      continue;
+    ++requirementCount;
+    positionCount = std::max(positionCount, requirement.coordinate.selectedPosition + 1);
+    if (fragmentType && fragmentType != requirement.publicationFragmentType)
+      return false;
+    fragmentType = requirement.publicationFragmentType;
+  }
+  auto fragment = dyn_cast<RankedTensorType>(fragmentType);
+  if (positionCount < 2 || requirementCount != positionCount || !fragment
+      || leaf.reconstructedType.getRank() != fragment.getRank() + 1
+      || leaf.reconstructedType.getDimSize(0) != positionCount
+      || leaf.reconstructedType.getShape().drop_front() != fragment.getShape())
+    return false;
+  SmallVector<bool> seen(positionCount);
+  for (size_t offset = 0; begin + offset < actions.size(); ++offset) {
+    const CanonicalAction& action = actions[begin + offset];
+    if (action.key.kind != SemanticKind::Availability || !action.locals.empty()
+        || action.key.exchange != exchange
+        || action.key.coordinate.leafIndex != leafIndex || action.key.fragmentType != fragmentType)
+      break;
+    unsigned position = action.key.coordinate.selectedPosition;
+    if (action.slices.empty() || position >= positionCount || seen[position])
+      return false;
+    seen[position] = true;
+    positions.push_back(position);
+  }
+  if (positions.size() < 2)
+    return false;
+  for (RequirementFamily& requirement : exchange->requirements) {
+    if (requirement.coordinate.leafIndex != leafIndex || !requirement.targetLanes.contains(representative)
+        || seen[requirement.coordinate.selectedPosition])
+      continue;
+    bool local = llvm::any_of(exchange->local, [&](const LocalAvailabilityFamily& availability) {
+      return availability.requirement == &requirement && availability.targetLanes.contains(representative);
+    });
+    if (!local)
+      return false;
+  }
+  return true;
+}
+
+static BoundaryInstructionList materializeInstructions(ArrayRef<CanonicalAction> actions,
+                                                  const LaneSet& lanes) {
+  BoundaryInstructionList result;
+  auto& instructions = result.instructions;
+  for (size_t index = 0; index < actions.size();) {
+    const CanonicalAction& action = actions[index];
+    if (action.key.kind == SemanticKind::Send) {
+      EmitSendRun run;
+      run.lanes = action.instructionLanes;
+      do {
+        llvm::append_range(run.slices, actions[index].slices);
+        run.lanes = run.lanes.unite(actions[index].instructionLanes);
+        ++index;
+      }
+      while (index < actions.size() && actions[index].key.kind == SemanticKind::Send
+             && actions[index].key.emission == action.key.emission);
+      instructions.push_back(std::move(run));
+      continue;
+    }
+    SmallVector<unsigned> assemblyEntries;
+    if (matchesAssembly(actions, index, assemblyEntries)) {
+      EmitReceiveAssemblyRun run;
+      run.lanes = lanes;
+      run.assemblyEntries = assemblyEntries;
+      run.entryOffsets.push_back(0);
+      for (size_t offset = 0; offset < assemblyEntries.size(); ++offset) {
+        llvm::append_range(run.slices, actions[index + offset].slices);
+        run.entryOffsets.push_back(run.slices.size());
+      }
+      instructions.push_back(std::move(run));
+      index += assemblyEntries.size();
+      continue;
+    }
+    unsigned projectionLeaf = 0;
+    SmallVector<unsigned> projectionPositions;
+    if (matchesProjectionAssembly(actions, index, lanes, projectionLeaf, projectionPositions)) {
+      EmitReceiveAssemblyRun run;
+      run.lanes = lanes;
+      run.projectionLeaf = projectionLeaf;
+      run.assemblyEntries = projectionPositions;
+      run.entryOffsets.push_back(0);
+      for (size_t offset = 0; offset < projectionPositions.size(); ++offset) {
+        llvm::append_range(run.slices, actions[index + offset].slices);
+        run.entryOffsets.push_back(run.slices.size());
+      }
+      instructions.push_back(std::move(run));
+      index += projectionPositions.size();
+      continue;
+    }
+    if (action.key.kind == SemanticKind::Availability) {
+      ResolveAvailability availability;
+      availability.exchange = action.key.exchange;
+      availability.coordinate = action.key.coordinate;
+      availability.fragmentType = action.key.fragmentType;
+      if (!action.slices.empty()) {
+        EmitReceiveRun receive;
+        receive.slices = action.slices;
+        receive.entryOffsets = {0, receive.slices.size()};
+        receive.lanes = action.receiveLanes;
+        availability.alternatives.push_back(
+          {receive.lanes, AvailabilitySource(std::move(receive))});
+      }
+      if (!action.locals.empty()) {
+        MaterializeLocalFamily local {action.locals, action.localLanes};
+        availability.alternatives.push_back(
+          {local.lanes, AvailabilitySource(std::move(local))});
+      }
+      instructions.push_back(std::move(availability));
+      ++index;
+      continue;
+    }
+    instructions.push_back(ProduceDeferredResult {action.result, lanes});
+    ++index;
+  }
+  return result;
+}
+
+static void addTokenToClasses(const LaneSet& active,
+                              unsigned semantic,
+                              SmallVectorImpl<IntervalClass>& classes,
+                              SmallVectorImpl<SequenceNode>& nodes,
+                              DenseMap<std::pair<unsigned, unsigned>, unsigned>& interned) {
+  SmallVector<IntervalClass> next;
+  for (const IntervalClass& current : classes) {
+    LaneSet intersection = current.lanes.intersect(active);
+    LaneSet difference = current.lanes.subtract(active);
+    if (!difference.empty())
+      next.push_back({difference, current.sequence});
+    if (intersection.empty())
+      continue;
+    auto key = std::make_pair(current.sequence, semantic);
+    auto [it, inserted] = interned.try_emplace(key, nodes.size());
+    if (inserted)
+      nodes.push_back({current.sequence, semantic});
+    next.push_back({intersection, it->second});
+  }
+  classes = std::move(next);
+}
+
+struct SequenceClass {
+  LaneSet lanes;
+  SmallVector<unsigned> sequence;
+};
+
+static SmallVector<SequenceClass>
+buildSequenceClasses(unsigned laneCount, ArrayRef<PendingToken> tokens) {
+  SmallVector<IntervalClass> classes {
+    {LaneSet::all(laneCount), 0}
+  };
+  SmallVector<SequenceNode> nodes {
+    {0, 0}
+  };
+  DenseMap<std::pair<unsigned, unsigned>, unsigned> interned;
+  for (const PendingToken& token : tokens)
+    addTokenToClasses(token.lanes, token.semantic, classes, nodes, interned);
+
+  SmallVector<SequenceClass> result;
+  DenseMap<unsigned, unsigned> classBySequence;
+  SmallVector<unsigned> classSequences;
+  for (const IntervalClass& interval : classes) {
+    auto [it, inserted] = classBySequence.try_emplace(
+      interval.sequence, result.size());
+    if (inserted) {
+      result.push_back({interval.lanes, {}});
+      classSequences.push_back(interval.sequence);
+    }
+    else
+      result[it->second].lanes = result[it->second].lanes.unite(interval.lanes);
+  }
+  for (auto [behavior, sequence] : llvm::zip_equal(result, classSequences)) {
+    for (unsigned node = sequence; node != 0; node = nodes[node].previous)
+      behavior.sequence.push_back(nodes[node].instruction);
+    std::reverse(behavior.sequence.begin(), behavior.sequence.end());
+  }
+  return result;
+}
+
+static SmallVector<DeferredExchangePlan*>
+getProducedExchanges(const BoundaryInstructionList& list) {
+  SmallVector<DeferredExchangePlan*> result;
+  for (const BoundaryInstruction& instruction : list.instructions)
+    if (auto produced = std::get_if<ProduceDeferredResult>(&instruction))
+      result.push_back(produced->exchange);
+  return result;
+}
+
+static LogicalResult buildCanonicalBoundary(
+  BoundaryProgram& boundary, ArrayRef<BoundaryEvent> events,
+  ArrayRef<PendingToken> tokens, ArrayRef<SemanticKey> semantics) {
+  unsigned laneCount = boundary.key.scheduled->cores.size();
+  SmallVector<SequenceClass> classes =
+    buildSequenceClasses(laneCount, tokens);
+  if (classes.empty())
+    return failure();
+  size_t prefix = classes.front().sequence.size();
+  for (const SequenceClass& behavior : ArrayRef(classes).drop_front()) {
+    prefix = std::min(prefix, behavior.sequence.size());
+    size_t index = 0;
+    while (index < prefix
+           && behavior.sequence[index] == classes.front().sequence[index])
+      ++index;
+    prefix = index;
+  }
+  auto prefixActions = collectCanonicalActions(
+    ArrayRef(classes.front().sequence).take_front(prefix), tokens, semantics,
+    events, LaneSet::all(laneCount));
+  if (failed(prefixActions))
+    return failure();
+  boundary.root = materializeInstructions(*prefixActions, LaneSet::all(laneCount));
+  if (classes.size() == 1)
+    return success();
+
+  auto dispatch = std::make_unique<LaneDispatch>();
+  SmallVector<int64_t> classIds(laneCount);
+  for (auto [classId, behavior] : llvm::enumerate(classes)) {
+    for (LaneInterval interval : behavior.lanes.intervals())
+      for (unsigned lane = interval.begin; lane < interval.end; ++lane)
+        classIds[lane] = classId;
+    auto actions = collectCanonicalActions(
+      behavior.sequence, tokens, semantics, events, behavior.lanes);
+    if (failed(actions))
+      return failure();
+    dispatch->branches.push_back(materializeInstructions(
+      ArrayRef(*actions).drop_front(prefix), behavior.lanes));
+    auto produced = getProducedExchanges(dispatch->branches.back());
+    if (classId == 0)
+      dispatch->producedExchanges = std::move(produced);
+    else if (produced != dispatch->producedExchanges)
+      return failure();
+  }
+  dispatch->branchByLane = StaticIntSequence::fromValues(classIds);
+  boundary.root.instructions.push_back(std::move(dispatch));
+  return success();
+}
+
+} // namespace
+
+FailureOr<DeferredBoundaryPlan> buildDeferredBoundaryPlan(DeferredTransferPlan& transfers,
+                                                          const ScheduledCommunicationPlan& schedule) {
+  DeferredBoundaryPlan result;
+  SmallVector<BoundaryWork> boundaries;
+  DenseMap<BoundaryKey, unsigned> boundaryIndices;
+  for (const ScheduledTransferSlice& slice : schedule.slices) {
+    ExternalTransferFamily& family = *slice.family;
+    unsigned sourceIndex =
+      getBoundaryIndex(boundaries, boundaryIndices,
+                       {family.sourceScheduled, slice.sourceInsertionStep});
+    unsigned targetIndex =
+      getBoundaryIndex(boundaries, boundaryIndices,
+                       {family.targetScheduled, slice.targetInsertionStep});
+    unsigned sourceLane = family.requirement->producer->scheduledLane;
+    boundaries[sourceIndex].events.push_back({BoundaryEventKind::Send,
+                                              slice,
+                                              LaneSet::range(sourceLane, sourceLane + 1),
+                                              getTransferEmissionSignature(family)});
+    LaneInterval familyLanes = family.targetLanes.intervals().front();
+    unsigned targetBegin = familyLanes.begin + slice.familyOffset;
+    boundaries[targetIndex].events.push_back({BoundaryEventKind::Receive,
+                                              slice,
+                                              LaneSet::range(targetBegin, targetBegin + slice.transferCount),
+                                              getTransferEmissionSignature(family)});
+  }
+
+  DenseMap<BoundaryKey, SmallVector<LocalAvailabilityFamily*>> locals;
+  DenseMap<BoundaryKey, SmallVector<DeferredExchangePlan*>> exchanges;
+  for (const std::unique_ptr<DeferredExchangePlan>& exchange : transfers.exchanges) {
+    unsigned step = getResultBoundary(*exchange, schedule);
+    BoundaryKey key {exchange->target, step};
+    getBoundaryIndex(boundaries, boundaryIndices, key);
+    for (LocalAvailabilityFamily& local : exchange->local)
+      locals[key].push_back(&local);
+    exchanges[key].push_back(exchange.get());
+    auto resultPlan = buildDeferredResultPlan(*exchange);
+    if (failed(resultPlan))
+      return exchange->deferred.emitOpError("cannot evaluate deferred result lane functions"), failure();
+    result.results.push_back(std::move(*resultPlan));
+  }
+
+  DenseMap<ScheduledInfo*, unsigned> scheduledOrder;
+  for (auto [index, scheduled] : llvm::enumerate(transfers.scheduled))
+    scheduledOrder[&scheduled] = index;
+  llvm::stable_sort(boundaries, [&](const BoundaryWork& lhs, const BoundaryWork& rhs) {
+    return std::tie(scheduledOrder[lhs.program.key.scheduled], lhs.program.key.insertionStep)
+         < std::tie(scheduledOrder[rhs.program.key.scheduled], rhs.program.key.insertionStep);
+  });
+  for (BoundaryWork& work : boundaries) {
+    BoundaryProgram& boundary = work.program;
+    SmallVector<PendingToken> tokens;
+    SmallVector<SemanticKey> semantics;
+    DenseMap<size_t, SmallVector<unsigned>> semanticsByHash;
+    DenseMap<unsigned, SmallVector<LocalAvailabilityFamily*>> localsBySemantic;
+    SmallPtrSet<LocalAvailabilityFamily*, 8> emittedLocals;
+    for (LocalAvailabilityFamily* local : locals[boundary.key]) {
+      unsigned semantic =
+        internSemantic(getLocalKey(*local), semantics, semanticsByHash);
+      localsBySemantic[semantic].push_back(local);
+    }
+    SmallVector<unsigned> eventSemantics;
+    DenseMap<unsigned, SmallVector<unsigned>> receivesBySemantic;
+    for (auto [eventId, event] : llvm::enumerate(work.events)) {
+      unsigned semantic =
+        internSemantic(getEventKey(event), semantics, semanticsByHash);
+      eventSemantics.push_back(semantic);
+      if (event.kind == BoundaryEventKind::Receive)
+        receivesBySemantic[semantic].push_back(eventId);
+    }
+    llvm::SmallDenseSet<unsigned, 8> emittedReceives;
+    for (auto [eventId, event] : llvm::enumerate(work.events)) {
+      unsigned semantic = eventSemantics[eventId];
+      if (event.kind == BoundaryEventKind::Send) {
+        tokens.push_back({LaneSet::all(
+                            boundary.key.scheduled->cores.size()),
+                          semantic, static_cast<unsigned>(eventId)});
+        continue;
+      }
+      if (!emittedReceives.insert(semantic).second)
+        continue;
+      for (unsigned receive : receivesBySemantic[semantic])
+        tokens.push_back({work.events[receive].activeLanes, semantic, receive});
+      for (LocalAvailabilityFamily* local : localsBySemantic[semantic]) {
+        tokens.push_back({local->targetLanes, semantic, local});
+        emittedLocals.insert(local);
+      }
+      localsBySemantic.erase(semantic);
+    }
+    for (LocalAvailabilityFamily* local : locals[boundary.key])
+      if (!emittedLocals.contains(local)) {
+        unsigned semantic =
+          internSemantic(getLocalKey(*local), semantics, semanticsByHash);
+        tokens.push_back({local->targetLanes, semantic, local});
+      }
+    for (DeferredExchangePlan* exchange : exchanges[boundary.key]) {
+      SemanticKey key;
+      key.kind = SemanticKind::Result;
+      key.exchange = exchange;
+      unsigned semantic = internSemantic(key, semantics, semanticsByHash);
+      tokens.push_back({LaneSet::all(exchange->targetLaneCount), semantic, exchange});
+    }
+    if (failed(buildCanonicalBoundary(boundary, work.events, tokens, semantics)))
+      return boundary.key.scheduled->op->emitOpError("cannot construct canonical boundary program"), failure();
+    result.boundaries.push_back(std::move(boundary));
+  }
+  return result;
+}
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,102 @@
+#pragma once
+
+#include "DeferredCommunicationScheduling.hpp"
+#include "DeferredResultRealization.hpp"
+
+namespace onnx_mlir::spatial {
+
+struct DeferredTransferPlan;
+
+struct BoundaryKey {
+  ScheduledInfo* scheduled = nullptr;
+  unsigned insertionStep = 0;
+
+  bool operator==(const BoundaryKey& other) const {
+    return scheduled == other.scheduled
+        && insertionStep == other.insertionStep;
+  }
+};
+
+struct EmitSendRun {
+  llvm::SmallVector<ScheduledTransferSlice> slices;
+  LaneSet lanes;
+};
+struct EmitReceiveRun {
+  llvm::SmallVector<ScheduledTransferSlice> slices;
+  llvm::SmallVector<size_t> entryOffsets;
+  LaneSet lanes;
+};
+struct EmitReceiveAssemblyRun {
+  llvm::SmallVector<ScheduledTransferSlice> slices;
+  llvm::SmallVector<size_t> entryOffsets;
+  llvm::SmallVector<unsigned> assemblyEntries;
+  std::optional<unsigned> projectionLeaf;
+  LaneSet lanes;
+};
+struct MaterializeLocalFamily {
+  llvm::SmallVector<LocalAvailabilityFamily*> families;
+  LaneSet lanes;
+};
+using AvailabilitySource =
+  std::variant<EmitReceiveRun, MaterializeLocalFamily>;
+struct AvailabilityAlternative {
+  LaneSet lanes;
+  AvailabilitySource source;
+};
+struct ResolveAvailability {
+  DeferredExchangePlan* exchange = nullptr;
+  RequirementCoordinate coordinate;
+  mlir::Type fragmentType;
+  llvm::SmallVector<AvailabilityAlternative, 2> alternatives;
+};
+struct ProduceDeferredResult {
+  DeferredExchangePlan* exchange = nullptr;
+  LaneSet lanes;
+};
+
+struct BoundaryInstructionList;
+struct LaneDispatch;
+using BoundaryInstruction = std::variant<EmitSendRun, ResolveAvailability,
+  EmitReceiveAssemblyRun, ProduceDeferredResult,
+  std::unique_ptr<LaneDispatch>>;
+struct BoundaryInstructionList {
+  llvm::SmallVector<BoundaryInstruction, 0> instructions;
+};
+struct LaneDispatch {
+  StaticIntSequence branchByLane = StaticIntSequence::uniform(0, 1);
+  llvm::SmallVector<BoundaryInstructionList> branches;
+  llvm::SmallVector<DeferredExchangePlan*> producedExchanges;
+};
+struct BoundaryProgram {
+  BoundaryKey key;
+  BoundaryInstructionList root;
+};
+
+struct DeferredBoundaryPlan {
+  llvm::SmallVector<BoundaryProgram, 0> boundaries;
+  llvm::SmallVector<DeferredResultPlan, 0> results;
+};
+
+mlir::FailureOr<DeferredBoundaryPlan> buildDeferredBoundaryPlan(DeferredTransferPlan& transfers,
+                                                                const ScheduledCommunicationPlan& schedule);
+
+} // namespace onnx_mlir::spatial
+
+namespace llvm {
+template <>
+struct DenseMapInfo<onnx_mlir::spatial::BoundaryKey> {
+  static onnx_mlir::spatial::BoundaryKey getEmptyKey() {
+    return {DenseMapInfo<onnx_mlir::spatial::ScheduledInfo*>::getEmptyKey(), 0};
+  }
+  static onnx_mlir::spatial::BoundaryKey getTombstoneKey() {
+    return {DenseMapInfo<onnx_mlir::spatial::ScheduledInfo*>::getTombstoneKey(), 0};
+  }
+  static unsigned getHashValue(const onnx_mlir::spatial::BoundaryKey& key) {
+    return hash_combine(key.scheduled, key.insertionStep);
+  }
+  static bool isEqual(const onnx_mlir::spatial::BoundaryKey& lhs,
+                      const onnx_mlir::spatial::BoundaryKey& rhs) {
+    return lhs == rhs;
+  }
+};
+} // namespace llvm

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

@@ -0,0 +1,824 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+
+#include "DeferredBoundaryRealization.hpp"
+#include "DeferredResultRealization.hpp"
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
+
+namespace onnx_mlir::spatial {
+using namespace mlir;
+namespace {
+
+struct LogicalTransferMetadataView {
+  StaticIntSequenceChain channels;
+  StaticIntSequenceChain parents;
+  StaticIntSequenceChain parentCounts;
+  StaticIntSequenceChain sourceCores;
+  StaticIntSequenceChain targetCores;
+  StaticIntSequenceChain targetLanes;
+  StaticIntSequenceChain localOffsets;
+
+  size_t size() const { return channels.size(); }
+};
+
+static void appendMetadata(const ScheduledTransferSlice& slice,
+                           LogicalTransferMetadataView& metadata) {
+  ExternalTransferFamily& family = *slice.family;
+  LaneInterval familyLanes = family.targetLanes.intervals().front();
+  LaneInterval requirementLanes = family.requirement->targetLanes.intervals().front();
+  size_t count = slice.transferCount;
+  size_t familyIndex = slice.familyOffset;
+  size_t targetLane = familyLanes.begin + familyIndex;
+  metadata.channels.append(family.channelIds, familyIndex, count);
+  metadata.parents.append(StaticIntSequence::uniform(
+    family.requirement->exchange->exchangeId, count));
+  metadata.parentCounts.append(StaticIntSequence::uniform(
+    family.requirement->exchange->externalTransferCount, count));
+  metadata.sourceCores.append(family.sourceCores, familyIndex, count);
+  metadata.targetCores.append(family.targetCores, familyIndex, count);
+  metadata.targetLanes.append(
+    StaticIntSequence::affine(targetLane, 1, count));
+  if (family.requirement->producerLocalOffsets)
+    metadata.localOffsets.append(
+      *family.requirement->producerLocalOffsets,
+      targetLane - requirementLanes.begin, count);
+  else
+    metadata.localOffsets.append(StaticIntSequence::uniform(0, count));
+}
+
+static LogicalTransferMetadataView
+buildMetadataView(ArrayRef<ScheduledTransferSlice> slices) {
+  LogicalTransferMetadataView metadata;
+  for (const ScheduledTransferSlice& slice : slices)
+    appendMetadata(slice, metadata);
+  return metadata;
+}
+
+static void setLogicalTransferMetadata(
+  Operation* op, const LogicalTransferMetadataView& metadata) {
+  size_t logicalCount = metadata.size();
+  OpBuilder builder(op);
+  if (logicalCount == 1) {
+    op->setAttr("raptor.exchange_id",
+                builder.getI64IntegerAttr(metadata.channels.valueAt(0)));
+    op->setAttr("raptor.channel_id",
+                builder.getI64IntegerAttr(metadata.channels.valueAt(0)));
+    op->setAttr("raptor.parent_exchange_id",
+                builder.getI64IntegerAttr(metadata.parents.valueAt(0)));
+    op->setAttr("raptor.parent_transfer_count",
+                builder.getI64IntegerAttr(metadata.parentCounts.valueAt(0)));
+    op->setAttr("raptor.source_core",
+                builder.getI64IntegerAttr(metadata.sourceCores.valueAt(0)));
+    op->setAttr("raptor.target_core",
+                builder.getI64IntegerAttr(metadata.targetCores.valueAt(0)));
+    return;
+  }
+  op->setAttr("raptor.batch_transfer_count", builder.getI64IntegerAttr(logicalCount));
+  setStaticIntSequenceAttr(op, "raptor.batch_channel_ids",
+    metadata.channels.canonicalize(), logicalCount);
+  setStaticIntSequenceAttr(op, "raptor.batch_source_cores",
+    metadata.sourceCores.canonicalize(), logicalCount);
+  setStaticIntSequenceAttr(op, "raptor.batch_target_cores",
+    metadata.targetCores.canonicalize(), logicalCount);
+  setStaticIntSequenceAttr(op, "raptor.batch_parent_exchange_ids",
+    metadata.parents.canonicalize(), logicalCount);
+  setStaticIntSequenceAttr(op, "raptor.batch_parent_transfer_counts",
+    metadata.parentCounts.canonicalize(), logicalCount);
+}
+
+static Value
+lookup(ArrayRef<int64_t> table, Value position, Operation* anchor, DeferredEmissionContext& context, Location loc) {
+  return emitStaticIntLookup(
+    StaticIntSequence::fromValues(table), position, anchor, context.constants, context.rewriter, loc);
+}
+
+static OpFoldResult lookupGeometry(
+  ArrayRef<int64_t> table, Value position, Operation* anchor, DeferredEmissionContext& context, Location loc) {
+  StaticIntSequence sequence = StaticIntSequence::fromValues(table);
+  if (sequence.getKind() == StaticIntSequenceKind::Uniform)
+    return context.rewriter.getIndexAttr(sequence.valueAt(0));
+  return emitStaticIntLookup(sequence, position, anchor, context.constants, context.rewriter, loc);
+}
+
+static FailureOr<Value> materializeSendPayload(const RequirementFamily& requirement,
+                                               Value localOffset,
+                                               DeferredEmissionContext& context,
+                                               Location loc) {
+  Value payload = requirement.producer->payload;
+  if (payload.getType() == requirement.publicationFragmentType)
+    return payload;
+  auto payloadType = dyn_cast<RankedTensorType>(payload.getType());
+  auto fragmentType = dyn_cast<RankedTensorType>(requirement.publicationFragmentType);
+  if (!payloadType || !fragmentType || !requirement.graphLanes || payloadType.getRank() != fragmentType.getRank() + 1)
+    return failure();
+  MixedSliceGeometry geometry;
+  geometry.offsets.assign(payloadType.getRank(), context.rewriter.getIndexAttr(0));
+  geometry.sizes.push_back(context.rewriter.getIndexAttr(1));
+  geometry.strides.assign(payloadType.getRank(), context.rewriter.getIndexAttr(1));
+  geometry.offsets.front() = localOffset;
+  for (int64_t dimension : fragmentType.getShape())
+    geometry.sizes.push_back(context.rewriter.getIndexAttr(dimension));
+  SmallVector<int64_t> unitShape {1};
+  llvm::append_range(unitShape, fragmentType.getShape());
+  auto unitType = RankedTensorType::get(unitShape, fragmentType.getElementType());
+  Value unit = extractMixedSliceOrIdentity(context.rewriter, loc, payload, unitType, geometry);
+  return removeLeadingUnitTensorDimension(context.rewriter, loc, unit, fragmentType);
+}
+
+static LogicalResult
+emitSendRun(const EmitSendRun& run, Value lane, unsigned laneCount, DeferredEmissionContext& context) {
+  SmallVector<LogicalTransferMetadataView, 0> metadataByLane(laneCount);
+  for (const ScheduledTransferSlice& slice : run.slices) {
+    unsigned sourceLane = slice.family->requirement->producer->scheduledLane;
+    appendMetadata(slice, metadataByLane[sourceLane]);
+  }
+  LogicalTransferMetadataView logical = buildMetadataView(run.slices);
+  size_t actionCount = 0;
+  for (const LogicalTransferMetadataView& laneMetadata : metadataByLane)
+    actionCount = std::max(actionCount, laneMetadata.size());
+  SmallVector<int64_t> channelTable(
+    actionCount * laneCount, logical.channels.valueAt(0));
+  SmallVector<int64_t> sourceTable(
+    actionCount * laneCount, logical.sourceCores.valueAt(0));
+  SmallVector<int64_t> targetTable(
+    actionCount * laneCount, logical.targetCores.valueAt(0));
+  SmallVector<int64_t> offsetTable(
+    actionCount * laneCount, logical.localOffsets.valueAt(0));
+  SmallVector<int64_t> counts(laneCount);
+  for (unsigned sourceLane = 0; sourceLane < laneCount; ++sourceLane) {
+    const LogicalTransferMetadataView& source = metadataByLane[sourceLane];
+    counts[sourceLane] = source.size();
+    for (size_t action = 0; action < source.size(); ++action) {
+      size_t index = action * laneCount + sourceLane;
+      channelTable[index] = source.channels.valueAt(action);
+      sourceTable[index] = source.sourceCores.valueAt(action);
+      targetTable[index] = source.targetCores.valueAt(action);
+      offsetTable[index] = source.localOffsets.valueAt(action);
+    }
+  }
+  ExternalTransferFamily& firstFamily = *run.slices.front().family;
+  RequirementFamily& requirement = *firstFamily.requirement;
+  Operation* anchor = requirement.exchange->deferred;
+  Location loc = requirement.exchange->deferred.getLoc();
+  auto emitOne = [&](Value position) -> LogicalResult {
+    Value localOffset = lookup(offsetTable, position, anchor, context, loc);
+    auto payload = materializeSendPayload(requirement, localOffset, context, loc);
+    if (failed(payload))
+      return failure();
+    auto send = SpatChannelSendOp::create(context.rewriter,
+                                          loc,
+                                          lookup(channelTable, position, anchor, context, loc),
+                                          lookup(sourceTable, position, anchor, context, loc),
+                                          lookup(targetTable, position, anchor, context, loc),
+                                          *payload);
+    setLogicalTransferMetadata(send, logical);
+    return success();
+  };
+  Value runtimeLane = lane ? lane : context.constants.getIndex(0);
+  if (actionCount == 1)
+    return emitOne(runtimeLane);
+  bool uniformCount = true;
+  std::optional<int64_t> firstCount;
+  for (LaneInterval interval : run.lanes.intervals())
+    for (unsigned activeLane = interval.begin; activeLane < interval.end; ++activeLane) {
+      if (!firstCount)
+        firstCount = counts[activeLane];
+      uniformCount &= counts[activeLane] == *firstCount;
+    }
+  Value count =
+    uniformCount ? context.constants.getIndex(*firstCount) : lookup(counts, runtimeLane, anchor, context, loc);
+  auto loop = buildNormalizedScfFor(context.rewriter,
+                                    loc,
+                                    context.constants.getIndex(0),
+                                    count,
+                                    context.constants.getIndex(1),
+                                    ValueRange {},
+                                    [&](OpBuilder&, Location, Value index, ValueRange, SmallVectorImpl<Value>&) {
+                                      Value base = affineMulConst(context.rewriter, loc, index, laneCount, anchor);
+                                      Value position = arith::AddIOp::create(context.rewriter, loc, base, runtimeLane);
+                                      return emitOne(position);
+                                    });
+  return success(succeeded(loop));
+}
+
+static FailureOr<Value>
+emitReceiveValue(const EmitReceiveRun& run, Value lane, unsigned laneCount, DeferredEmissionContext& context) {
+  LogicalTransferMetadataView metadata = buildMetadataView(run.slices);
+  RequirementFamily& requirement = *run.slices.front().family->requirement;
+  Operation* anchor = requirement.exchange->deferred;
+  Location loc = requirement.exchange->deferred.getLoc();
+  SmallVector<int64_t> channelTable(laneCount, metadata.channels.valueAt(0));
+  SmallVector<int64_t> sourceTable(laneCount, metadata.sourceCores.valueAt(0));
+  SmallVector<int64_t> targetTable(laneCount, metadata.targetCores.valueAt(0));
+  for (size_t index = 0; index < metadata.size(); ++index) {
+    unsigned targetLane = metadata.targetLanes.valueAt(index);
+    channelTable[targetLane] = metadata.channels.valueAt(index);
+    sourceTable[targetLane] = metadata.sourceCores.valueAt(index);
+    targetTable[targetLane] = metadata.targetCores.valueAt(index);
+  }
+  Value position = lane ? lane : context.constants.getIndex(0);
+  auto receive = SpatChannelReceiveOp::create(context.rewriter,
+                                              loc,
+                                              requirement.publicationFragmentType,
+                                              lookup(channelTable, position, anchor, context, loc),
+                                              lookup(sourceTable, position, anchor, context, loc),
+                                              lookup(targetTable, position, anchor, context, loc));
+  setLogicalTransferMetadata(receive, metadata);
+  return receive.getOutput();
+}
+
+static LogicalResult
+emitConditionalSendRun(const EmitSendRun& run, Value lane, unsigned laneCount, DeferredEmissionContext& context) {
+  if (run.lanes.size() == laneCount)
+    return emitSendRun(run, lane, laneCount, context);
+  if (!lane)
+    return failure();
+  SmallVector<int64_t> active(laneCount);
+  for (LaneInterval interval : run.lanes.intervals())
+    for (unsigned index = interval.begin; index < interval.end; ++index)
+      active[index] = 1;
+  Location loc = run.slices.front().family->requirement->exchange->deferred.getLoc();
+  Value selected = lookup(active, lane, run.slices.front().family->sourceScheduled->op, context, loc);
+  Value condition =
+    arith::CmpIOp::create(context.rewriter, loc, arith::CmpIPredicate::ne, selected, context.constants.getIndex(0));
+  auto conditional = scf::IfOp::create(context.rewriter, loc, TypeRange {}, condition, false);
+  Block& block = conditional.getThenRegion().front();
+  auto yield = cast<scf::YieldOp>(block.getTerminator());
+  OpBuilder::InsertionGuard guard(context.rewriter);
+  context.rewriter.setInsertionPoint(yield);
+  return emitSendRun(run, lane, laneCount, context);
+}
+
+static FailureOr<Value> materializeLocalValue(const MaterializeLocalFamily& local,
+                                              Value lane,
+                                              unsigned laneCount,
+                                              DeferredEmissionContext& context) {
+  RequirementFamily& reference = *local.families.front()->requirement;
+  Value fragment = reference.producer->payload;
+  if (fragment.getType() != reference.publicationFragmentType) {
+    SmallVector<int64_t> offsets(laneCount);
+    for (LocalAvailabilityFamily* family : local.families) {
+      RequirementFamily& requirement = *family->requirement;
+      LaneInterval requirementLanes = requirement.targetLanes.intervals().front();
+      for (LaneInterval interval : family->targetLanes.intervals())
+        for (unsigned targetLane = interval.begin; targetLane < interval.end; ++targetLane)
+          offsets[targetLane] = requirement.producerLocalOffsets->valueAt(targetLane - requirementLanes.begin);
+    }
+    Location loc = reference.exchange->deferred.getLoc();
+    Value position = lane ? lane : context.constants.getIndex(0);
+    auto materialized = materializeSendPayload(
+      reference, lookup(offsets, position, reference.exchange->deferred, context, loc), context, loc);
+    if (failed(materialized))
+      return failure();
+    fragment = *materialized;
+  }
+  return fragment;
+}
+
+static const DeferredResultPlan* findResultPlan(ArrayRef<DeferredResultPlan> results, DeferredExchangePlan* exchange) {
+  auto it = llvm::find_if(results, [&](const DeferredResultPlan& result) { return result.exchange == exchange; });
+  return it == results.end() ? nullptr : &*it;
+}
+
+static FailureOr<Value> applyAssemblyTransform(Value source,
+                                               const DeferredInsertAssemblyEntryTemplate& entry,
+                                               DeferredEmissionContext& context,
+                                               Location loc) {
+  auto sourceType = dyn_cast<RankedTensorType>(source.getType());
+  auto targetType = entry.sourceType;
+  if (!sourceType)
+    return failure();
+  if (sourceType == targetType)
+    return source;
+  if (targetType.getRank() == sourceType.getRank() + 1 && targetType.getDimSize(0) == 1
+      && targetType.getShape().drop_front() == sourceType.getShape())
+    return addLeadingUnitTensorDimension(context.rewriter, loc, source);
+  if (sourceType.getRank() != targetType.getRank() + 1 || sourceType.getDimSize(0) != 1
+      || sourceType.getShape().drop_front() != targetType.getShape())
+    return failure();
+  return removeLeadingUnitTensorDimension(context.rewriter, loc, source, targetType);
+}
+
+static FailureOr<Value> emitProjectionAssemblyRun(const EmitReceiveAssemblyRun& run,
+                                                  Value lane,
+                                                  unsigned laneCount,
+                                                  DeferredEmissionContext& context) {
+  if (!run.projectionLeaf || run.slices.empty() || run.entryOffsets.size() != run.assemblyEntries.size() + 1)
+    return failure();
+  DeferredExchangePlan* exchange = run.slices.front().family->requirement->exchange;
+  unsigned leafIndex = *run.projectionLeaf;
+  if (leafIndex >= exchange->program.leaves.size())
+    return failure();
+  const DeferredProjectionLeafTemplate& leaf = exchange->program.leaves[leafIndex];
+  size_t entryCount = run.assemblyEntries.size();
+  LogicalTransferMetadataView metadata = buildMetadataView(run.slices);
+  SmallVector<int64_t> channelTable(
+    entryCount * laneCount, metadata.channels.valueAt(0));
+  SmallVector<int64_t> sourceTable(
+    entryCount * laneCount, metadata.sourceCores.valueAt(0));
+  SmallVector<int64_t> targetTable(
+    entryCount * laneCount, metadata.targetCores.valueAt(0));
+  SmallVector<int64_t> positions(entryCount * laneCount);
+  for (size_t entry = 0; entry < entryCount; ++entry) {
+    ArrayRef<ScheduledTransferSlice> slices =
+      ArrayRef(run.slices).slice(run.entryOffsets[entry], run.entryOffsets[entry + 1] - run.entryOffsets[entry]);
+    LogicalTransferMetadataView item = buildMetadataView(slices);
+    for (size_t index = 0; index < item.size(); ++index) {
+      size_t tableIndex = entry * laneCount + item.targetLanes.valueAt(index);
+      channelTable[tableIndex] = item.channels.valueAt(index);
+      sourceTable[tableIndex] = item.sourceCores.valueAt(index);
+      targetTable[tableIndex] = item.targetCores.valueAt(index);
+      positions[tableIndex] = run.assemblyEntries[entry];
+    }
+  }
+  Location loc = exchange->deferred.getLoc();
+  Value initial = tensor::EmptyOp::create(
+    context.rewriter, loc, leaf.reconstructedType.getShape(), leaf.reconstructedType.getElementType());
+  auto emitEntry = [&](Value entry, Value current) -> FailureOr<Value> {
+    Value tableIndex = entry;
+    if (lane) {
+      Value base = affineMulConst(context.rewriter, loc, entry, laneCount, exchange->deferred);
+      tableIndex = arith::AddIOp::create(context.rewriter, loc, base, lane);
+    }
+    Type fragmentType = run.slices.front().family->requirement->publicationFragmentType;
+    auto receive = SpatChannelReceiveOp::create(context.rewriter,
+                                                loc,
+                                                fragmentType,
+                                                lookup(channelTable, tableIndex, exchange->deferred, context, loc),
+                                                lookup(sourceTable, tableIndex, exchange->deferred, context, loc),
+                                                lookup(targetTable, tableIndex, exchange->deferred, context, loc));
+    setLogicalTransferMetadata(receive, metadata);
+    auto source = addLeadingUnitTensorDimension(context.rewriter, loc, receive.getOutput());
+    if (failed(source))
+      return failure();
+    auto sourceType = cast<RankedTensorType>(source->getType());
+    MixedSliceGeometry geometry;
+    geometry.offsets.assign(sourceType.getRank(), context.rewriter.getIndexAttr(0));
+    geometry.offsets.front() = lookupGeometry(positions, tableIndex, exchange->deferred, context, loc);
+    for (int64_t dimension : sourceType.getShape())
+      geometry.sizes.push_back(context.rewriter.getIndexAttr(dimension));
+    geometry.strides.assign(sourceType.getRank(), context.rewriter.getIndexAttr(1));
+    return insertMixedSlice(context.rewriter, loc, *source, current, geometry);
+  };
+  auto loop = buildNormalizedScfFor(
+    context.rewriter,
+    loc,
+    context.constants.getIndex(0),
+    context.constants.getIndex(entryCount),
+    context.constants.getIndex(1),
+    ValueRange {initial},
+    [&](OpBuilder&, Location, Value entry, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) -> LogicalResult {
+      auto value = emitEntry(entry, iterArgs.front());
+      if (failed(value))
+        return failure();
+      yielded.push_back(*value);
+      return success();
+    });
+  if (failed(loop))
+    return failure();
+  Value result = loop->results.front();
+  context.projectionAssemblies[{exchange, leafIndex}] = result;
+  return result;
+}
+
+static FailureOr<Value> emitAssemblyRun(const EmitReceiveAssemblyRun& run,
+                                        Value lane,
+                                        unsigned laneCount,
+                                        ArrayRef<DeferredResultPlan> results,
+                                        DeferredEmissionContext& context) {
+  DeferredExchangePlan* exchange = run.slices.front().family->requirement->exchange;
+  const DeferredResultPlan* resultPlan = findResultPlan(results, exchange);
+  if (!resultPlan)
+    return failure();
+  auto& assembly = exchange->program.insertAssembly;
+  size_t entryCount = run.entryOffsets.empty() ? 0 : run.entryOffsets.size() - 1;
+  if (!assembly || assembly->entries.size() != entryCount)
+    return failure();
+  if (run.assemblyEntries.size() != entryCount)
+    return failure();
+  LogicalTransferMetadataView metadata = buildMetadataView(run.slices);
+  SmallVector<int64_t> channelTable(
+    entryCount * laneCount, metadata.channels.valueAt(0));
+  SmallVector<int64_t> sourceTable(
+    entryCount * laneCount, metadata.sourceCores.valueAt(0));
+  SmallVector<int64_t> targetTable(
+    entryCount * laneCount, metadata.targetCores.valueAt(0));
+  unsigned rank = assembly->resultType.getRank();
+  SmallVector<SmallVector<int64_t>> geometryOffsets(rank, SmallVector<int64_t>(entryCount * laneCount));
+  SmallVector<SmallVector<int64_t>> geometrySizes(rank, SmallVector<int64_t>(entryCount * laneCount, 1));
+  SmallVector<SmallVector<int64_t>> geometryStrides(rank, SmallVector<int64_t>(entryCount * laneCount, 1));
+  for (size_t entryIndex = 0; entryIndex < entryCount; ++entryIndex) {
+    unsigned assemblyEntry = run.assemblyEntries[entryIndex];
+    if (assemblyEntry >= resultPlan->assemblyGeometry.size())
+      return failure();
+    ArrayRef<ScheduledTransferSlice> entrySlices =
+      ArrayRef(run.slices)
+        .slice(run.entryOffsets[entryIndex], run.entryOffsets[entryIndex + 1] - run.entryOffsets[entryIndex]);
+    LogicalTransferMetadataView entry = buildMetadataView(entrySlices);
+    for (size_t index = 0; index < entry.size(); ++index) {
+      size_t tableIndex =
+        entryIndex * laneCount + entry.targetLanes.valueAt(index);
+      channelTable[tableIndex] = entry.channels.valueAt(index);
+      sourceTable[tableIndex] = entry.sourceCores.valueAt(index);
+      targetTable[tableIndex] = entry.targetCores.valueAt(index);
+      unsigned targetLane = entry.targetLanes.valueAt(index);
+      for (unsigned dimension = 0; dimension < rank; ++dimension) {
+        geometryOffsets[dimension][tableIndex] =
+          resultPlan->assemblyGeometry[assemblyEntry].offsets[dimension].valueAt(targetLane);
+        geometrySizes[dimension][tableIndex] =
+          resultPlan->assemblyGeometry[assemblyEntry].sizes[dimension].valueAt(targetLane);
+        geometryStrides[dimension][tableIndex] =
+          resultPlan->assemblyGeometry[assemblyEntry].strides[dimension].valueAt(targetLane);
+      }
+    }
+  }
+  Location loc = exchange->deferred.getLoc();
+  Operation* initial = context.rewriter.clone(*assembly->initialValue);
+  Value assembled = initial->getResult(0);
+  auto emitEntry = [&](Value entryIndex, unsigned staticEntry, Value current) -> FailureOr<Value> {
+    Value tableIndex = entryIndex;
+    if (lane) {
+      Value base = affineMulConst(context.rewriter, loc, entryIndex, laneCount, exchange->deferred);
+      AffineExpr d0 = context.rewriter.getAffineDimExpr(0);
+      AffineExpr d1 = context.rewriter.getAffineDimExpr(1);
+      tableIndex = createOrFoldAffineApply(context.rewriter, loc, d0 + d1, ValueRange {base, lane}, exchange->deferred);
+    }
+    Type fragmentType = run.slices[run.entryOffsets[staticEntry]].family->requirement->publicationFragmentType;
+    auto receive = SpatChannelReceiveOp::create(context.rewriter,
+                                                loc,
+                                                fragmentType,
+                                                lookup(channelTable, tableIndex, exchange->deferred, context, loc),
+                                                lookup(sourceTable, tableIndex, exchange->deferred, context, loc),
+                                                lookup(targetTable, tableIndex, exchange->deferred, context, loc));
+    setLogicalTransferMetadata(receive, metadata);
+    auto source =
+      applyAssemblyTransform(receive.getOutput(), assembly->entries[run.assemblyEntries[staticEntry]], context, loc);
+    if (failed(source))
+      return failure();
+    MixedSliceGeometry geometry;
+    for (unsigned dimension = 0; dimension < rank; ++dimension) {
+      geometry.offsets.push_back(
+        lookupGeometry(geometryOffsets[dimension], tableIndex, exchange->deferred, context, loc));
+      geometry.sizes.push_back(lookupGeometry(geometrySizes[dimension], tableIndex, exchange->deferred, context, loc));
+      geometry.strides.push_back(
+        lookupGeometry(geometryStrides[dimension], tableIndex, exchange->deferred, context, loc));
+    }
+    return insertMixedSlice(context.rewriter, loc, *source, current, geometry);
+  };
+  if (entryCount == 1) {
+    auto value = emitEntry(context.constants.getIndex(0), 0, assembled);
+    if (failed(value))
+      return failure();
+    context.assemblies[exchange] = *value;
+    return *value;
+  }
+  const DeferredInsertAssemblyEntryTemplate& firstEntry = assembly->entries[run.assemblyEntries.front()];
+  bool sameTransform = llvm::all_of(ArrayRef(run.assemblyEntries).drop_front(), [&](unsigned entryIndex) {
+    const DeferredInsertAssemblyEntryTemplate& entry = assembly->entries[entryIndex];
+    return entry.sourceTransform == firstEntry.sourceTransform && entry.sourceType == firstEntry.sourceType;
+  });
+  if (!sameTransform)
+    return failure();
+  auto loop = buildNormalizedScfFor(
+    context.rewriter,
+    loc,
+    context.constants.getIndex(0),
+    context.constants.getIndex(entryCount),
+    context.constants.getIndex(1),
+    ValueRange {assembled},
+    [&](OpBuilder&, Location, Value index, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) -> LogicalResult {
+      auto value = emitEntry(index, 0, iterArgs.front());
+      if (failed(value))
+        return failure();
+      yielded.push_back(*value);
+      return success();
+    });
+  if (failed(loop))
+    return failure();
+  context.assemblies[exchange] = loop->results.front();
+  return loop->results.front();
+}
+
+static RequirementFamily* getAvailabilityRequirement(const AvailabilityAlternative& alternative) {
+  return std::visit(
+    [](const auto& value) -> RequirementFamily* {
+      using Alternative = std::decay_t<decltype(value)>;
+      if constexpr (std::is_same_v<Alternative, EmitReceiveRun>)
+        return value.slices.front().family->requirement;
+      else
+        return value.families.front()->requirement;
+    },
+    alternative.source);
+}
+
+static FailureOr<Value> emitAvailabilityAlternative(const AvailabilityAlternative& alternative,
+                                                    Value lane,
+                                                    unsigned laneCount,
+                                                    DeferredEmissionContext& context) {
+  if (auto receive = std::get_if<EmitReceiveRun>(&alternative.source))
+    return emitReceiveValue(*receive, lane, laneCount, context);
+  return materializeLocalValue(
+    std::get<MaterializeLocalFamily>(alternative.source), lane, laneCount,
+    context);
+}
+
+static void recordAvailability(const ResolveAvailability& availability,
+                               Value value,
+                               DeferredEmissionContext& context) {
+  for (const AvailabilityAlternative& alternative : availability.alternatives) {
+    if (auto receive = std::get_if<EmitReceiveRun>(&alternative.source)) {
+      for (const ScheduledTransferSlice& slice : receive->slices)
+        context.receives[slice.family->requirement] = value;
+      continue;
+    }
+    for (LocalAvailabilityFamily* family :
+         std::get<MaterializeLocalFamily>(alternative.source).families)
+      context.receives[family->requirement] = value;
+  }
+}
+
+static FailureOr<SmallVector<Value>> emitSelection(
+  Value branch, ArrayRef<Type> resultTypes, size_t branchCount, Location loc,
+  DeferredEmissionContext& context,
+  llvm::function_ref<LogicalResult(Region&, size_t)> emitRegion) {
+  SmallVector<Value> results;
+  if (branchCount == 2) {
+    Value condition = arith::CmpIOp::create(context.rewriter, loc,
+      arith::CmpIPredicate::eq, branch, context.constants.getIndex(0));
+    auto selection = scf::IfOp::create(
+      context.rewriter, loc, resultTypes, condition, true);
+    if (failed(emitRegion(selection.getThenRegion(), 0))
+        || failed(emitRegion(selection.getElseRegion(), 1)))
+      return failure();
+    llvm::append_range(results, selection.getResults());
+    return results;
+  }
+  SmallVector<int64_t> cases;
+  for (size_t index = 0; index + 1 < branchCount; ++index)
+    cases.push_back(index);
+  auto selection = scf::IndexSwitchOp::create(
+    context.rewriter, loc, resultTypes, branch, cases, cases.size());
+  for (auto [index, region] : llvm::enumerate(selection.getCaseRegions()))
+    if (failed(emitRegion(region, index)))
+      return failure();
+  if (failed(emitRegion(selection.getDefaultRegion(), branchCount - 1)))
+    return failure();
+  llvm::append_range(results, selection.getResults());
+  return results;
+}
+
+static LogicalResult emitAvailability(const ResolveAvailability& availability,
+                                      Value lane,
+                                      unsigned laneCount,
+                                      DeferredEmissionContext& context) {
+  if (availability.alternatives.empty())
+    return failure();
+  if (availability.alternatives.size() == 1) {
+    auto value = emitAvailabilityAlternative(availability.alternatives.front(), lane, laneCount, context);
+    if (failed(value))
+      return failure();
+    recordAvailability(availability, *value, context);
+    return success();
+  }
+  if (!lane)
+    return failure();
+  const AvailabilityAlternative& first = availability.alternatives.front();
+  RequirementFamily* requirement = getAvailabilityRequirement(first);
+  Location loc = requirement->exchange->deferred.getLoc();
+  SmallVector<Type> resultTypes {requirement->publicationFragmentType};
+  auto emitRegion = [&](Region& region, const AvailabilityAlternative& alternative) {
+    OpBuilder::InsertionGuard guard(context.rewriter);
+    Block& block = region.empty() ? *context.rewriter.createBlock(&region)
+                                  : region.front();
+    auto yield = block.empty() ? scf::YieldOp() : dyn_cast<scf::YieldOp>(&block.back());
+    if (yield)
+      context.rewriter.setInsertionPoint(yield);
+    else
+      context.rewriter.setInsertionPointToEnd(&block);
+    auto value = emitAvailabilityAlternative(alternative, lane, laneCount, context);
+    if (failed(value))
+      return failure();
+    if (yield)
+      yield->setOperands(ValueRange {*value});
+    else
+      scf::YieldOp::create(context.rewriter, loc, ValueRange {*value});
+    return success();
+  };
+  SmallVector<int64_t> branchByLane(laneCount);
+  SmallVector<bool> assigned(laneCount);
+  for (auto [branch, alternative] : llvm::enumerate(availability.alternatives))
+    for (LaneInterval interval : alternative.lanes.intervals())
+      for (unsigned activeLane = interval.begin; activeLane < interval.end;
+           ++activeLane) {
+        if (assigned[activeLane])
+          return failure();
+        assigned[activeLane] = true;
+        branchByLane[activeLane] = branch;
+      }
+  Value branch = emitStaticIntLookup(StaticIntSequence::fromValues(branchByLane),
+                                     lane,
+                                     requirement->exchange->deferred,
+                                     context.constants,
+                                     context.rewriter,
+                                     loc);
+  auto selected = emitSelection(branch, resultTypes,
+    availability.alternatives.size(), loc, context,
+    [&](Region& region, size_t index) {
+      return emitRegion(region, availability.alternatives[index]);
+    });
+  if (failed(selected))
+    return failure();
+  Value value = selected->front();
+  recordAvailability(availability, value, context);
+  return success();
+}
+
+static FailureOr<SmallVector<Value>> emitInstructions(
+  const BoundaryInstructionList& list, Value lane, unsigned laneCount,
+  ArrayRef<DeferredResultPlan> results, DeferredEmissionContext& context);
+
+static FailureOr<SmallVector<Value>> emitDispatch(
+  const LaneDispatch& dispatch, Value lane, unsigned laneCount,
+  ArrayRef<DeferredResultPlan> results, DeferredEmissionContext& context) {
+  if (dispatch.branches.empty())
+    return failure();
+  if (dispatch.branches.size() == 1)
+    return emitInstructions(
+      dispatch.branches.front(), lane, laneCount, results, context);
+  if (!lane)
+    return failure();
+
+  SmallVector<Type> resultTypes;
+  for (DeferredExchangePlan* exchange : dispatch.producedExchanges)
+    resultTypes.push_back(exchange->deferred.getOutput().getType());
+  Location loc = dispatch.producedExchanges.empty()
+                   ? lane.getLoc()
+                   : dispatch.producedExchanges.front()->deferred.getLoc();
+  Value branch = emitStaticIntLookup(dispatch.branchByLane,
+                                     lane,
+                                     dispatch.producedExchanges.empty()
+                                       ? lane.getDefiningOp()
+                                       : dispatch.producedExchanges.front()->deferred,
+                                     context.constants,
+                                     context.rewriter,
+                                     loc);
+  auto receives = context.receives;
+  auto assemblies = context.assemblies;
+  auto projectionAssemblies = context.projectionAssemblies;
+  auto emitRegion = [&](Region& region,
+                        const BoundaryInstructionList& instructions) {
+    context.receives = receives;
+    context.assemblies = assemblies;
+    context.projectionAssemblies = projectionAssemblies;
+    OpBuilder::InsertionGuard guard(context.rewriter);
+    Block& block = region.empty() ? *context.rewriter.createBlock(&region)
+                                  : region.front();
+    auto yield = block.empty() ? scf::YieldOp()
+                               : dyn_cast<scf::YieldOp>(&block.back());
+    if (yield)
+      context.rewriter.setInsertionPoint(yield);
+    else
+      context.rewriter.setInsertionPointToEnd(&block);
+    auto values = emitInstructions(
+      instructions, lane, laneCount, results, context);
+    if (failed(values))
+      return failure();
+    if (yield)
+      yield->setOperands(*values);
+    else
+      scf::YieldOp::create(context.rewriter, loc, *values);
+    return success();
+  };
+
+  auto produced = emitSelection(branch, resultTypes, dispatch.branches.size(),
+    loc, context, [&](Region& region, size_t index) {
+      return emitRegion(region, dispatch.branches[index]);
+    });
+  if (failed(produced))
+    return failure();
+  context.receives = std::move(receives);
+  context.assemblies = std::move(assemblies);
+  context.projectionAssemblies = std::move(projectionAssemblies);
+  return *produced;
+}
+
+static FailureOr<SmallVector<Value>> emitInstructions(
+  const BoundaryInstructionList& list, Value lane, unsigned laneCount,
+  ArrayRef<DeferredResultPlan> results, DeferredEmissionContext& context) {
+  SmallVector<Value> produced;
+  for (const BoundaryInstruction& instruction : list.instructions) {
+    if (auto send = std::get_if<EmitSendRun>(&instruction)) {
+      if (failed(emitConditionalSendRun(*send, lane, laneCount, context)))
+        return failure();
+      continue;
+    }
+    if (auto availability = std::get_if<ResolveAvailability>(&instruction)) {
+      if (failed(emitAvailability(*availability, lane, laneCount, context)))
+        return failure();
+      continue;
+    }
+    if (auto assembly = std::get_if<EmitReceiveAssemblyRun>(&instruction)) {
+      auto value = assembly->projectionLeaf
+                     ? emitProjectionAssemblyRun(
+                         *assembly, lane, laneCount, context)
+                     : emitAssemblyRun(
+                         *assembly, lane, laneCount, results, context);
+      if (failed(value))
+        return failure();
+      continue;
+    }
+    if (auto result = std::get_if<ProduceDeferredResult>(&instruction)) {
+      const DeferredResultPlan* plan = findResultPlan(results, result->exchange);
+      if (!plan)
+        return failure();
+      auto value = realizeDeferredResult(*plan, result->lanes, lane, context);
+      if (failed(value))
+        return failure();
+      produced.push_back(*value);
+      continue;
+    }
+    auto values = emitDispatch(
+      **std::get_if<std::unique_ptr<LaneDispatch>>(&instruction), lane,
+      laneCount, results, context);
+    if (failed(values))
+      return failure();
+    llvm::append_range(produced, *values);
+  }
+  return produced;
+}
+
+static SmallVector<DeferredExchangePlan*> getProducedExchanges(
+  const BoundaryInstructionList& list) {
+  SmallVector<DeferredExchangePlan*> exchanges;
+  for (const BoundaryInstruction& instruction : list.instructions) {
+    if (auto result = std::get_if<ProduceDeferredResult>(&instruction))
+      exchanges.push_back(result->exchange);
+    else if (auto dispatch =
+               std::get_if<std::unique_ptr<LaneDispatch>>(&instruction))
+      llvm::append_range(exchanges, (*dispatch)->producedExchanges);
+  }
+  return exchanges;
+}
+
+static void setInsertionAtBoundary(IRRewriter& rewriter, const BoundaryKey& key) {
+  ScheduledInfo& scheduled = *key.scheduled;
+  if (key.insertionStep < scheduled.stepAnchors.size() && scheduled.stepAnchors[key.insertionStep]->getBlock()) {
+    rewriter.setInsertionPoint(scheduled.stepAnchors[key.insertionStep]);
+    return;
+  }
+  rewriter.setInsertionPoint(scheduled.blocks.front()->getTerminator());
+}
+
+static LogicalResult
+replaceResults(ArrayRef<DeferredExchangePlan*> exchanges, ValueRange replacements, DeferredEraseSet& erase) {
+  if (exchanges.size() != replacements.size())
+    return failure();
+  for (auto [exchange, replacement] : llvm::zip_equal(exchanges, replacements)) {
+    exchange->deferred.getOutput().replaceAllUsesWith(replacement);
+    erase.insert(exchange->deferred);
+  }
+  return success();
+}
+
+static LogicalResult emitBoundary(const BoundaryProgram& boundary,
+                                  ArrayRef<DeferredResultPlan> results,
+                                  DeferredEmissionContext& context,
+                                  DeferredEraseSet& erase) {
+  setInsertionAtBoundary(context.rewriter, boundary.key);
+  unsigned laneCount = boundary.key.scheduled->cores.size();
+  Value lane;
+  if (auto batch = dyn_cast<SpatScheduledComputeBatch>(boundary.key.scheduled->op))
+    lane = *batch.getLaneArgument();
+  SmallVector<DeferredExchangePlan*> exchanges =
+    getProducedExchanges(boundary.root);
+  auto values = emitInstructions(
+    boundary.root, lane, laneCount, results, context);
+  return failed(values) ? failure()
+                        : replaceResults(exchanges, *values, erase);
+}
+
+} // namespace
+
+LogicalResult realizeDeferredBoundaries(ArrayRef<BoundaryProgram> boundaries,
+                                        ArrayRef<DeferredResultPlan> results,
+                                        DeferredEmissionContext& context,
+                                        DeferredEraseSet& erase) {
+  ScheduledInfo* scheduled = nullptr;
+  for (const BoundaryProgram& boundary : boundaries) {
+    if (scheduled != boundary.key.scheduled) {
+      context.receives.clear();
+      context.assemblies.clear();
+      context.projectionAssemblies.clear();
+      scheduled = boundary.key.scheduled;
+    }
+    if (failed(emitBoundary(boundary, results, context, erase)))
+      return boundary.key.scheduled->op->emitOpError("phase 2 failed to realize a communication boundary");
+  }
+  return success();
+}
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,30 @@
+#pragma once
+
+#include "mlir/IR/PatternMatch.h"
+
+#include "DeferredBoundaryPlanning.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+
+namespace onnx_mlir::spatial {
+
+struct DeferredEmissionContext {
+  DeferredEmissionContext(mlir::IRRewriter& rewriter,
+                          ConstantPool& constants)
+    : rewriter(rewriter), constants(constants) {}
+
+  mlir::IRRewriter& rewriter;
+  ConstantPool& constants;
+  llvm::DenseMap<RequirementFamily*, mlir::Value> receives;
+  llvm::DenseMap<DeferredExchangePlan*, mlir::Value> assemblies;
+  llvm::DenseMap<std::pair<DeferredExchangePlan*, unsigned>, mlir::Value>
+    projectionAssemblies;
+};
+
+using DeferredEraseSet = llvm::SetVector<mlir::Operation*>;
+
+mlir::LogicalResult realizeDeferredBoundaries(mlir::ArrayRef<BoundaryProgram> boundaries,
+                                              mlir::ArrayRef<DeferredResultPlan> results,
+                                              DeferredEmissionContext& context,
+                                              DeferredEraseSet& erase);
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,392 @@
+#include "DeferredCommunicationDeadlock.hpp"
+
+#include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+#include "llvm/ADT/DenseMap.h"
+
+#include <map>
+
+namespace onnx_mlir::spatial {
+using namespace mlir;
+namespace {
+
+enum class EventKind { Compute, Send, Receive };
+
+struct Event {
+  EventKind kind = EventKind::Compute;
+  uint64_t channel = 0;
+};
+
+struct PlannedStreamCursor {
+  unsigned step = 0;
+  size_t slice = 0;
+  size_t offset = 0;
+};
+
+static std::optional<Event> getPlannedHead(
+    unsigned stream, PlannedStreamCursor &cursor, unsigned stepCount,
+    const ScheduledCommunicationPlan &plan) {
+  while (cursor.slice < plan.slices.size()) {
+    const ScheduledTransferSlice &slice = plan.slices[cursor.slice];
+    ExternalTransferFamily &family = *slice.family;
+    size_t begin = slice.familyOffset + cursor.offset;
+    size_t length = slice.transferCount - cursor.offset;
+    auto source = family.sourceStreams.find(stream, begin, length);
+    auto target = family.targetStreams.find(stream, begin, length);
+    std::optional<size_t> index = source;
+    EventKind kind = EventKind::Send;
+    unsigned insertionStep = slice.sourceInsertionStep;
+    if (target && (!index || *target < *index)) {
+      index = *target;
+      kind = EventKind::Receive;
+      insertionStep = slice.targetInsertionStep;
+    }
+    if (!index) {
+      ++cursor.slice;
+      cursor.offset = 0;
+      continue;
+    }
+    cursor.offset = *index - slice.familyOffset;
+    if (cursor.step < insertionStep)
+      return Event {EventKind::Compute, 0};
+    return Event {kind, static_cast<uint64_t>(
+      family.channelIds.valueAt(*index))};
+  }
+  if (cursor.step < stepCount)
+    return Event {EventKind::Compute, 0};
+  return std::nullopt;
+}
+
+static LogicalResult simulatePlanned(
+    Operation *anchor, ArrayRef<unsigned> stepCounts,
+    const ScheduledCommunicationPlan &plan) {
+  SmallVector<PlannedStreamCursor> cursors(stepCounts.size());
+  DenseMap<uint64_t, unsigned> sends, receives;
+  SmallVector<unsigned> readyComputes;
+  SmallVector<uint64_t> readyChannels;
+  size_t computeCursor = 0, channelCursor = 0;
+  unsigned finished = 0;
+  auto registerHead = [&](unsigned stream) {
+    auto head = getPlannedHead(
+      stream, cursors[stream], stepCounts[stream], plan);
+    if (!head) {
+      ++finished;
+      return;
+    }
+    if (head->kind == EventKind::Compute) {
+      readyComputes.push_back(stream);
+      return;
+    }
+    auto &own = head->kind == EventKind::Send ? sends : receives;
+    auto &peer = head->kind == EventKind::Send ? receives : sends;
+    own[head->channel] = stream;
+    if (peer.contains(head->channel))
+      readyChannels.push_back(head->channel);
+  };
+  for (unsigned stream = 0; stream < stepCounts.size(); ++stream)
+    registerHead(stream);
+  while (computeCursor != readyComputes.size()
+         || channelCursor != readyChannels.size()) {
+    if (computeCursor != readyComputes.size()) {
+      unsigned stream = readyComputes[computeCursor++];
+      ++cursors[stream].step;
+      registerHead(stream);
+      continue;
+    }
+    uint64_t channel = readyChannels[channelCursor++];
+    auto send = sends.find(channel);
+    auto receive = receives.find(channel);
+    if (send == sends.end() || receive == receives.end())
+      continue;
+    unsigned source = send->second;
+    unsigned target = receive->second;
+    sends.erase(send);
+    receives.erase(receive);
+    ++cursors[source].offset;
+    ++cursors[target].offset;
+    registerHead(source);
+    registerHead(target);
+  }
+  if (finished == stepCounts.size())
+    return success();
+  InFlightDiagnostic diagnostic = anchor->emitError(
+    "planned communication rendezvous simulation made no progress");
+  unsigned reported = 0;
+  for (unsigned stream = 0;
+       stream < stepCounts.size() && reported < 8; ++stream) {
+    auto head = getPlannedHead(
+      stream, cursors[stream], stepCounts[stream], plan);
+    if (!head)
+      continue;
+    diagnostic << (reported++ == 0 ? "; blocked " : ", ")
+               << "stream " << stream << " at channel " << head->channel;
+  }
+  return failure();
+}
+
+struct RealizedOperation {
+  Operation *op = nullptr;
+  bool send = false;
+  StaticIntSequence channels;
+  StaticIntSequence parents;
+  StaticIntSequence counts;
+  StaticIntSequence sources;
+  StaticIntSequence targets;
+};
+
+static FailureOr<size_t> getBatchTransferCount(Operation *op) {
+  if (auto count = op->getAttrOfType<IntegerAttr>(
+        "raptor.batch_transfer_count")) {
+    if (count.getInt() > 0)
+      return count.getInt();
+    return op->emitOpError("has invalid compact transfer count"), failure();
+  }
+  if (op->hasAttr("raptor.batch_channel_ids_encoding"))
+    return op->emitOpError("is missing compact transfer count"), failure();
+  Attribute channels = op->getAttr("raptor.batch_channel_ids");
+  if (auto array = dyn_cast_or_null<DenseI64ArrayAttr>(channels))
+    return array.empty()
+      ? FailureOr<size_t>(failure())
+      : FailureOr<size_t>(array.size());
+  if (auto elements = dyn_cast_or_null<DenseIntElementsAttr>(channels);
+      elements && elements.getNumElements() > 0)
+    return elements.getNumElements();
+  return op->emitOpError("has invalid legacy compact transfer metadata"),
+         failure();
+}
+
+static FailureOr<RealizedOperation> parseRealizedOperation(Operation *op) {
+  bool scalar = op->hasAttr("raptor.channel_id");
+  bool batch = op->hasAttr("raptor.batch_channel_ids");
+  if (scalar == batch) {
+    op->emitOpError(
+      "must have exactly one scalar or compact metadata form");
+    return failure();
+  }
+  auto batchCount = scalar ? FailureOr<size_t>(1)
+                           : getBatchTransferCount(op);
+  if (failed(batchCount))
+    return failure();
+  size_t size = *batchCount;
+  auto channels = getStaticIntSequenceAttr(
+    op, scalar ? "raptor.channel_id" : "raptor.batch_channel_ids", size);
+  auto parents = getStaticIntSequenceAttr(
+    op, scalar ? "raptor.parent_exchange_id"
+               : "raptor.batch_parent_exchange_ids", size);
+  auto counts = getStaticIntSequenceAttr(
+    op, scalar ? "raptor.parent_transfer_count"
+               : "raptor.batch_parent_transfer_counts", size);
+  auto sources = getStaticIntSequenceAttr(
+    op, scalar ? "raptor.source_core" : "raptor.batch_source_cores", size);
+  auto targets = getStaticIntSequenceAttr(
+    op, scalar ? "raptor.target_core" : "raptor.batch_target_cores", size);
+  if (failed(channels) || failed(parents) || failed(counts)
+      || failed(sources) || failed(targets))
+    return failure();
+  if (scalar) {
+    auto exchange = op->getAttrOfType<IntegerAttr>("raptor.exchange_id");
+    if (!exchange || exchange.getInt() != channels->valueAt(0)) {
+      op->emitOpError("has inconsistent scalar exchange metadata");
+      return failure();
+    }
+  }
+  return RealizedOperation {op, isa<SpatChannelSendOp>(op),
+    *channels, *parents, *counts, *sources, *targets};
+}
+
+struct CoreTransferSequences {
+  DenseMap<int64_t, StaticIntSequenceChain> sends;
+  DenseMap<int64_t, StaticIntSequenceChain> receives;
+};
+
+struct ExpectedFamily {
+  ExternalTransferFamily *family = nullptr;
+  int64_t firstChannel = 0;
+  int64_t endChannel = 0;
+};
+
+static void appendByCore(DenseMap<int64_t, StaticIntSequenceChain> &result,
+                         const StaticIntSequence &channels,
+                         const StaticIntSequence &cores, size_t begin,
+                         size_t count) {
+  size_t end = begin + count;
+  cores.forEachEqualRun(
+    [&](int64_t core, size_t runBegin, size_t runCount) {
+      size_t selectedBegin = std::max(begin, runBegin);
+      size_t selectedEnd = std::min(end, runBegin + runCount);
+      if (selectedBegin < selectedEnd)
+        result[core].append(
+          channels, selectedBegin, selectedEnd - selectedBegin);
+    });
+}
+
+static LogicalResult compareSequences(
+    func::FuncOp funcOp,
+    const DenseMap<int64_t, StaticIntSequenceChain> &expected,
+    const DenseMap<int64_t, StaticIntSequenceChain> &actual, StringRef kind) {
+  if (expected.size() != actual.size())
+    return funcOp.emitOpError()
+           << "realized " << kind << " stream set differs from plan";
+  for (const auto &[core, sequence] : expected) {
+    auto found = actual.find(core);
+    if (found == actual.end())
+      return funcOp.emitOpError() << "realized " << kind
+             << " stream is missing on core " << core;
+    StaticIntSequenceChainCursor expectedCursor(sequence);
+    StaticIntSequenceChainCursor actualCursor(found->second);
+    uint64_t ordinal = 0;
+    while (!expectedCursor.done() && !actualCursor.done()
+           && expectedCursor.value() == actualCursor.value()) {
+      expectedCursor.advance();
+      actualCursor.advance();
+      ++ordinal;
+    }
+    if (expectedCursor.done() && actualCursor.done())
+      continue;
+    return funcOp.emitOpError()
+           << "realized " << kind << " logical order differs on core "
+           << core << " at ordinal " << ordinal << ": expected channel "
+           << (expectedCursor.done() ? -1 : expectedCursor.value())
+           << ", actual channel "
+           << (actualCursor.done() ? -1 : actualCursor.value());
+  }
+  return success();
+}
+
+} // namespace
+
+LogicalResult verifyPlannedCommunicationDeadlockFree(
+    Operation *anchor, ArrayRef<unsigned> stepCounts,
+    const ScheduledCommunicationPlan &plan) {
+  for (const ScheduledTransferSlice &slice : plan.slices) {
+    ExternalTransferFamily &family = *slice.family;
+    for (size_t offset = 0; offset < slice.transferCount; ++offset) {
+      size_t index = slice.familyOffset + offset;
+      unsigned source = family.sourceStreams.valueAt(index);
+      unsigned target = family.targetStreams.valueAt(index);
+      if (source >= stepCounts.size() || target >= stepCounts.size()
+          || slice.sourceInsertionStep > stepCounts[source]
+          || slice.targetInsertionStep > stepCounts[target]
+          || slice.sourceInsertionStep <= family.requirement->producer->step
+          || slice.targetInsertionStep
+               > family.requirement->exchange->consumerStep)
+        return anchor->emitError(
+          "communication plan references an invalid stream step");
+    }
+  }
+  return simulatePlanned(anchor, stepCounts, plan);
+}
+
+LogicalResult verifyRealizedCommunicationDeadlockFree(
+    func::FuncOp funcOp, const ScheduledCommunicationPlan &plan) {
+  SmallVector<ExpectedFamily> families;
+  DenseMap<ExternalTransferFamily *, unsigned> familyIndex;
+  for (const ScheduledTransferSlice &slice : plan.slices) {
+    ExternalTransferFamily *family = slice.family;
+    if (familyIndex.count(family))
+      continue;
+    size_t count = family->channelIds.size();
+    int64_t first = family->channelIds.valueAt(0);
+    for (size_t index = 1; index < count; ++index)
+      if (family->channelIds.valueAt(index)
+          != first + static_cast<int64_t>(index))
+        return funcOp.emitOpError(
+          "planned communication family has non-consecutive channels");
+    familyIndex[family] = families.size();
+    families.push_back({family, first, first + static_cast<int64_t>(count)});
+  }
+  llvm::sort(families, [](const ExpectedFamily &lhs,
+                          const ExpectedFamily &rhs) {
+    return lhs.firstChannel < rhs.firstChannel;
+  });
+  familyIndex.clear();
+  std::map<int64_t, unsigned> familyByFirstChannel;
+  int64_t nextChannel = 0;
+  for (auto [index, expected] : llvm::enumerate(families)) {
+    if (expected.firstChannel != nextChannel)
+      return funcOp.emitOpError(
+        "planned communication channels are not exactly contiguous");
+    nextChannel = expected.endChannel;
+    familyIndex[expected.family] = index;
+    familyByFirstChannel.emplace(expected.firstChannel, index);
+  }
+  if (static_cast<uint64_t>(nextChannel) != plan.logicalTransferCount)
+    return funcOp.emitOpError(
+      "planned communication channel count is inconsistent");
+
+  CoreTransferSequences expected;
+  for (const ScheduledTransferSlice &slice : plan.slices) {
+    ExternalTransferFamily &family = *slice.family;
+    appendByCore(expected.sends, family.channelIds, family.sourceCores,
+                 slice.familyOffset, slice.transferCount);
+    appendByCore(expected.receives, family.channelIds, family.targetCores,
+                 slice.familyOffset, slice.transferCount);
+  }
+
+  CoreTransferSequences actual;
+  SmallVector<std::unique_ptr<StaticIntSequence>> actualChannels;
+  bool invalid = false;
+  funcOp.walk([&](Operation *op) {
+    if (invalid || !isa<SpatChannelSendOp, SpatChannelReceiveOp>(op))
+      return;
+    auto realized = parseRealizedOperation(op);
+    if (failed(realized)) {
+      invalid = true;
+      return;
+    }
+    Type payloadType = realized->send
+      ? cast<SpatChannelSendOp>(op).getInput().getType()
+      : cast<SpatChannelReceiveOp>(op).getOutput().getType();
+    for (size_t index = 0; index < realized->channels.size(); ++index) {
+      int64_t channel = realized->channels.valueAt(index);
+      auto upper = familyByFirstChannel.upper_bound(channel);
+      if (upper == familyByFirstChannel.begin()) {
+        op->emitOpError("references an unknown logical channel");
+        invalid = true;
+        return;
+      }
+      ExpectedFamily &expected = families[std::prev(upper)->second];
+      if (channel >= expected.endChannel) {
+        op->emitOpError("references an unknown logical channel");
+        invalid = true;
+        return;
+      }
+      ExternalTransferFamily &family = *expected.family;
+      size_t familyOffset = channel - expected.firstChannel;
+      RequirementFamily &requirement = *family.requirement;
+      if (realized->parents.valueAt(index)
+              != static_cast<int64_t>(requirement.exchange->exchangeId)
+          || realized->counts.valueAt(index)
+              != requirement.exchange->externalTransferCount
+          || realized->sources.valueAt(index)
+              != family.sourceCores.valueAt(familyOffset)
+          || realized->targets.valueAt(index)
+              != family.targetCores.valueAt(familyOffset)
+          || payloadType != requirement.publicationFragmentType) {
+        op->emitOpError(
+          "logical transfer metadata differs from its symbolic family");
+        invalid = true;
+        return;
+      }
+    }
+    if (invalid)
+      return;
+    actualChannels.push_back(
+      std::make_unique<StaticIntSequence>(std::move(realized->channels)));
+    appendByCore(realized->send ? actual.sends : actual.receives,
+                 *actualChannels.back(),
+                 realized->send ? realized->sources : realized->targets,
+                 0, actualChannels.back()->size());
+  });
+  if (invalid)
+    return failure();
+  if (failed(compareSequences(
+        funcOp, expected.sends, actual.sends, "send"))
+      || failed(compareSequences(
+        funcOp, expected.receives, actual.receives, "receive")))
+    return failure();
+  return success();
+}
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,18 @@
+#pragma once
+
+#include "DeferredCommunicationScheduling.hpp"
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+
+namespace onnx_mlir::spatial {
+
+mlir::LogicalResult verifyPlannedCommunicationDeadlockFree(
+  mlir::Operation *anchor,
+  mlir::ArrayRef<unsigned> stepCounts,
+  const ScheduledCommunicationPlan &plan);
+
+mlir::LogicalResult verifyRealizedCommunicationDeadlockFree(
+  mlir::func::FuncOp funcOp,
+  const ScheduledCommunicationPlan &plan);
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,236 @@
+#pragma once
+
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/Operation.h"
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/ADT/SmallVector.h"
+
+#include <memory>
+#include <optional>
+#include <variant>
+
+#include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+namespace onnx_mlir::spatial {
+
+struct LaneInterval {
+  unsigned begin = 0;
+  unsigned end = 0;
+
+  bool operator==(const LaneInterval& other) const { return begin == other.begin && end == other.end; }
+};
+
+class LaneSet {
+public:
+  static LaneSet all(unsigned laneCount) { return range(0, laneCount); }
+  static LaneSet range(unsigned begin, unsigned end) {
+    LaneSet lanes;
+    if (begin < end)
+      lanes.ranges.push_back({begin, end});
+    return lanes;
+  }
+
+  bool empty() const { return ranges.empty(); }
+  size_t size() const {
+    size_t result = 0;
+    for (LaneInterval range : ranges)
+      result += range.end - range.begin;
+    return result;
+  }
+  bool contains(unsigned lane) const {
+    return llvm::any_of(ranges, [&](LaneInterval range) { return range.begin <= lane && lane < range.end; });
+  }
+  llvm::ArrayRef<LaneInterval> intervals() const { return ranges; }
+
+  LaneSet intersect(const LaneSet& other) const {
+    LaneSet result;
+    for (LaneInterval lhs : ranges)
+      for (LaneInterval rhs : other.ranges) {
+        unsigned begin = std::max(lhs.begin, rhs.begin);
+        unsigned end = std::min(lhs.end, rhs.end);
+        if (begin < end)
+          result.append(begin, end);
+      }
+    return result;
+  }
+
+  LaneSet subtract(const LaneSet& other) const {
+    LaneSet result;
+    for (LaneInterval source : ranges) {
+      unsigned cursor = source.begin;
+      for (LaneInterval removed : other.ranges) {
+        if (removed.end <= cursor || removed.begin >= source.end)
+          continue;
+        if (cursor < removed.begin)
+          result.append(cursor, std::min(removed.begin, source.end));
+        cursor = std::max(cursor, removed.end);
+        if (cursor >= source.end)
+          break;
+      }
+      if (cursor < source.end)
+        result.append(cursor, source.end);
+    }
+    return result;
+  }
+
+  LaneSet unite(const LaneSet& other) const {
+    llvm::SmallVector<LaneInterval, 4> combined(ranges.begin(), ranges.end());
+    llvm::append_range(combined, other.ranges);
+    llvm::sort(combined, [](LaneInterval lhs, LaneInterval rhs) { return lhs.begin < rhs.begin; });
+    LaneSet normalized;
+    for (LaneInterval range : combined)
+      normalized.append(range.begin, range.end);
+    return normalized;
+  }
+
+  bool operator==(const LaneSet& other) const { return ranges == other.ranges; }
+
+private:
+  void append(unsigned begin, unsigned end) {
+    if (begin >= end)
+      return;
+    if (!ranges.empty() && begin <= ranges.back().end) {
+      ranges.back().end = std::max(ranges.back().end, end);
+      return;
+    }
+    ranges.push_back({begin, end});
+  }
+
+  llvm::SmallVector<LaneInterval, 2> ranges;
+};
+
+struct RequirementCoordinate {
+  unsigned leafIndex = 0;
+  unsigned selectedPosition = 0;
+
+  bool operator==(const RequirementCoordinate& other) const {
+    return leafIndex == other.leafIndex && selectedPosition == other.selectedPosition;
+  }
+};
+
+enum class DeferredLeafForm {
+  DirectSource,
+  GraphBatchProjection
+};
+enum class DeferredAssemblySourceTransform {
+  Identity,
+  AddLeadingUnitDimension,
+  RemoveLeadingUnitDimension
+};
+
+struct DeferredSliceTemplate {
+  llvm::SmallVector<mlir::OpFoldResult> offsets;
+  llvm::SmallVector<mlir::OpFoldResult> sizes;
+  llvm::SmallVector<mlir::OpFoldResult> strides;
+};
+
+struct DeferredProjectionLeafTemplate {
+  DeferredLeafForm form = DeferredLeafForm::DirectSource;
+  mlir::Value sourceRoot;
+  mlir::Value replacementRoot;
+  mlir::tensor::ExtractSliceOp leadingProjection;
+  DeferredSliceTemplate leadingGeometry;
+  DeferredSliceTemplate innerGeometry;
+  mlir::RankedTensorType reconstructedType;
+};
+
+struct DeferredInsertAssemblyEntryTemplate {
+  RequirementCoordinate coordinate;
+  DeferredAssemblySourceTransform sourceTransform = DeferredAssemblySourceTransform::Identity;
+  mlir::RankedTensorType sourceType;
+  DeferredSliceTemplate targetGeometry;
+};
+
+struct DeferredInsertAssemblyTemplate {
+  mlir::tensor::EmptyOp initialValue;
+  mlir::RankedTensorType resultType;
+  llvm::SmallVector<DeferredInsertAssemblyEntryTemplate> entries;
+};
+
+struct DeferredProgramTemplate {
+  SpatDeferredCommunicationOp deferred;
+  mlir::Value scheduledLane;
+  mlir::Value yieldedValue;
+  llvm::SmallVector<DeferredProjectionLeafTemplate, 0> leaves;
+  llvm::SmallVector<mlir::Operation*> residualOps;
+  std::optional<DeferredInsertAssemblyTemplate> insertAssembly;
+};
+
+struct ScheduledInfo;
+
+struct ProducedValue {
+  ScheduledInfo* scheduled = nullptr;
+  unsigned step = 0;
+  unsigned resultIndex = 0;
+  int64_t graphId = -1;
+  int64_t core = -1;
+  int64_t laneStart = 0;
+  int64_t laneCount = 1;
+  unsigned scheduledLane = 0;
+  int64_t publishedSlotStart = 0;
+  int64_t publishedSlotCount = 1;
+  mlir::Value payload;
+  mlir::Value published;
+};
+
+struct ScheduledInfo {
+  mlir::Operation* op = nullptr;
+  llvm::SmallVector<mlir::Block*> blocks;
+  llvm::SmallVector<mlir::Operation*> stepAnchors;
+  llvm::SmallVector<int64_t> cores;
+  llvm::SmallVector<int64_t> stepSourceIds;
+  llvm::SmallVector<int64_t> resultOffsets;
+  llvm::SmallVector<int64_t> resultCounts;
+  llvm::SmallVector<ProducedValue*> produced;
+  llvm::SmallVector<unsigned> streamIds;
+
+  bool isBatch() const { return mlir::isa<SpatScheduledComputeBatch>(op); }
+};
+
+struct DeferredExchangePlan;
+
+struct RequirementFamily {
+  DeferredExchangePlan* exchange = nullptr;
+  RequirementCoordinate coordinate;
+  LaneSet targetLanes;
+  ProducedValue* producer = nullptr;
+  mlir::Type publicationFragmentType;
+  std::optional<StaticIntSequence> graphLanes;
+  std::optional<StaticIntSequence> producerLocalOffsets;
+};
+
+struct LocalAvailabilityFamily {
+  RequirementFamily* requirement = nullptr;
+  LaneSet targetLanes;
+};
+
+struct ExternalTransferFamily {
+  RequirementFamily* requirement = nullptr;
+  LaneSet targetLanes;
+  ScheduledInfo* sourceScheduled = nullptr;
+  ScheduledInfo* targetScheduled = nullptr;
+  StaticIntSequence sourceStreams = StaticIntSequence::uniform(0, 1);
+  StaticIntSequence targetStreams = StaticIntSequence::uniform(0, 1);
+  StaticIntSequence sourceCores = StaticIntSequence::uniform(0, 1);
+  StaticIntSequence targetCores = StaticIntSequence::uniform(0, 1);
+  StaticIntSequence channelIds = StaticIntSequence::uniform(0, 1);
+};
+
+struct DeferredExchangePlan {
+  uint64_t exchangeId = 0;
+  SpatDeferredCommunicationOp deferred;
+  ScheduledInfo* target = nullptr;
+  unsigned consumerStep = 0;
+  unsigned targetLaneCount = 1;
+  DeferredProgramTemplate program;
+  llvm::SmallVector<RequirementFamily, 0> requirements;
+  llvm::SmallVector<LocalAvailabilityFamily> local;
+  llvm::SmallVector<ExternalTransferFamily, 0> external;
+  unsigned externalTransferCount = 0;
+};
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,432 @@
+#include "DeferredCommunicationPlanning.hpp"
+
+#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/IRMapping.h"
+
+#include "llvm/ADT/SmallPtrSet.h"
+
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
+
+namespace onnx_mlir::spatial {
+using namespace mlir;
+namespace {
+
+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 SmallVector<Value> getBlueprintFragments(SpatBlueprintOp blueprint) {
+  SmallVector<Value> fragments {blueprint.getInput()};
+  llvm::append_range(fragments, blueprint.getFragments());
+  return fragments;
+}
+
+static FailureOr<Value> buildBlueprintReconstruction(
+    OpBuilder &builder, Location loc, SpatBlueprintOp blueprint,
+    ValueRange sourceBlockArgs) {
+  auto resultType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
+  auto operandIndices = blueprint.getFragmentOperandIndices();
+  auto sourceSlots = blueprint.getFragmentSourceSlots();
+  auto sourceOffsets = blueprint.getFragmentSourceOffsets();
+  auto strides = blueprint.getFragmentStrides();
+  if (!resultType || !resultType.hasStaticShape() || !operandIndices ||
+      !sourceSlots || !sourceOffsets || !strides)
+    return blueprint.emitOpError("phase 1 requires complete static fragment assembly metadata"), failure();
+  int64_t rank = resultType.getRank();
+  ArrayRef<int64_t> offsets = blueprint.getFragmentOffsets();
+  ArrayRef<int64_t> sizes = blueprint.getFragmentSizes();
+  if (offsets.size() != sizes.size() || offsets.size() != strides->size() ||
+      offsets.size() != operandIndices->size() * rank ||
+      sourceSlots->size() != operandIndices->size() ||
+      sourceOffsets->size() != operandIndices->size())
+    return blueprint.emitOpError("phase 1 fragment assembly metadata has inconsistent sizes"), failure();
+
+  Value result = tensor::EmptyOp::create(builder, loc, resultType.getShape(),
+                                         resultType.getElementType());
+  for (auto [fragmentIndex, operandIndex] : llvm::enumerate(*operandIndices)) {
+    if (operandIndex < 0 || operandIndex >= static_cast<int64_t>(sourceBlockArgs.size()))
+      return blueprint.emitOpError("phase 1 fragment assembly operand index is out of range"), failure();
+    auto physicalType = dyn_cast<RankedTensorType>(sourceBlockArgs[operandIndex].getType());
+    if (!physicalType || !physicalType.hasStaticShape() || physicalType.getRank() != rank + 1)
+      return blueprint.emitOpError("phase 1 fragment assembly source is not a physical fragment batch"), failure();
+    SmallVector<int64_t> fragmentShape(physicalType.getShape().drop_front());
+    int64_t linearOffset = (*sourceOffsets)[fragmentIndex];
+    SmallVector<int64_t> sourceCoordinates(rank);
+    for (int64_t dim = rank - 1; dim >= 0; --dim) {
+      sourceCoordinates[dim] = linearOffset % fragmentShape[dim];
+      linearOffset /= fragmentShape[dim];
+    }
+    if (linearOffset != 0)
+      return blueprint.emitOpError("phase 1 fragment source offset is out of range"), failure();
+
+    SmallVector<OpFoldResult> sliceOffsets, sliceSizes, sliceStrides;
+    sliceOffsets.push_back(builder.getIndexAttr((*sourceSlots)[fragmentIndex]));
+    sliceSizes.push_back(builder.getIndexAttr(1));
+    sliceStrides.push_back(builder.getIndexAttr(1));
+    SmallVector<int64_t> selectedShape {1};
+    for (int64_t dim = 0; dim < rank; ++dim) {
+      int64_t index = fragmentIndex * rank + dim;
+      int64_t size = sizes[index];
+      if ((*strides)[index] != 1 || sourceCoordinates[dim] < 0 || size <= 0 ||
+          sourceCoordinates[dim] + size > fragmentShape[dim])
+        return blueprint.emitOpError("phase 1 fragment geometry is unsupported"), failure();
+      sliceOffsets.push_back(builder.getIndexAttr(sourceCoordinates[dim]));
+      sliceSizes.push_back(builder.getIndexAttr(size));
+      sliceStrides.push_back(builder.getIndexAttr(1));
+      selectedShape.push_back(size);
+    }
+    auto selectedType = RankedTensorType::get(selectedShape, resultType.getElementType());
+    Value selected = tensor::ExtractSliceOp::create(
+        builder, loc, selectedType, sourceBlockArgs[operandIndex], sliceOffsets,
+        sliceSizes, sliceStrides);
+    SmallVector<int64_t> fragmentResultShape(selectedShape.begin() + 1,
+                                             selectedShape.end());
+    auto fragmentType = RankedTensorType::get(fragmentResultShape,
+                                              resultType.getElementType());
+    SmallVector<ReassociationIndices> reassociation {{0, 1}};
+    for (int64_t dim = 1; dim < rank; ++dim)
+      reassociation.push_back({dim + 1});
+    Value fragment = tensor::CollapseShapeOp::create(
+        builder, loc, fragmentType, selected, reassociation);
+    SmallVector<OpFoldResult> targetOffsets, targetSizes, targetStrides;
+    for (int64_t dim = 0; dim < rank; ++dim) {
+      int64_t index = fragmentIndex * rank + dim;
+      targetOffsets.push_back(builder.getIndexAttr(offsets[index]));
+      targetSizes.push_back(builder.getIndexAttr(sizes[index]));
+      targetStrides.push_back(builder.getIndexAttr((*strides)[index]));
+    }
+    result = tensor::InsertSliceOp::create(builder, loc, fragment, result,
+                                           targetOffsets, targetSizes,
+                                           targetStrides);
+  }
+  return result;
+}
+
+static FailureOr<Value> buildIndexSwitchSelection(OpBuilder &builder, Location loc,
+                                                  Value selector, ValueRange candidates,
+                                                  Operation *diagnosticOwner) {
+  if (candidates.empty())
+    return diagnosticOwner->emitOpError("direct selection requires at least one candidate"), failure();
+  Type type = candidates.front().getType();
+  if (llvm::any_of(candidates, [&](Value candidate) { return candidate.getType() != type; }))
+    return diagnosticOwner->emitOpError("direct selection requires identical candidate types"), failure();
+  if (candidates.size() == 1)
+    return candidates.front();
+
+  SmallVector<int64_t> cases;
+  for (int64_t index = 0; index < static_cast<int64_t>(candidates.size()) - 1; ++index)
+    cases.push_back(index);
+  auto selection = scf::IndexSwitchOp::create(
+    builder, loc, TypeRange {type}, selector, cases, cases.size());
+  auto buildYield = [&](Region &region, Value candidate) {
+    OpBuilder::InsertionGuard guard(builder);
+    Block *block = builder.createBlock(&region);
+    builder.setInsertionPointToEnd(block);
+    scf::YieldOp::create(builder, loc, candidate);
+  };
+  for (auto [region, candidate] : llvm::zip(selection.getCaseRegions(), candidates.drop_back()))
+    buildYield(region, candidate);
+  // The scheduled-lane verifier guarantees an in-range selector, so default is
+  // the final lane without an otherwise-unreachable extra branch.
+  buildYield(selection.getDefaultRegion(), candidates.back());
+  return selection.getResult(0);
+}
+
+static FailureOr<Value> buildSelectedDeferredSource(OpBuilder &builder, Location loc,
+                                                    SpatDeferredCommunicationOp transfer,
+                                                    Value scheduledLane,
+                                                    ValueRange sourceBlockArgs,
+                                                    ArrayRef<int64_t> sourceOperandForScheduledLane) {
+  if (sourceBlockArgs.size() == 1)
+    return sourceBlockArgs.front();
+  if (!scheduledLane || sourceOperandForScheduledLane.empty())
+    return transfer.emitOpError("multiple deferred sources require the enclosing scheduled lane"), failure();
+  auto scheduled = transfer->getParentOfType<SpatScheduledComputeBatch>();
+  if (!scheduled || sourceOperandForScheduledLane.size() != static_cast<size_t>(scheduled.getLaneCount()))
+    return transfer.emitOpError("deferred source mapping must cover every scheduled lane"), failure();
+  SmallVector<Value> candidates;
+  candidates.reserve(sourceOperandForScheduledLane.size());
+  for (int64_t sourceIndex : sourceOperandForScheduledLane) {
+    if (sourceIndex < 0 || sourceIndex >= static_cast<int64_t>(sourceBlockArgs.size()))
+      return transfer.emitOpError("deferred source mapping operand is out of range"), failure();
+    candidates.push_back(sourceBlockArgs[sourceIndex]);
+  }
+  return buildIndexSwitchSelection(builder, loc, scheduledLane, candidates, transfer.getOperation());
+}
+
+static bool isDeferredPayloadCandidateOp(Operation *op) {
+  return isShapingOnlyOp(op) || isCompileTimeOp(op) || isPureIndexComputationOp(op);
+}
+
+static bool isTopLevelDeferredOperation(Operation *op, Block &body,
+                                        const DeferredInputPlan &plan) {
+  (void)plan;
+  return op->getBlock() == &body
+      && (isDeferredPayloadCandidateOp(op) || isa<scf::ForOp>(op));
+}
+
+static bool isEligible(Value value, Block &body, const DeferredInputPlan &plan,
+                       llvm::SmallPtrSetImpl<Operation *> &seen) {
+  if (value == plan.graphInput || value == plan.graphLane || value == plan.scheduledLane)
+    return true;
+  auto arg = dyn_cast<BlockArgument>(value);
+  if (arg)
+    return false;
+  Operation *op = value.getDefiningOp();
+  if (op && op->hasTrait<OpTrait::ConstantLike>())
+    return true;
+  if (!op || !isTopLevelDeferredOperation(op, body, plan) || !seen.insert(op).second)
+    return op && seen.contains(op);
+  return llvm::all_of(op->getOperands(), [&](Value operand) { return isEligible(operand, body, plan, seen); });
+}
+
+static FailureOr<Value> clonePayloadRoot(Value root, Block &body, const DeferredInputPlan &plan,
+                                         OpBuilder &builder, SpatDeferredCommunicationOp transfer,
+                                         Value selectedSource, Value boundGraphLane) {
+  IRMapping mapping;
+  mapping.map(plan.graphInput, selectedSource);
+  std::function<FailureOr<Value>(Value)> cloneScheduledLane = [&](Value value) -> FailureOr<Value> {
+    if (mapping.contains(value)) return mapping.lookup(value);
+    if (value == plan.scheduledLane) return value;
+    if (isa<BlockArgument>(value))
+      return transfer.emitOpError("phase 1 payload shaping captures an unsupported block argument"), failure();
+    Operation *op = value.getDefiningOp();
+    if (!op || (!isDeferredPayloadCandidateOp(op) && !op->hasTrait<OpTrait::ConstantLike>()))
+      return transfer.emitOpError("phase 1 cannot clone the scheduled graph-lane expression"), failure();
+    for (Value operand : op->getOperands()) if (failed(cloneScheduledLane(operand))) return failure();
+    Operation *copy = builder.clone(*op, mapping);
+    for (auto pair : llvm::zip(op->getResults(), copy->getResults())) mapping.map(std::get<0>(pair), std::get<1>(pair));
+    return mapping.lookup(value);
+  };
+  std::function<FailureOr<Value>(Value)> clone = [&](Value value) -> FailureOr<Value> {
+    if (mapping.contains(value)) return mapping.lookup(value);
+    if (value == plan.graphLane) {
+      auto mappedLane = cloneScheduledLane(boundGraphLane ? boundGraphLane : plan.scheduledGraphLane);
+      if (failed(mappedLane)) return failure();
+      mapping.map(value, *mappedLane);
+      return *mappedLane;
+    }
+    if (isa<BlockArgument>(value))
+      return transfer.emitOpError("phase 1 payload shaping captures an unsupported block argument"), failure();
+    Operation *op = value.getDefiningOp();
+    if (!op || (!isTopLevelDeferredOperation(op, body, plan) && !op->hasTrait<OpTrait::ConstantLike>()))
+      return transfer.emitOpError("phase 1 payload shaping contains an unsupported operation"), failure();
+    for (Value operand : op->getOperands()) if (failed(clone(operand))) return failure();
+    Operation *copy = builder.clone(*op, mapping);
+    for (auto pair : llvm::zip(op->getResults(), copy->getResults())) mapping.map(std::get<0>(pair), std::get<1>(pair));
+    return mapping.lookup(value);
+  };
+  return clone(root);
+}
+
+static bool dependsOnGraphLane(Value value, Value graphLane, Block &body,
+                               const DeferredInputPlan &plan,
+                               llvm::SmallPtrSetImpl<Operation *> &seen) {
+  if (value == graphLane)
+    return true;
+  Operation *op = value.getDefiningOp();
+  if (!op || !isTopLevelDeferredOperation(op, body, plan) || !seen.insert(op).second)
+    return false;
+  if (auto loop = dyn_cast<scf::ForOp>(op)) {
+    bool depends = false;
+    loop.getRegion().walk([&](Operation *nested) {
+      depends |= llvm::is_contained(nested->getOperands(), graphLane);
+    });
+    if (depends)
+      return true;
+  }
+  return llvm::any_of(op->getOperands(), [&](Value operand) {
+    return dependsOnGraphLane(operand, graphLane, body, plan, seen);
+  });
+}
+
+static void collectClosure(Value value, Block &body, const DeferredInputPlan &plan,
+                           llvm::SmallPtrSetImpl<Operation *> &ops) {
+  Operation *op = value.getDefiningOp();
+  if (!op || !isTopLevelDeferredOperation(op, body, plan) || !ops.insert(op).second) return;
+  if (auto loop = dyn_cast<scf::ForOp>(op))
+    loop.getRegion().walk([&](Operation *nested) { ops.insert(nested); });
+  for (Value operand : op->getOperands())
+    if (operand != plan.graphInput && operand != plan.graphLane) collectClosure(operand, body, plan, ops);
+}
+
+} // namespace
+
+bool isDeferredFragmentAssemblyInput(
+    Value input, const ComputeInstance &consumerInstance) {
+  auto blueprint = input.getDefiningOp<SpatBlueprintOp>();
+  if (!blueprint || blueprint.getMode() != "fragment_assembly")
+    return false;
+  return llvm::all_of(getBlueprintFragments(blueprint), [&](Value fragment) {
+    return getProducerValueRef(fragment, &consumerInstance).has_value();
+  });
+}
+
+LogicalResult prepareSingleCpuInput(OpBuilder &, Location loc, Value input, BlockArgument graphInput,
+                                    const ComputeInstance &consumerInstance, const MergeScheduleResult &,
+                                    ValueRange scheduledInputs, Block &block, unsigned firstInputArgument,
+                                    ArrayRef<ProducerValueKey> carriedKeys, Value graphLane, Value scheduledGraphLane,
+                                    DeferredInputPlan &plan) {
+  plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, {}, {}, {}, {}, 1, nullptr};
+  if (isDeferredFragmentAssemblyInput(input, consumerInstance)) {
+    plan.blueprint = input.getDefiningOp<SpatBlueprintOp>();
+    plan.originalSources = getBlueprintFragments(plan.blueprint);
+    return success();
+  }
+  auto producer = getProducerValueRef(input, &consumerInstance);
+  if (!producer) { plan.availableValue = getBlockOperand(block, scheduledInputs, input, firstInputArgument); return success(); }
+  ProducerValueKey key {producer->instance, producer->resultIndex};
+  auto carried = llvm::find(carriedKeys, key);
+  if (carried != carriedKeys.end()) {
+    plan.availableValue = block.getArgument(firstInputArgument + scheduledInputs.size() + std::distance(carriedKeys.begin(), carried));
+    return success();
+  }
+  auto source = getOriginalProducerValue(*producer);
+  if (failed(source)) return emitError(loc) << "cannot resolve original graph producer value";
+  plan.originalSources.push_back(*source);
+  return success();
+}
+
+LogicalResult prepareMultiCpuTupleInput(OpBuilder &, Location loc, Value input, BlockArgument graphInput,
+                                        const ComputeStepTuple &tuple, const PeftClassPlan &,
+                                        const MergeScheduleResult &, ValueRange scheduledInputs, Block &block,
+                                        unsigned firstInputArgument, Value graphLane, Value scheduledGraphLane, Value scheduledLane,
+                                        DeferredInputPlan &plan) {
+  const ComputeInstance &representative = tuple.instances.front();
+  plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, scheduledLane, {}, {}, {}, 1, nullptr};
+  if (isDeferredFragmentAssemblyInput(input, representative)) {
+    plan.blueprint = input.getDefiningOp<SpatBlueprintOp>();
+    plan.originalSources = getBlueprintFragments(plan.blueprint);
+    return success();
+  }
+  auto producer = getProducerValueRef(input, &representative);
+  if (!producer) { plan.availableValue = getBlockOperand(block, scheduledInputs, input, firstInputArgument); return success(); }
+  auto inputs = getComputeInstanceInputs(representative);
+  auto it = llvm::find(inputs, input);
+  if (it == inputs.end()) return emitError(loc) << "cannot resolve scheduled batch step input";
+  unsigned inputIndex = std::distance(inputs.begin(), it);
+  for (const ComputeInstance &instance : tuple.instances) {
+    auto laneInputs = getComputeInstanceInputs(instance);
+    if (inputIndex >= laneInputs.size()) return emitError(loc) << "scheduled batch step input out of range";
+    auto laneProducer = getProducerValueRef(laneInputs[inputIndex], &instance);
+    if (!laneProducer) return emitError(loc) << "scheduled batch step mixes host and producer inputs";
+    auto source = getOriginalProducerValue(*laneProducer);
+    if (failed(source)) return emitError(loc) << "cannot resolve original graph producer value";
+    auto sourceIt = llvm::find(plan.originalSources, *source);
+    if (sourceIt == plan.originalSources.end()) { plan.sourceOperandForScheduledLane.push_back(plan.originalSources.size()); plan.originalSources.push_back(*source); }
+    else plan.sourceOperandForScheduledLane.push_back(std::distance(plan.originalSources.begin(), sourceIt));
+  }
+  return success();
+}
+
+LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc, Block &body,
+                                                 ArrayRef<DeferredInputPlan> plans, IRMapping &mapper,
+                                                 llvm::SmallPtrSetImpl<Operation *> &absorbed) {
+  for (const DeferredInputPlan &plan : plans) {
+    if (plan.availableValue) { mapper.map(plan.graphInput, plan.availableValue); continue; }
+    SmallVector<Value> roots;
+    bool needsIdentity = false;
+    SmallVector<Value> worklist {plan.graphInput};
+    llvm::SmallDenseSet<Value, 32> seen;
+    while (!worklist.empty()) {
+      Value value = worklist.pop_back_val();
+      if (!seen.insert(value).second) continue;
+      for (OpOperand &use : value.getUses()) {
+        Operation *user = use.getOwner();
+        if (!isTopLevelDeferredOperation(user, body, plan)) { needsIdentity = true; continue; }
+        llvm::SmallPtrSet<Operation *, 16> eligibility;
+        if (!isEligible(user->getResult(0), body, plan, eligibility)) { needsIdentity = true; continue; }
+        for (Value result : user->getResults()) {
+          bool hasShapingUse = llvm::any_of(result.getUses(), [&](OpOperand &next) { return isTopLevelDeferredOperation(next.getOwner(), body, plan); });
+          bool hasOtherUse = llvm::any_of(result.getUses(), [&](OpOperand &next) { return !isTopLevelDeferredOperation(next.getOwner(), body, plan); });
+          if (hasOtherUse) roots.push_back(result);
+          if (hasShapingUse) worklist.push_back(result);
+        }
+      }
+    }
+    if (needsIdentity) roots.push_back(plan.graphInput);
+    llvm::sort(roots, [](Value a, Value b) { return a.getAsOpaquePointer() < b.getAsOpaquePointer(); });
+    roots.erase(std::unique(roots.begin(), roots.end()), roots.end());
+    for (Value root : roots) {
+      llvm::SmallPtrSet<Operation *, 16> laneDependencies;
+      bool scalarize = plan.scalarizedGraphLaneBase
+                    && dependsOnGraphLane(root, plan.graphLane, body, plan, laneDependencies);
+      OpBuilder::InsertPoint restore = builder.saveInsertionPoint();
+      Operation *loop = nullptr;
+      if (scalarize) {
+        loop = builder.getInsertionBlock()->getParentOp();
+        if (loop && !isa<scf::ForOp>(loop))
+          loop = loop->getParentOfType<scf::ForOp>();
+        if (loop)
+          builder.setInsertionPoint(loop);
+        else if (plan.scalarizedHoistBlock)
+          builder.setInsertionPointToEnd(plan.scalarizedHoistBlock);
+        else
+          return emitError(loc) << "phase 1 scalarized deferred payload is missing a hoist point";
+      }
+      SmallVector<Value> payloads;
+      unsigned count = scalarize ? plan.scalarizedLaneCount : 1;
+      for (unsigned offset = 0; offset < count; ++offset) {
+      auto transfer = SpatDeferredCommunicationOp::create(builder, loc, root.getType(), plan.originalSources);
+      Block *deferred = builder.createBlock(&transfer.getBody(), transfer.getBody().end(),
+        TypeRange {transfer.getSources().getTypes()}, SmallVector<Location>(transfer.getSources().size(), loc));
+      builder.setInsertionPointToStart(deferred);
+      auto selected = plan.blueprint
+                        ? buildBlueprintReconstruction(builder, loc, plan.blueprint,
+                                                       deferred->getArguments())
+                        : buildSelectedDeferredSource(builder, loc, transfer,
+                                                      plan.scheduledLane,
+                                                      deferred->getArguments(),
+                                                      plan.sourceOperandForScheduledLane);
+      if (failed(selected)) return failure();
+      Value boundGraphLane;
+      if (scalarize) {
+        boundGraphLane = affineAddConst(
+          builder, loc, plan.scalarizedGraphLaneBase, offset, transfer.getOperation());
+      }
+      auto payload = clonePayloadRoot(root, body, plan, builder, transfer, *selected, boundGraphLane);
+      if (failed(payload)) return failure();
+      SpatYieldOp::create(builder, loc, *payload);
+      payloads.push_back(transfer.getOutput());
+      builder.setInsertionPointAfter(transfer);
+      }
+      if (scalarize) {
+        builder.restoreInsertionPoint(restore);
+        auto selected = buildIndexSwitchSelection(
+          builder, loc, plan.scalarizedLocalLane, payloads, root.getDefiningOp());
+        if (failed(selected)) return failure();
+        mapper.map(root, *selected);
+      } else {
+        mapper.map(root, payloads.front());
+      }
+      collectClosure(root, body, plan, absorbed);
+    }
+  }
+  SmallVector<Operation *> notFullyAbsorbed;
+  for (Operation *op : absorbed) {
+    bool allResultsMapped = llvm::all_of(op->getResults(), [&](Value result) {
+      return mapper.contains(result) || llvm::all_of(result.getUses(), [&](OpOperand &use) { return absorbed.contains(use.getOwner()); });
+    });
+    if (!allResultsMapped)
+      notFullyAbsorbed.push_back(op);
+  }
+  for (Operation *op : notFullyAbsorbed)
+    absorbed.erase(op);
+  return success();
+}
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,54 @@
+#pragma once
+
+#include "ScheduledComputePlan.hpp"
+
+namespace onnx_mlir {
+namespace spatial {
+
+// A graph input is either already available in the scheduled block, or is a
+// graph result whose individual deterministic payloads are materialized later.
+struct DeferredInputPlan {
+  BlockArgument graphInput;
+  Value availableValue;
+  SmallVector<Value> originalSources;
+  SmallVector<int64_t> sourceOperandForScheduledLane;
+  Value graphLane;
+  Value scheduledGraphLane;
+  Value scheduledLane;
+  SpatBlueprintOp blueprint;
+  Value scalarizedLocalLane;
+  Value scalarizedGraphLaneBase;
+  int64_t scalarizedLaneCount = 1;
+  Block *scalarizedHoistBlock = nullptr;
+};
+
+bool isDeferredFragmentAssemblyInput(Value input,
+                                     const ComputeInstance &consumerInstance);
+
+LogicalResult prepareSingleCpuInput(OpBuilder &builder, Location loc, Value input,
+                                    BlockArgument graphInput,
+                                    const ComputeInstance &consumerInstance,
+                                    const MergeScheduleResult &schedule,
+                                    ValueRange scheduledInputs, Block &block,
+                                    unsigned firstInputArgument,
+                                    ArrayRef<ProducerValueKey> carriedKeys,
+                                    Value graphLane, Value scheduledGraphLane,
+                                    DeferredInputPlan &plan);
+
+LogicalResult prepareMultiCpuTupleInput(OpBuilder &builder, Location loc, Value input,
+                                        BlockArgument graphInput,
+                                        const ComputeStepTuple &stepTuple,
+                                        const PeftClassPlan &peftClassPlan,
+                                        const MergeScheduleResult &schedule,
+                                        ValueRange scheduledInputs, Block &block,
+                                        unsigned firstInputArgument, Value graphLane, Value scheduledGraphLane,
+                                        Value scheduledLane, DeferredInputPlan &plan);
+
+LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc,
+                                                 Block &graphBody,
+                                                 ArrayRef<DeferredInputPlan> plans,
+                                                 IRMapping &mapper,
+                                                 llvm::SmallPtrSetImpl<Operation *> &absorbed);
+
+} // namespace spatial
+} // namespace onnx_mlir

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

@@ -0,0 +1,272 @@
+#include "DeferredCommunicationRealization.hpp"
+
+#include "DeferredBoundaryPlanning.hpp"
+#include "DeferredBoundaryRealization.hpp"
+#include "DeferredCommunicationDeadlock.hpp"
+#include "DeferredCommunicationScheduling.hpp"
+#include "DeferredTransferPlanning.hpp"
+
+#include "mlir/IR/Dominance.h"
+
+namespace onnx_mlir::spatial {
+using namespace mlir;
+namespace {
+
+static LogicalResult validateScalarLinearization(ScheduledInfo &info) {
+  auto scheduled = cast<SpatScheduledCompute>(info.op);
+  for (unsigned index = 1; index < info.blocks.size(); ++index) {
+    auto previous = dyn_cast<SpatBlockYieldOp>(
+      info.blocks[index - 1]->getTerminator());
+    if (!previous
+        || previous.getOutputs().size()
+             != info.blocks[index]->getNumArguments())
+      return scheduled.emitOpError(
+        "phase 2 cannot linearize malformed scalar scheduled blocks");
+    for (auto [argument, value] : llvm::zip(
+           info.blocks[index]->getArguments(), previous.getOutputs())) {
+      if (argument.getType() == value.getType())
+        continue;
+      for (Operation *user : argument.getUsers())
+        if (!isa<SpatBlockYieldOp>(user))
+          return scheduled.emitOpError(
+            "phase 2 cannot linearize a live mismatched carried value");
+    }
+  }
+  return success();
+}
+
+static LogicalResult linearizeScalar(ScheduledInfo &info,
+                                     IRRewriter &rewriter) {
+  auto scheduled = cast<SpatScheduledCompute>(info.op);
+  if (failed(validateScalarLinearization(info)))
+    return failure();
+  Block *first = info.blocks.front();
+  SmallVector<SmallVector<Value>> incoming(info.blocks.size());
+  for (unsigned index = 1; index < info.blocks.size(); ++index) {
+    auto previous = cast<SpatBlockYieldOp>(
+      info.blocks[index - 1]->getTerminator());
+    incoming[index].assign(previous.getOutputs().begin(),
+                           previous.getOutputs().end());
+  }
+  IRMapping carried;
+  for (unsigned index = 1; index < info.blocks.size(); ++index)
+    for (auto [argument, value] : llvm::zip(
+           info.blocks[index]->getArguments(), incoming[index])) {
+      Value resolved = carried.lookupOrDefault(value);
+      if (argument.getType() == resolved.getType()) {
+        carried.map(argument, resolved);
+        argument.replaceAllUsesWith(resolved);
+      }
+    }
+  for (unsigned index = 1; index < info.blocks.size(); ++index)
+    for (Operation &op : llvm::make_early_inc_range(
+           info.blocks[index]->without_terminator()))
+      op.moveBefore(first->getTerminator());
+  for (Block *block : info.blocks)
+    cast<SpatBlockYieldOp>(block->getTerminator()).erase();
+  SmallVector<Value> outputs(scheduled.getNumResults());
+  for (ProducedValue *produced : info.produced) {
+    unsigned result = info.resultOffsets[produced->step]
+        + produced->resultIndex;
+    outputs[result] = produced->payload;
+  }
+  if (llvm::any_of(outputs, [](Value output) { return !output; }))
+    return scheduled.emitOpError(
+      "phase 2 cannot recover every scheduled scalar result");
+  rewriter.setInsertionPointToEnd(first);
+  SpatYieldOp::create(rewriter, scheduled.getLoc(), outputs);
+  scheduled->setAttr("scheduled.realized", rewriter.getBoolAttr(true));
+  for (unsigned index = 1; index < info.blocks.size(); ++index) {
+    if (llvm::any_of(info.blocks[index]->getArguments(),
+                     [](BlockArgument argument) {
+                       return !argument.use_empty();
+                     }))
+      return scheduled.emitOpError(
+        "phase 2 scalar linearization left a live block argument");
+    info.blocks[index]->erase();
+  }
+  info.blocks.assign(1, first);
+  return success();
+}
+
+static LogicalResult linearizeBatch(ScheduledInfo &info,
+                                    IRRewriter &rewriter) {
+  auto scheduled = cast<SpatScheduledComputeBatch>(info.op);
+  Block *first = info.blocks.front();
+  SmallVector<SpatInParallelOp> terminators;
+  for (Block *block : info.blocks) {
+    auto parallel = dyn_cast<SpatInParallelOp>(block->getTerminator());
+    if (!parallel)
+      return scheduled.emitOpError(
+        "phase 2 cannot linearize a batch block without spat.in_parallel");
+    terminators.push_back(parallel);
+  }
+  for (unsigned index = 1; index < info.blocks.size(); ++index)
+    for (auto [argument, firstArgument] : llvm::zip(
+           info.blocks[index]->getArguments(), first->getArguments())) {
+      if (argument.getType() != firstArgument.getType())
+        return scheduled.emitOpError(
+          "phase 2 cannot linearize incompatible batch block arguments");
+      argument.replaceAllUsesWith(firstArgument);
+    }
+  for (unsigned index = 1; index < info.blocks.size(); ++index)
+    for (Operation &op : llvm::make_early_inc_range(
+           info.blocks[index]->without_terminator()))
+      op.moveBefore(first->getTerminator());
+  rewriter.setInsertionPoint(first->getTerminator());
+  auto combined = SpatInParallelOp::create(rewriter, scheduled.getLoc());
+  Block &combinedBlock = combined.getRegion().front();
+  for (SpatInParallelOp parallel : terminators)
+    for (Operation &op : llvm::make_early_inc_range(
+           parallel.getRegion().front()))
+      op.moveBefore(&combinedBlock, combinedBlock.end());
+  for (SpatInParallelOp parallel : terminators)
+    parallel.erase();
+  scheduled->setAttr("scheduled.realized", rewriter.getBoolAttr(true));
+  for (unsigned index = 1; index < info.blocks.size(); ++index) {
+    if (llvm::any_of(info.blocks[index]->getArguments(),
+                     [](BlockArgument argument) {
+                       return !argument.use_empty();
+                     }))
+      return scheduled.emitOpError(
+        "phase 2 batch linearization left a live block argument");
+    info.blocks[index]->erase();
+  }
+  info.blocks.assign(1, first);
+  return success();
+}
+
+static LogicalResult linearizeScheduled(
+    DeferredTransferPlan &plan, IRRewriter &rewriter) {
+  for (ScheduledInfo &info : plan.scheduled) {
+    LogicalResult result = info.isBatch()
+      ? linearizeBatch(info, rewriter) : linearizeScalar(info, rewriter);
+    if (failed(result))
+      return failure();
+  }
+  return success();
+}
+
+static LogicalResult replaceFinalGraphPublications(
+    func::FuncOp funcOp, DeferredTransferPlan &plan) {
+  for (Operation &op : funcOp.getOps()) {
+    if (!isa<SpatGraphCompute, SpatGraphComputeBatch>(op))
+      continue;
+    auto graphId = op.getAttrOfType<IntegerAttr>("scheduled.graph_id");
+    if (!graphId)
+      continue;
+    for (auto [resultIndex, result] : llvm::enumerate(op.getResults())) {
+      SmallVector<OpOperand *> externalUses;
+      for (OpOperand &use : result.getUses())
+        if (!isa<SpatGraphCompute, SpatGraphComputeBatch>(use.getOwner()))
+          externalUses.push_back(&use);
+      if (externalUses.empty())
+        continue;
+      SmallVector<Value> exact;
+      for (ProducedValue *produced :
+           plan.producedByGraph.lookup(graphId.getInt()))
+        if (produced->resultIndex == resultIndex
+            && produced->published.getType() == result.getType()
+            && !llvm::is_contained(exact, produced->published))
+          exact.push_back(produced->published);
+      if (exact.size() != 1)
+        return op.emitOpError(
+          "phase 2 final publication ownership changed after planning");
+      for (OpOperand *use : externalUses) {
+        Operation *consumer = use->getOwner();
+        Operation *producer = exact.front().getDefiningOp();
+        if (consumer->getBlock() == producer->getBlock()
+            && consumer->isBeforeInBlock(producer))
+          consumer->moveAfter(producer);
+        use->set(exact.front());
+      }
+    }
+  }
+  return success();
+}
+
+static LogicalResult eraseOldGraph(func::FuncOp funcOp,
+                                   IRRewriter &rewriter) {
+  SmallVector<Operation *> graphOps;
+  for (Operation &op : funcOp.getOps())
+    if (isa<SpatGraphCompute, SpatGraphComputeBatch>(op))
+      graphOps.push_back(&op);
+  for (Operation *op : llvm::reverse(graphOps)) {
+    if (!op->use_empty())
+      return op->emitOpError(
+        "phase 2 cannot erase an old graph compute with live results");
+    rewriter.eraseOp(op);
+  }
+  SmallVector<SpatBlueprintOp> deadBlueprints;
+  funcOp.walk([&](SpatBlueprintOp blueprint) {
+    if (blueprint.getOutput().use_empty())
+      deadBlueprints.push_back(blueprint);
+  });
+  for (SpatBlueprintOp blueprint : deadBlueprints)
+    rewriter.eraseOp(blueprint);
+  return success();
+}
+
+static LogicalResult verifyDominance(func::FuncOp funcOp) {
+  DominanceInfo dominance(funcOp);
+  WalkResult result = funcOp.walk([&](Operation *op) {
+    for (auto [index, operand] : llvm::enumerate(op->getOperands()))
+      if (!dominance.dominates(operand, op)) {
+        op->emitOpError() << "phase 2 produced non-dominating operand "
+                          << index << ": " << operand;
+        return WalkResult::interrupt();
+      }
+    return WalkResult::advance();
+  });
+  return success(!result.wasInterrupted());
+}
+
+} // namespace
+
+LogicalResult realizeDeferredCommunication(func::FuncOp funcOp) {
+  auto transfers = buildDeferredTransferPlan(funcOp);
+  if (failed(transfers))
+    return funcOp.emitOpError(
+      "phase 2 failed to build symbolic transfer families");
+  auto schedule = scheduleDeferredCommunication(funcOp, *transfers);
+  if (failed(schedule)
+      || failed(verifyPlannedCommunicationDeadlockFree(
+        funcOp, transfers->stepCounts, *schedule)))
+    return funcOp.emitOpError(
+      "phase 2 failed to schedule symbolic communication");
+  auto boundaries = buildDeferredBoundaryPlan(*transfers, *schedule);
+  if (failed(boundaries))
+    return funcOp.emitOpError(
+      "phase 2 failed to build sparse boundary programs");
+
+  IRRewriter rewriter(funcOp.getContext());
+  if (failed(linearizeScheduled(*transfers, rewriter)))
+    return failure();
+  ConstantPool constants(funcOp, rewriter);
+  DeferredEmissionContext context(rewriter, constants);
+  DeferredEraseSet erase;
+  if (failed(realizeDeferredBoundaries(
+        boundaries->boundaries, boundaries->results, context, erase)))
+    return failure();
+  for (Operation *op : erase) {
+    if (!op->use_empty())
+      return op->emitOpError(
+        "phase 2 cannot erase deferred communication with live uses");
+    rewriter.eraseOp(op);
+  }
+  if (failed(retargetDeferredPublications(funcOp, *transfers))
+      || failed(replaceFinalGraphPublications(funcOp, *transfers))
+      || failed(eraseOldGraph(funcOp, rewriter))
+      || failed(verifyDominance(funcOp))
+      || failed(verifyRealizedCommunicationDeadlockFree(funcOp, *schedule)))
+    return failure();
+  bool deferredRemains = false;
+  funcOp.walk([&](SpatDeferredCommunicationOp deferred) {
+    deferred.emitOpError(
+      "phase 2 left an unrealized deferred communication");
+    deferredRemains = true;
+  });
+  return success(!deferredRemains);
+}
+
+} // namespace onnx_mlir::spatial

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

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

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

@@ -0,0 +1,409 @@
+#include "llvm/ADT/MapVector.h"
+
+#include <limits>
+#include <queue>
+
+#include "DeferredCommunicationScheduling.hpp"
+#include "DeferredTransferPlanning.hpp"
+
+namespace onnx_mlir::spatial {
+using namespace mlir;
+namespace {
+
+struct StreamThreshold {
+  unsigned stream = 0;
+  unsigned completedStep = 0;
+};
+
+struct TransferGroup {
+  DeferredExchangePlan* exchange = nullptr;
+  SmallVector<ScheduledTransferSlice> ordered;
+  SmallVector<StreamThreshold> dependencies;
+  unsigned unsatisfied = 0;
+  bool ready = false;
+  bool scheduled = false;
+  BoundaryCost cost;
+  std::tuple<unsigned, unsigned, unsigned, uint64_t> originalPriority;
+  std::optional<TransferEmissionSignature> firstSignature;
+};
+
+struct StreamProgress {
+  unsigned completedStep = 0;
+  SmallVector<unsigned> pendingAtStep;
+  SmallVector<SmallVector<unsigned>> unlockedAtStep;
+  bool queued = false;
+};
+
+static size_t hashSignature(const TransferEmissionSignature& signature) {
+  return llvm::hash_combine(signature.scheduled,
+                            signature.payload.getAsOpaquePointer(),
+                            signature.fragmentType.getAsOpaquePointer(),
+                            signature.hasGraphLane,
+                            signature.sourceIsBatch);
+}
+
+static bool betterStaticPriority(const TransferGroup& lhs, const TransferGroup& rhs) {
+  auto left = std::tuple(lhs.cost.instructionCount,
+                         lhs.cost.branchRegions,
+                         std::numeric_limits<unsigned>::max()
+                           - lhs.cost.absorbedTransfers,
+                         lhs.cost.lookupEntries,
+                         lhs.originalPriority);
+  auto right = std::tuple(rhs.cost.instructionCount,
+                          rhs.cost.branchRegions,
+                          std::numeric_limits<unsigned>::max()
+                            - rhs.cost.absorbedTransfers,
+                          rhs.cost.lookupEntries,
+                          rhs.originalPriority);
+  return left < right;
+}
+
+static bool orderPermutationCycles(SmallVectorImpl<ScheduledTransferSlice>& slices) {
+  if (slices.size() < 2)
+    return false;
+  DenseMap<unsigned, unsigned> edgeBySource;
+  DenseMap<unsigned, unsigned> edgeByTarget;
+  for (auto [index, slice] : llvm::enumerate(slices)) {
+    ExternalTransferFamily& family = *slice.family;
+    if (slice.transferCount != 1 || family.sourceScheduled != family.targetScheduled)
+      return false;
+    unsigned source = family.requirement->producer->scheduledLane;
+    unsigned target = family.targetLanes.intervals().front().begin + slice.familyOffset;
+    if (!edgeBySource.try_emplace(source, index).second || !edgeByTarget.try_emplace(target, index).second)
+      return false;
+  }
+  SmallVector<ScheduledTransferSlice> ordered;
+  SmallVector<bool> emitted(slices.size());
+  auto appendChain = [&](unsigned source) {
+    while (true) {
+      auto edge = edgeBySource.find(source);
+      if (edge == edgeBySource.end() || emitted[edge->second])
+        break;
+      unsigned index = edge->second;
+      emitted[index] = true;
+      ordered.push_back(slices[index]);
+      ExternalTransferFamily& family = *slices[index].family;
+      source = family.targetLanes.intervals().front().begin + slices[index].familyOffset;
+    }
+  };
+  SmallVector<unsigned> pathStarts;
+  for (auto [source, index] : edgeBySource)
+    if (!edgeByTarget.count(source))
+      pathStarts.push_back(source);
+  llvm::sort(pathStarts);
+  for (unsigned source : pathStarts)
+    appendChain(source);
+  while (ordered.size() != slices.size()) {
+    unsigned source = std::numeric_limits<unsigned>::max();
+    for (auto [candidate, index] : edgeBySource)
+      if (!emitted[index])
+        source = std::min(source, candidate);
+    appendChain(source);
+  }
+  slices = std::move(ordered);
+  return true;
+}
+
+static void orderGroupSlices(TransferGroup& group) {
+  SmallVector<TransferEmissionSignature> signatures;
+  DenseMap<size_t, SmallVector<unsigned>> signatureIdsByHash;
+  SmallVector<SmallVector<ScheduledTransferSlice>> slicesBySignature;
+  SmallVector<unsigned> order;
+  for (ExternalTransferFamily& family : group.exchange->external) {
+    ScheduledTransferSlice slice {&family, 0, family.targetLanes.size()};
+    TransferEmissionSignature signature = getTransferEmissionSignature(family);
+    size_t hash = hashSignature(signature);
+    std::optional<unsigned> id;
+    for (unsigned candidate : signatureIdsByHash.lookup(hash))
+      if (signatures[candidate] == signature) {
+        id = candidate;
+        break;
+      }
+    if (!id) {
+      id = signatures.size();
+      signatures.push_back(signature);
+      signatureIdsByHash[hash].push_back(*id);
+      slicesBySignature.emplace_back();
+      order.push_back(*id);
+    }
+    slicesBySignature[*id].push_back(slice);
+  }
+  llvm::stable_sort(order, [&](unsigned lhs, unsigned rhs) {
+    auto count = [&](unsigned id) {
+      size_t result = 0;
+      for (const ScheduledTransferSlice& slice : slicesBySignature[id])
+        result += slice.transferCount;
+      return result;
+    };
+    return count(lhs) > count(rhs);
+  });
+  for (unsigned id : order) {
+    if (!orderPermutationCycles(slicesBySignature[id]))
+      llvm::stable_sort(slicesBySignature[id], [](const ScheduledTransferSlice& lhs, const ScheduledTransferSlice& rhs) {
+        ExternalTransferFamily& left = *lhs.family;
+        ExternalTransferFamily& right = *rhs.family;
+        return std::tuple(left.sourceStreams.valueAt(lhs.familyOffset), left.targetLanes.intervals().front().begin)
+             > std::tuple(right.sourceStreams.valueAt(rhs.familyOffset), right.targetLanes.intervals().front().begin);
+      });
+    llvm::append_range(group.ordered, slicesBySignature[id]);
+  }
+  group.cost = estimateCanonicalBoundaryCost({}, group.ordered);
+  if (!group.ordered.empty())
+    group.firstSignature = getTransferEmissionSignature(*group.ordered.front().family);
+}
+
+static SmallVector<TransferGroup> buildGroups(DeferredTransferPlan& plan) {
+  SmallVector<TransferGroup> groups;
+  for (const std::unique_ptr<DeferredExchangePlan>& exchange : plan.exchanges) {
+    if (exchange->external.empty())
+      continue;
+    TransferGroup group;
+    group.exchange = exchange.get();
+    DenseMap<unsigned, unsigned> thresholdByStream;
+    unsigned minSource = std::numeric_limits<unsigned>::max();
+    unsigned minTarget = std::numeric_limits<unsigned>::max();
+    for (ExternalTransferFamily& family : exchange->external) {
+      unsigned source = family.sourceStreams.valueAt(0);
+      thresholdByStream[source] = std::max(thresholdByStream.lookup(source), family.requirement->producer->step + 1);
+      minSource = std::min(minSource, source);
+      for (size_t index = 0; index < family.targetStreams.size(); ++index)
+        minTarget = std::min<unsigned>(minTarget, family.targetStreams.valueAt(index));
+    }
+    for (LocalAvailabilityFamily& local : exchange->local)
+      for (LaneInterval interval : local.targetLanes.intervals())
+        for (unsigned lane = interval.begin; lane < interval.end; ++lane) {
+          unsigned stream = exchange->target->streamIds[lane];
+          thresholdByStream[stream] = std::max(thresholdByStream.lookup(stream), local.requirement->producer->step + 1);
+        }
+    SmallVector<unsigned> streams;
+    for (auto [stream, threshold] : thresholdByStream)
+      streams.push_back(stream);
+    llvm::sort(streams);
+    for (unsigned stream : streams)
+      group.dependencies.push_back({stream, thresholdByStream.lookup(stream)});
+    group.unsatisfied =
+      llvm::count_if(group.dependencies, [](StreamThreshold dependency) { return dependency.completedStep != 0; });
+    group.originalPriority = {exchange->consumerStep, minSource, minTarget, exchange->exchangeId};
+    orderGroupSlices(group);
+    groups.push_back(std::move(group));
+  }
+  return groups;
+}
+
+static unsigned tailExtension(const TransferGroup& group,
+                              const TransferEmissionSignature& tail,
+                              unsigned tailBoundary,
+                              ArrayRef<StreamProgress> streams) {
+  unsigned extension = 0;
+  for (const ScheduledTransferSlice& slice : group.ordered) {
+    ExternalTransferFamily& family = *slice.family;
+    if (!(getTransferEmissionSignature(family) == tail))
+      break;
+    for (size_t index = 0; index < slice.transferCount; ++index) {
+      size_t familyIndex = slice.familyOffset + index;
+      if (streams[family.sourceStreams.valueAt(familyIndex)].completedStep
+          != tailBoundary)
+        return extension;
+      ++extension;
+    }
+  }
+  return extension;
+}
+
+} // namespace
+
+TransferEmissionSignature getTransferEmissionSignature(const ExternalTransferFamily& family) {
+  ProducedValue* producer = family.requirement->producer;
+  return {producer->scheduled,
+          producer->payload,
+          family.requirement->publicationFragmentType,
+          family.requirement->graphLanes.has_value(),
+          producer->scheduled->isBatch()};
+}
+
+BoundaryCost estimateCanonicalBoundaryCost(
+    ArrayRef<ScheduledTransferSlice> existingTail,
+    ArrayRef<ScheduledTransferSlice> candidate) {
+  BoundaryCost cost;
+  std::optional<TransferEmissionSignature> previous;
+  RequirementFamily* previousRequirement = nullptr;
+  if (!existingTail.empty()) {
+    previous = getTransferEmissionSignature(*existingTail.back().family);
+    previousRequirement = existingTail.back().family->requirement;
+  }
+  for (const ScheduledTransferSlice& slice : candidate) {
+    TransferEmissionSignature signature = getTransferEmissionSignature(*slice.family);
+    if (!previous || !(*previous == signature))
+      ++cost.instructionCount;
+    else
+      cost.absorbedTransfers += slice.transferCount;
+    previous = signature;
+    if (previousRequirement != slice.family->requirement)
+      ++cost.instructionCount;
+    previousRequirement = slice.family->requirement;
+    cost.lookupEntries += slice.transferCount;
+  }
+  return cost;
+}
+
+FailureOr<ScheduledCommunicationPlan> scheduleDeferredCommunication(func::FuncOp funcOp, DeferredTransferPlan& plan) {
+  SmallVector<TransferGroup> groups = buildGroups(plan);
+  SmallVector<StreamProgress> streams(plan.stepCounts.size());
+  for (auto [stream, progress] : llvm::enumerate(streams)) {
+    progress.pendingAtStep.resize(plan.stepCounts[stream]);
+    progress.unlockedAtStep.resize(plan.stepCounts[stream] + 1);
+  }
+  for (const std::unique_ptr<DeferredExchangePlan>& exchange : plan.exchanges)
+    for (ExternalTransferFamily& family : exchange->external)
+      for (size_t index = 0; index < family.targetStreams.size(); ++index)
+        ++streams[family.targetStreams.valueAt(index)].pendingAtStep[exchange->consumerStep];
+  for (auto [groupIndex, group] : llvm::enumerate(groups))
+    for (StreamThreshold dependency : group.dependencies)
+      streams[dependency.stream].unlockedAtStep[dependency.completedStep].push_back(groupIndex);
+
+  auto compare = [&](unsigned lhs, unsigned rhs) { return betterStaticPriority(groups[rhs], groups[lhs]); };
+  using Heap = std::priority_queue<unsigned, std::vector<unsigned>, decltype(compare)>;
+  Heap ready(compare);
+  DenseMap<size_t, std::unique_ptr<Heap>> bySignature;
+  auto addReady = [&](unsigned index) {
+    TransferGroup& group = groups[index];
+    if (group.ready || group.scheduled)
+      return;
+    group.ready = true;
+    ready.push(index);
+    if (group.firstSignature) {
+      auto& bucket = bySignature[hashSignature(*group.firstSignature)];
+      if (!bucket)
+        bucket = std::make_unique<Heap>(compare);
+      bucket->push(index);
+    }
+  };
+  for (auto [index, group] : llvm::enumerate(groups))
+    if (group.unsatisfied == 0)
+      addReady(index);
+
+  std::queue<unsigned> advanceable;
+  auto enqueue = [&](unsigned stream) {
+    StreamProgress& progress = streams[stream];
+    if (!progress.queued && progress.completedStep < plan.stepCounts[stream]
+        && progress.pendingAtStep[progress.completedStep] == 0) {
+      progress.queued = true;
+      advanceable.push(stream);
+    }
+  };
+  for (unsigned stream = 0; stream < streams.size(); ++stream)
+    enqueue(stream);
+  auto advance = [&] {
+    bool changed = false;
+    while (!advanceable.empty()) {
+      unsigned stream = advanceable.front();
+      advanceable.pop();
+      StreamProgress& progress = streams[stream];
+      progress.queued = false;
+      if (progress.completedStep == plan.stepCounts[stream] || progress.pendingAtStep[progress.completedStep] != 0)
+        continue;
+      ++progress.completedStep;
+      changed = true;
+      for (unsigned groupIndex : progress.unlockedAtStep[progress.completedStep])
+        if (--groups[groupIndex].unsatisfied == 0)
+          addReady(groupIndex);
+      enqueue(stream);
+    }
+    return changed;
+  };
+
+  ScheduledCommunicationPlan result;
+  unsigned finishedGroups = 0;
+  while (finishedGroups != groups.size()) {
+    bool progressed = advance();
+    std::optional<unsigned> chosen;
+    unsigned bestExtension = 0;
+    if (!result.slices.empty()) {
+      ScheduledTransferSlice& tailSlice = result.slices.back();
+      ExternalTransferFamily& tailFamily = *tailSlice.family;
+      TransferEmissionSignature tail = getTransferEmissionSignature(tailFamily);
+      size_t hash = hashSignature(tail);
+      auto bucket = bySignature.find(hash);
+      SmallVector<unsigned> inspected;
+      while (bucket != bySignature.end() && !bucket->second->empty()) {
+        unsigned candidate = bucket->second->top();
+        bucket->second->pop();
+        TransferGroup& group = groups[candidate];
+        if (!group.ready || group.scheduled)
+          continue;
+        inspected.push_back(candidate);
+        unsigned extension = tailExtension(
+          group, tail, tailSlice.sourceInsertionStep, streams);
+        if (extension > bestExtension
+            || (extension == bestExtension && extension != 0
+                && (!chosen || betterStaticPriority(group, groups[*chosen])))) {
+          chosen = candidate;
+          bestExtension = extension;
+        }
+      }
+      for (unsigned candidate : inspected)
+        if (!chosen || candidate != *chosen)
+          bucket->second->push(candidate);
+    }
+    while (!chosen && !ready.empty()) {
+      unsigned candidate = ready.top();
+      ready.pop();
+      if (groups[candidate].ready && !groups[candidate].scheduled)
+        chosen = candidate;
+    }
+    if (chosen) {
+      TransferGroup& group = groups[*chosen];
+      group.ready = false;
+      group.scheduled = true;
+      ++finishedGroups;
+      for (ScheduledTransferSlice slice : group.ordered) {
+        ExternalTransferFamily& family = *slice.family;
+        size_t end = slice.familyOffset + slice.transferCount;
+        for (size_t begin = slice.familyOffset; begin < end;) {
+          unsigned sourceStep = streams[
+            family.sourceStreams.valueAt(begin)].completedStep;
+          unsigned targetStep = streams[
+            family.targetStreams.valueAt(begin)].completedStep;
+          size_t next = begin + 1;
+          while (next < end
+                 && streams[family.sourceStreams.valueAt(next)].completedStep
+                      == sourceStep
+                 && streams[family.targetStreams.valueAt(next)].completedStep
+                      == targetStep)
+            ++next;
+          result.slices.push_back(
+            {&family, begin, next - begin, sourceStep, targetStep});
+          result.logicalTransferCount += next - begin;
+          begin = next;
+        }
+        for (size_t index = slice.familyOffset; index < end; ++index) {
+          unsigned stream = family.targetStreams.valueAt(index);
+          unsigned& pending = streams[stream].pendingAtStep[group.exchange->consumerStep];
+          --pending;
+          if (pending == 0 && streams[stream].completedStep == group.exchange->consumerStep)
+            enqueue(stream);
+        }
+      }
+      progressed = true;
+    }
+    if (progressed)
+      continue;
+    InFlightDiagnostic diagnostic =
+      funcOp.emitOpError("global communication scheduler made no progress before IR mutation");
+    unsigned reported = 0;
+    for (const TransferGroup& group : groups) {
+      if (group.scheduled || reported++ == 8)
+        continue;
+      diagnostic << "; blocked parent " << group.exchange->exchangeId;
+      auto dependency = llvm::find_if(group.dependencies, [&](StreamThreshold item) {
+        return streams[item.stream].completedStep < item.completedStep;
+      });
+      if (dependency != group.dependencies.end())
+        diagnostic << " stream " << dependency->stream << " requires " << dependency->completedStep;
+    }
+    return failure();
+  }
+  return result;
+}
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,55 @@
+#pragma once
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+
+#include "DeferredCommunicationModel.hpp"
+
+namespace onnx_mlir::spatial {
+
+struct DeferredTransferPlan;
+
+struct ScheduledTransferSlice {
+  ExternalTransferFamily* family = nullptr;
+  size_t familyOffset = 0;
+  size_t transferCount = 0;
+  unsigned sourceInsertionStep = 0;
+  unsigned targetInsertionStep = 0;
+};
+
+struct ScheduledCommunicationPlan {
+  llvm::SmallVector<ScheduledTransferSlice> slices;
+  uint64_t logicalTransferCount = 0;
+};
+
+struct TransferEmissionSignature {
+  ScheduledInfo* scheduled = nullptr;
+  mlir::Value payload;
+  mlir::Type fragmentType;
+  bool hasGraphLane = false;
+  bool sourceIsBatch = false;
+
+  bool operator==(const TransferEmissionSignature& other) const {
+    return scheduled == other.scheduled && payload == other.payload
+        && fragmentType == other.fragmentType
+        && hasGraphLane == other.hasGraphLane
+        && sourceIsBatch == other.sourceIsBatch;
+  }
+};
+
+struct BoundaryCost {
+  unsigned instructionCount = 0;
+  unsigned branchRegions = 0;
+  unsigned absorbedTransfers = 0;
+  unsigned lookupEntries = 0;
+};
+
+TransferEmissionSignature getTransferEmissionSignature(const ExternalTransferFamily& family);
+
+BoundaryCost estimateCanonicalBoundaryCost(
+  mlir::ArrayRef<ScheduledTransferSlice> existingTail,
+  mlir::ArrayRef<ScheduledTransferSlice> candidate);
+
+mlir::FailureOr<ScheduledCommunicationPlan> scheduleDeferredCommunication(mlir::func::FuncOp funcOp,
+                                                                          DeferredTransferPlan& plan);
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,638 @@
+#include "DeferredProjectionAnalysis.hpp"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/IR/Matchers.h"
+
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallSet.h"
+
+#include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
+
+namespace onnx_mlir::spatial {
+using namespace mlir;
+namespace {
+
+static FailureOr<int64_t> getSignedInt64(IntegerAttr value) {
+  return value.getValue().isSignedIntN(64)
+    ? FailureOr<int64_t>(value.getValue().getSExtValue())
+    : FailureOr<int64_t>(failure());
+}
+
+static FailureOr<int64_t> evaluate(
+    Value value, const StaticIndexEnvironment &environment,
+    llvm::SmallDenseSet<Value, 16> &visiting);
+
+static FailureOr<int64_t> evaluateDenseExtract(
+    tensor::ExtractOp extract, const StaticIndexEnvironment &environment,
+    llvm::SmallDenseSet<Value, 16> &visiting) {
+  auto constant = extract.getTensor().getDefiningOp<arith::ConstantOp>();
+  auto elements = constant
+    ? dyn_cast<DenseIntElementsAttr>(constant.getValue())
+    : DenseIntElementsAttr();
+  auto type = elements
+    ? dyn_cast<RankedTensorType>(elements.getType()) : RankedTensorType();
+  if (!elements || !type || !type.hasStaticShape()
+      || extract.getIndices().size() != static_cast<size_t>(type.getRank()))
+    return failure();
+  int64_t linear = 0;
+  for (auto [index, dim] : llvm::zip(extract.getIndices(), type.getShape())) {
+    auto folded = evaluate(index, environment, visiting);
+    if (failed(folded) || *folded < 0 || *folded >= dim
+        || llvm::MulOverflow(linear, dim, linear)
+        || llvm::AddOverflow(linear, *folded, linear))
+      return failure();
+  }
+  APInt value = elements.getValues<APInt>()[linear];
+  return value.isSignedIntN(64)
+    ? FailureOr<int64_t>(value.getSExtValue())
+    : FailureOr<int64_t>(failure());
+}
+
+static FailureOr<int64_t> evaluate(
+    Value value, const StaticIndexEnvironment &environment,
+    llvm::SmallDenseSet<Value, 16> &visiting) {
+  if (auto it = environment.bindings.find(value);
+      it != environment.bindings.end())
+    return it->second;
+  Attribute constant;
+  if (matchPattern(value, m_Constant(&constant)))
+    if (auto integer = dyn_cast_or_null<IntegerAttr>(constant))
+      return getSignedInt64(integer);
+  if (isa<BlockArgument>(value) || !value.getDefiningOp()
+      || !visiting.insert(value).second)
+    return failure();
+  if (auto extract = value.getDefiningOp<tensor::ExtractOp>()) {
+    auto result = evaluateDenseExtract(extract, environment, visiting);
+    visiting.erase(value);
+    return result;
+  }
+  Operation *op = value.getDefiningOp();
+  if (op->getNumRegions() != 0 || !isPureIndexComputationOp(op)) {
+    visiting.erase(value);
+    return failure();
+  }
+  SmallVector<Attribute> operands;
+  Builder builder(op->getContext());
+  for (Value operand : op->getOperands()) {
+    auto folded = evaluate(operand, environment, visiting);
+    if (failed(folded)) {
+      visiting.erase(value);
+      return failure();
+    }
+    operands.push_back(builder.getIntegerAttr(operand.getType(), *folded));
+  }
+  SmallVector<OpFoldResult> results;
+  if (failed(op->fold(operands, results)) || results.size() != 1) {
+    visiting.erase(value);
+    return failure();
+  }
+  FailureOr<int64_t> result = failure();
+  if (auto attr = dyn_cast<Attribute>(results.front())) {
+    if (auto integer = dyn_cast<IntegerAttr>(attr))
+      result = getSignedInt64(integer);
+  } else if (auto folded = dyn_cast<Value>(results.front())) {
+    result = evaluate(folded, environment, visiting);
+  }
+  visiting.erase(value);
+  return result;
+}
+
+static FailureOr<std::optional<unsigned>> sourceArgument(
+    Value value, SpatDeferredCommunicationOp deferred,
+    const StaticIndexEnvironment &environment) {
+  while (auto cast = value.getDefiningOp<tensor::CastOp>())
+    value = cast.getSource();
+  if (auto argument = dyn_cast<BlockArgument>(value);
+      argument && argument.getOwner() == &deferred.getBody().front()
+      && argument.getArgNumber() < deferred.getSources().size())
+    return std::optional<unsigned>(argument.getArgNumber());
+  auto result = dyn_cast<OpResult>(value);
+  auto selection = result
+    ? dyn_cast<scf::IndexSwitchOp>(result.getOwner()) : scf::IndexSwitchOp();
+  if (!selection || result.getResultNumber() != 0
+      || selection.getNumResults() != 1)
+    return std::optional<unsigned>();
+  auto selector = evaluateDeferredIndex(selection.getArg(), environment);
+  if (failed(selector))
+    return failure();
+  Region *selected = &selection.getDefaultRegion();
+  for (auto [caseValue, region] :
+       llvm::zip(selection.getCases(), selection.getCaseRegions()))
+    if (caseValue == *selector) {
+      selected = &region;
+      break;
+    }
+  auto yield = selected->hasOneBlock()
+    ? dyn_cast<scf::YieldOp>(selected->front().getTerminator())
+    : scf::YieldOp();
+  return yield && yield.getResults().size() == 1
+    ? sourceArgument(yield.getResults().front(), deferred, environment)
+    : FailureOr<std::optional<unsigned>>(failure());
+}
+
+static void collectSourceArguments(Value value,
+                                   SpatDeferredCommunicationOp deferred,
+                                   llvm::SmallSet<unsigned, 4> &indices) {
+  while (auto cast = value.getDefiningOp<tensor::CastOp>())
+    value = cast.getSource();
+  if (auto argument = dyn_cast<BlockArgument>(value)) {
+    if (argument.getOwner() == &deferred.getBody().front()
+        && argument.getArgNumber() < deferred.getSources().size())
+      indices.insert(argument.getArgNumber());
+    return;
+  }
+  auto result = dyn_cast<OpResult>(value);
+  auto selection = result
+    ? dyn_cast<scf::IndexSwitchOp>(result.getOwner()) : scf::IndexSwitchOp();
+  if (!selection || result.getResultNumber() != 0)
+    return;
+  for (Region &region : selection.getCaseRegions())
+    collectSourceArguments(
+      cast<scf::YieldOp>(region.front().getTerminator()).getResults().front(),
+      deferred, indices);
+  collectSourceArguments(
+    cast<scf::YieldOp>(selection.getDefaultRegion().front().getTerminator())
+      .getResults().front(), deferred, indices);
+}
+
+static Value getEnclosingScheduledLane(
+    SpatDeferredCommunicationOp deferred,
+    SpatScheduledComputeBatch scheduled) {
+  Block *block = deferred->getBlock();
+  while (block && block->getParentOp() != scheduled) {
+    Operation *parent = block->getParentOp();
+    block = parent ? parent->getBlock() : nullptr;
+  }
+  return block && !block->empty() ? block->getArgument(0) : Value();
+}
+
+static bool isAllowedStaticIndexExpression(
+    Value value, Value scheduledLane,
+    llvm::SmallDenseSet<Value, 16> &visiting) {
+  if (value == scheduledLane)
+    return true;
+  Attribute constant;
+  if (matchPattern(value, m_Constant(&constant)))
+    return true;
+  if (isa<BlockArgument>(value) || !value.getDefiningOp()
+      || !visiting.insert(value).second)
+    return false;
+  Operation *op = value.getDefiningOp();
+  bool allowed = op->getNumRegions() == 0
+      && (isPureIndexComputationOp(op) || isCompileTimeOp(op))
+      && llvm::all_of(op->getOperands(), [&](Value operand) {
+           return isAllowedStaticIndexExpression(
+             operand, scheduledLane, visiting);
+         });
+  visiting.erase(value);
+  return allowed;
+}
+
+static bool isAllowedStaticIndexExpression(Value value,
+                                           Value scheduledLane) {
+  llvm::SmallDenseSet<Value, 16> visiting;
+  return isAllowedStaticIndexExpression(value, scheduledLane, visiting);
+}
+
+static bool originatesFromDeferredSource(
+    Value value, SpatDeferredCommunicationOp deferred,
+    llvm::SmallDenseSet<Value, 16> &visited) {
+  if (!visited.insert(value).second)
+    return false;
+  if (auto argument = dyn_cast<BlockArgument>(value))
+    return argument.getOwner() == &deferred.getBody().front()
+        && argument.getArgNumber() < deferred.getSources().size();
+  Operation *op = value.getDefiningOp();
+  return op && (isa<scf::IndexSwitchOp>(op)
+      || llvm::any_of(op->getOperands(), [&](Value operand) {
+           return originatesFromDeferredSource(operand, deferred, visited);
+         }));
+}
+
+static bool originatesFromDeferredSource(
+    Value value, SpatDeferredCommunicationOp deferred) {
+  llvm::SmallDenseSet<Value, 16> visited;
+  return originatesFromDeferredSource(value, deferred, visited);
+}
+
+static LogicalResult verifyCanonicalSourceSelector(
+    scf::IndexSwitchOp selection, SpatDeferredCommunicationOp deferred,
+    Value scheduledLane, int64_t laneCount) {
+  if (selection->getBlock() != &deferred.getBody().front()
+      || selection.getNumResults() != 1
+      || !selection.getArg().getType().isIndex()
+      || !isAllowedStaticIndexExpression(selection.getArg(), scheduledLane))
+    return selection.emitOpError(
+      "is not a canonical deferred source selector");
+  if (laneCount < 2
+      || selection.getCases().size() != static_cast<size_t>(laneCount - 1))
+    return selection.emitOpError(
+      "must cover every non-default scheduled lane");
+  for (auto [index, caseValue] : llvm::enumerate(selection.getCases()))
+    if (caseValue != static_cast<int64_t>(index))
+      return selection.emitOpError(
+        "must use consecutive scheduled-lane cases starting at zero");
+  auto verifyRegion = [&](Region &region) -> LogicalResult {
+    auto yield = region.hasOneBlock()
+      ? dyn_cast<scf::YieldOp>(region.front().getTerminator()) : scf::YieldOp();
+    if (!yield || yield.getResults().size() != 1)
+      return selection.emitOpError(
+        "source-selector region must yield one result");
+    for (Operation &op : region.front().without_terminator())
+      if (!isa<tensor::CastOp>(op))
+        return selection.emitOpError(
+          "source-selector regions may contain only tensor.cast");
+    llvm::SmallSet<unsigned, 4> indices;
+    collectSourceArguments(yield.getResults().front(), deferred, indices);
+    return success(indices.size() == 1);
+  };
+  for (Region &region : selection.getCaseRegions())
+    if (failed(verifyRegion(region)))
+      return failure();
+  return verifyRegion(selection.getDefaultRegion());
+}
+
+static bool isInsideDeferredLoop(Operation *op,
+                                 SpatDeferredCommunicationOp deferred) {
+  for (Operation *parent = op->getParentOp(); parent && parent != deferred;
+       parent = parent->getParentOp())
+    if (isa<scf::ForOp>(parent))
+      return true;
+  return false;
+}
+
+static bool haveLeadingUnitDifference(RankedTensorType larger,
+                                      RankedTensorType smaller) {
+  return larger && smaller && larger.hasStaticShape() && smaller.hasStaticShape()
+      && larger.getRank() == smaller.getRank() + 1
+      && larger.getDimSize(0) == 1
+      && larger.getShape().drop_front() == smaller.getShape();
+}
+
+static std::optional<DeferredAssemblySourceTransform> getSourceTransform(
+    RankedTensorType publication, RankedTensorType inserted) {
+  if (publication == inserted)
+    return DeferredAssemblySourceTransform::Identity;
+  if (haveLeadingUnitDifference(inserted, publication))
+    return DeferredAssemblySourceTransform::AddLeadingUnitDimension;
+  if (haveLeadingUnitDifference(publication, inserted))
+    return DeferredAssemblySourceTransform::RemoveLeadingUnitDimension;
+  return std::nullopt;
+}
+
+static FailureOr<std::optional<DeferredInsertAssemblyTemplate>>
+analyzeInsertAssembly(const DeferredProgramTemplate &program) {
+  auto finalInsert = program.yieldedValue.getDefiningOp<tensor::InsertSliceOp>();
+  if (!finalInsert || program.leaves.empty())
+    return std::optional<DeferredInsertAssemblyTemplate>();
+  DenseMap<Value, unsigned> leafByRoot;
+  for (auto [index, leaf] : llvm::enumerate(program.leaves))
+    if (!leafByRoot.try_emplace(leaf.replacementRoot, index).second)
+      return std::optional<DeferredInsertAssemblyTemplate>();
+
+  SmallPtrSet<Operation *, 32> consumed;
+  SmallVector<DeferredInsertAssemblyEntryTemplate> reverseEntries;
+  llvm::SmallDenseSet<unsigned, 16> insertedLeaves;
+  Value current = program.yieldedValue;
+  while (auto insert = current.getDefiningOp<tensor::InsertSliceOp>()) {
+    Value source = insert.getSource();
+    Operation *shape = source.getDefiningOp();
+    if (isa_and_nonnull<tensor::CollapseShapeOp, tensor::ExpandShapeOp>(shape))
+      source = shape->getOperand(0);
+    if (source.getDefiningOp<tensor::CollapseShapeOp>()
+        || source.getDefiningOp<tensor::ExpandShapeOp>())
+      return std::optional<DeferredInsertAssemblyTemplate>();
+    auto leaf = leafByRoot.find(source);
+    if (leaf == leafByRoot.end()
+        || !insertedLeaves.insert(leaf->second).second)
+      return std::optional<DeferredInsertAssemblyTemplate>();
+    auto publicationType = dyn_cast<RankedTensorType>(source.getType());
+    auto insertedType = dyn_cast<RankedTensorType>(insert.getSourceType());
+    auto transform = getSourceTransform(publicationType, insertedType);
+    if (!transform)
+      return std::optional<DeferredInsertAssemblyTemplate>();
+    DeferredInsertAssemblyEntryTemplate entry;
+    entry.coordinate = {leaf->second, 0};
+    entry.sourceTransform = *transform;
+    entry.sourceType = insertedType;
+    entry.targetGeometry = {SmallVector<OpFoldResult>(insert.getMixedOffsets()),
+                            SmallVector<OpFoldResult>(insert.getMixedSizes()),
+                            SmallVector<OpFoldResult>(insert.getMixedStrides())};
+    reverseEntries.push_back(std::move(entry));
+    consumed.insert(insert);
+    if (shape)
+      consumed.insert(shape);
+    current = insert.getDest();
+  }
+  auto initial = current.getDefiningOp<tensor::EmptyOp>();
+  auto resultType = dyn_cast<RankedTensorType>(program.yieldedValue.getType());
+  if (!initial || !initial.getDynamicSizes().empty() || !resultType
+      || initial.getType() != resultType
+      || insertedLeaves.size() != program.leaves.size())
+    return std::optional<DeferredInsertAssemblyTemplate>();
+  consumed.insert(initial);
+  if (llvm::any_of(program.residualOps,
+                   [&](Operation *op) { return !consumed.contains(op); }))
+    return std::optional<DeferredInsertAssemblyTemplate>();
+  DeferredInsertAssemblyTemplate assembly;
+  assembly.initialValue = initial;
+  assembly.resultType = resultType;
+  assembly.entries.assign(reverseEntries.rbegin(), reverseEntries.rend());
+  return std::optional<DeferredInsertAssemblyTemplate>(std::move(assembly));
+}
+
+} // namespace
+
+LogicalResult verifyDeferredProgramContract(
+    SpatDeferredCommunicationOp deferred) {
+  if (!deferred.getBody().hasOneBlock())
+    return deferred.emitOpError(
+      "deferred program must have exactly one body block");
+  Block &body = deferred.getBody().front();
+  auto terminator = dyn_cast<SpatYieldOp>(body.getTerminator());
+  if (!terminator || terminator.getOutputs().size() != 1)
+    return deferred.emitOpError(
+      "deferred program must have exactly one yielded value");
+  Value scheduledLane;
+  int64_t laneCount = 1;
+  if (auto scheduled = deferred->getParentOfType<SpatScheduledComputeBatch>()) {
+    scheduledLane = getEnclosingScheduledLane(deferred, scheduled);
+    laneCount = scheduled.getLaneCount();
+  }
+  bool invalid = false;
+  deferred.getBody().walk([&](Operation *op) {
+    auto reject = [&](StringRef message) {
+      op->emitOpError(message);
+      invalid = true;
+    };
+    if (invalid || isa<SpatYieldOp, scf::YieldOp>(op))
+      return;
+    if (auto selection = dyn_cast<scf::IndexSwitchOp>(op)) {
+      if (failed(verifyCanonicalSourceSelector(
+            selection, deferred, scheduledLane, laneCount)))
+        invalid = true;
+      return;
+    }
+    if (auto loop = dyn_cast<scf::ForOp>(op)) {
+      for (Value bound : {loop.getLowerBound(), loop.getUpperBound(),
+                          loop.getStep()})
+        if (!isAllowedStaticIndexExpression(bound, scheduledLane))
+          return reject("has a non-static deferred loop bound");
+      for (int64_t lane = 0; lane < laneCount; ++lane) {
+        StaticIndexEnvironment environment;
+        if (scheduledLane)
+          environment.bindings[scheduledLane] = lane;
+        auto lower = evaluateDeferredIndex(loop.getLowerBound(), environment);
+        auto upper = evaluateDeferredIndex(loop.getUpperBound(), environment);
+        auto step = evaluateDeferredIndex(loop.getStep(), environment);
+        if (failed(lower) || failed(upper) || failed(step) || *step <= 0)
+          return reject("has a loop bound that does not specialize");
+      }
+      return;
+    }
+    if (op->getNumRegions() != 0)
+      return reject("contains an unsupported region operation");
+    if (!isShapingOnlyOp(op) && !isCompileTimeOp(op)
+        && !isPureIndexComputationOp(op))
+      return reject("contains an unsupported deferred operation");
+    for (Value operand : op->getOperands())
+      if (isInsideDeferredLoop(op, deferred)
+          && originatesFromDeferredSource(operand, deferred))
+        return reject("projects a deferred source inside a residual loop");
+  });
+  return success(!invalid);
+}
+
+FailureOr<int64_t> evaluateDeferredIndex(
+    Value value, const StaticIndexEnvironment &environment) {
+  llvm::SmallDenseSet<Value, 16> visiting;
+  return evaluate(value, environment, visiting);
+}
+
+FailureOr<int64_t> evaluateDeferredIndex(
+    OpFoldResult value, const StaticIndexEnvironment &environment) {
+  if (auto attr = dyn_cast<Attribute>(value))
+    if (auto integer = dyn_cast<IntegerAttr>(attr))
+      return getSignedInt64(integer);
+  if (auto dynamic = dyn_cast<Value>(value))
+    return evaluateDeferredIndex(dynamic, environment);
+  return failure();
+}
+
+FailureOr<SmallVector<unsigned>> getPossibleDeferredSourceOperandIndices(
+    Value sourceRoot, SpatDeferredCommunicationOp deferred) {
+  llvm::SmallSet<unsigned, 4> indices;
+  collectSourceArguments(sourceRoot, deferred, indices);
+  if (indices.empty())
+    return failure();
+  return SmallVector<unsigned>(indices.begin(), indices.end());
+}
+
+DeferredLaneValueEvaluator::DeferredLaneValueEvaluator(
+    const DeferredProgramTemplate &program, unsigned laneCount)
+  : program(program), laneCount(laneCount) {}
+
+FailureOr<StaticIntSequence> DeferredLaneValueEvaluator::evaluate(Value value) {
+  if (auto it = values.find(value); it != values.end())
+    return it->second;
+  if (value == program.scheduledLane) {
+    StaticIntSequence result = StaticIntSequence::affine(0, 1, laneCount);
+    values.try_emplace(value, result);
+    return result;
+  }
+  SmallVector<int64_t> evaluated;
+  evaluated.reserve(laneCount);
+  for (unsigned lane = 0; lane < laneCount; ++lane) {
+    StaticIndexEnvironment environment;
+    if (program.scheduledLane)
+      environment.bindings[program.scheduledLane] = lane;
+    auto result = evaluateDeferredIndex(value, environment);
+    if (failed(result))
+      return failure();
+    evaluated.push_back(*result);
+  }
+  StaticIntSequence result = StaticIntSequence::fromValues(evaluated);
+  values.try_emplace(value, result);
+  return result;
+}
+
+FailureOr<StaticIntSequence> DeferredLaneValueEvaluator::evaluate(
+    OpFoldResult value) {
+  if (auto attr = dyn_cast<Attribute>(value)) {
+    auto integer = dyn_cast<IntegerAttr>(attr);
+    if (!integer)
+      return failure();
+    auto number = getSignedInt64(integer);
+    return succeeded(number)
+      ? FailureOr<StaticIntSequence>(
+          StaticIntSequence::uniform(*number, laneCount))
+      : FailureOr<StaticIntSequence>(failure());
+  }
+  return evaluate(cast<Value>(value));
+}
+
+FailureOr<StaticIntSequence>
+DeferredLaneValueEvaluator::resolveSourceOperandIndices(Value sourceRoot) {
+  if (auto it = sourceOperands.find(sourceRoot); it != sourceOperands.end())
+    return it->second;
+  SmallVector<int64_t> indices;
+  indices.reserve(laneCount);
+  for (unsigned lane = 0; lane < laneCount; ++lane) {
+    StaticIndexEnvironment environment;
+    if (program.scheduledLane)
+      environment.bindings[program.scheduledLane] = lane;
+    auto index = sourceArgument(sourceRoot, program.deferred, environment);
+    if (failed(index) || !*index)
+      return failure();
+    indices.push_back(**index);
+  }
+  StaticIntSequence result = StaticIntSequence::fromValues(indices);
+  sourceOperands.try_emplace(sourceRoot, result);
+  return result;
+}
+
+FailureOr<DeferredProgramTemplate> analyzeDeferredProgramTemplate(
+    SpatDeferredCommunicationOp deferred) {
+  Block &body = deferred.getBody().front();
+  auto yield = dyn_cast<SpatYieldOp>(body.getTerminator());
+  if (!yield || yield.getOutputs().size() != 1)
+    return deferred.emitOpError(
+      "requires one deferred yielded value"), failure();
+  DeferredProgramTemplate program;
+  program.deferred = deferred;
+  program.yieldedValue = yield.getOutputs().front();
+  if (auto scheduled = deferred->getParentOfType<SpatScheduledComputeBatch>())
+    program.scheduledLane = getEnclosingScheduledLane(deferred, scheduled);
+
+  llvm::SmallDenseSet<Value, 32> visited;
+  std::function<LogicalResult(Value)> visit = [&](Value value) -> LogicalResult {
+    if (!visited.insert(value).second)
+      return success();
+    if (auto slice = value.getDefiningOp<tensor::ExtractSliceOp>()) {
+      auto sources = getPossibleDeferredSourceOperandIndices(
+        slice.getSource(), deferred);
+      bool graphProjection = succeeded(sources)
+        && llvm::all_of(*sources, [&](unsigned index) {
+             auto result = dyn_cast<OpResult>(deferred.getSources()[index]);
+             return result && isa<SpatGraphComputeBatch>(result.getOwner());
+           });
+      if (graphProjection) {
+        DeferredProjectionLeafTemplate leaf;
+        leaf.form = DeferredLeafForm::GraphBatchProjection;
+        leaf.sourceRoot = slice.getSource();
+        leaf.replacementRoot = value;
+        leaf.leadingProjection = slice;
+        leaf.leadingGeometry = {
+          SmallVector<OpFoldResult>(slice.getMixedOffsets()),
+          SmallVector<OpFoldResult>(slice.getMixedSizes()),
+          SmallVector<OpFoldResult>(slice.getMixedStrides())};
+        leaf.innerGeometry = {
+          SmallVector<OpFoldResult>(
+            ArrayRef(slice.getMixedOffsets()).drop_front()),
+          SmallVector<OpFoldResult>(
+            ArrayRef(slice.getMixedSizes()).drop_front()),
+          SmallVector<OpFoldResult>(
+            ArrayRef(slice.getMixedStrides()).drop_front())};
+        leaf.reconstructedType = cast<RankedTensorType>(value.getType());
+        program.leaves.push_back(std::move(leaf));
+        return success();
+      }
+    }
+    if (succeeded(getPossibleDeferredSourceOperandIndices(value, deferred))) {
+      auto type = dyn_cast<RankedTensorType>(value.getType());
+      if (!type)
+        return deferred.emitOpError(
+          "deferred source is not a ranked tensor");
+      program.leaves.push_back({DeferredLeafForm::DirectSource, value, value,
+                                {}, {}, {}, type});
+      return success();
+    }
+    if (value.getType().isIndex() || isa<IntegerType>(value.getType()))
+      return success();
+    Operation *op = value.getDefiningOp();
+    if (!op || (op->getBlock() != &body && !isa<scf::ForOp>(op)))
+      return deferred.emitOpError(
+        "deferred residual escapes its verified body");
+    if (auto loop = dyn_cast<scf::ForOp>(op)) {
+      for (Value init : loop.getInitArgs())
+        if (failed(visit(init)))
+          return failure();
+    } else {
+      for (Value operand : op->getOperands())
+        if (failed(visit(operand)))
+          return failure();
+    }
+    program.residualOps.push_back(op);
+    return success();
+  };
+  if (failed(visit(program.yieldedValue)))
+    return failure();
+  auto assembly = analyzeInsertAssembly(program);
+  if (failed(assembly))
+    return failure();
+  program.insertAssembly = std::move(*assembly);
+  return program;
+}
+
+FailureOr<const GraphBatchPublicationMap *> getGraphBatchPublicationMap(
+    SpatGraphComputeBatch graphBatch, unsigned resultIndex,
+    GraphBatchPublicationCache &cache) {
+  GraphBatchPublicationKey key {graphBatch.getOperation(), resultIndex};
+  if (auto it = cache.find(key); it != cache.end())
+    return &it->second;
+  auto resultType = dyn_cast<RankedTensorType>(
+    graphBatch.getResult(resultIndex).getType());
+  auto output = graphBatch.getOutputArgument(resultIndex);
+  auto lane = graphBatch.getLaneArgument();
+  if (!resultType || !output || !lane || resultType.getRank() == 0)
+    return graphBatch.emitOpError(
+      "graph batch publication is malformed"), failure();
+  tensor::ParallelInsertSliceOp publication;
+  auto parallel = dyn_cast<SpatInParallelOp>(
+    graphBatch.getBody().front().getTerminator());
+  if (!parallel)
+    return graphBatch.emitOpError(
+      "graph batch lacks publication region"), failure();
+  for (Operation &op : parallel.getRegion().front())
+    if (auto insert = dyn_cast<tensor::ParallelInsertSliceOp>(op);
+        insert && insert.getDest() == *output)
+      publication = insert;
+  auto fragment = publication
+    ? dyn_cast<RankedTensorType>(publication.getSource().getType())
+    : RankedTensorType();
+  if (!publication || !fragment
+      || resultType.getRank() != fragment.getRank() + 1)
+    return graphBatch.emitOpError(
+      "graph publication fragment type is invalid"), failure();
+  GraphBatchPublicationMap map;
+  map.physicalResultType = resultType;
+  map.publicationFragmentType = fragment;
+  map.graphLaneToPhysicalSlot.resize(graphBatch.getLaneCount(), -1);
+  map.physicalSlotToGraphLane.resize(resultType.getDimSize(0), -1);
+  for (int64_t index = 0; index < graphBatch.getLaneCount(); ++index) {
+    StaticIndexEnvironment environment;
+    environment.bindings[*lane] = index;
+    auto slot = evaluateDeferredIndex(
+      publication.getMixedOffsets().front(), environment);
+    auto size = evaluateDeferredIndex(
+      publication.getMixedSizes().front(), environment);
+    auto stride = evaluateDeferredIndex(
+      publication.getMixedStrides().front(), environment);
+    if (failed(slot) || failed(size) || failed(stride)
+        || *size != 1 || *stride != 1 || *slot < 0
+        || *slot >= resultType.getDimSize(0)
+        || map.physicalSlotToGraphLane[*slot] != -1)
+      return graphBatch.emitOpError(
+        "graph publication leading geometry is invalid"), failure();
+    map.graphLaneToPhysicalSlot[index] = *slot;
+    map.physicalSlotToGraphLane[*slot] = index;
+  }
+  if (llvm::is_contained(map.physicalSlotToGraphLane, -1))
+    return graphBatch.emitOpError(
+      "graph publication has a missing physical slot"), failure();
+  return &cache.try_emplace(key, std::move(map)).first->second;
+}
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,87 @@
+#pragma once
+
+#include "DeferredCommunicationModel.hpp"
+
+#include "llvm/ADT/DenseMap.h"
+
+namespace onnx_mlir::spatial {
+
+struct StaticIndexEnvironment {
+  llvm::DenseMap<mlir::Value, int64_t> bindings;
+};
+
+mlir::FailureOr<int64_t> evaluateDeferredIndex(
+  mlir::Value value, const StaticIndexEnvironment &environment);
+mlir::FailureOr<int64_t> evaluateDeferredIndex(
+  mlir::OpFoldResult value, const StaticIndexEnvironment &environment);
+
+mlir::LogicalResult verifyDeferredProgramContract(
+  SpatDeferredCommunicationOp deferred);
+
+mlir::FailureOr<DeferredProgramTemplate> analyzeDeferredProgramTemplate(
+  SpatDeferredCommunicationOp deferred);
+
+class DeferredLaneValueEvaluator {
+public:
+  DeferredLaneValueEvaluator(const DeferredProgramTemplate &program,
+                             unsigned laneCount);
+
+  mlir::FailureOr<StaticIntSequence> evaluate(mlir::Value value);
+  mlir::FailureOr<StaticIntSequence> evaluate(mlir::OpFoldResult value);
+  mlir::FailureOr<StaticIntSequence> resolveSourceOperandIndices(
+    mlir::Value sourceRoot);
+
+private:
+  const DeferredProgramTemplate &program;
+  unsigned laneCount;
+  llvm::DenseMap<mlir::Value, StaticIntSequence> values;
+  llvm::DenseMap<mlir::Value, StaticIntSequence> sourceOperands;
+};
+
+mlir::FailureOr<llvm::SmallVector<unsigned>>
+getPossibleDeferredSourceOperandIndices(
+  mlir::Value sourceRoot, SpatDeferredCommunicationOp deferred);
+
+struct GraphBatchPublicationMap {
+  mlir::RankedTensorType physicalResultType;
+  mlir::RankedTensorType publicationFragmentType;
+  llvm::SmallVector<int64_t> graphLaneToPhysicalSlot;
+  llvm::SmallVector<int64_t> physicalSlotToGraphLane;
+};
+
+struct GraphBatchPublicationKey {
+  mlir::Operation *graphBatch = nullptr;
+  unsigned resultIndex = 0;
+  bool operator==(const GraphBatchPublicationKey &other) const {
+    return graphBatch == other.graphBatch && resultIndex == other.resultIndex;
+  }
+};
+
+using GraphBatchPublicationCache =
+  llvm::DenseMap<GraphBatchPublicationKey, GraphBatchPublicationMap>;
+
+mlir::FailureOr<const GraphBatchPublicationMap *> getGraphBatchPublicationMap(
+  SpatGraphComputeBatch graphBatch, unsigned resultIndex,
+  GraphBatchPublicationCache &cache);
+
+} // namespace onnx_mlir::spatial
+
+namespace llvm {
+template <> struct DenseMapInfo<onnx_mlir::spatial::GraphBatchPublicationKey> {
+  static onnx_mlir::spatial::GraphBatchPublicationKey getEmptyKey() {
+    return {DenseMapInfo<mlir::Operation *>::getEmptyKey(), 0};
+  }
+  static onnx_mlir::spatial::GraphBatchPublicationKey getTombstoneKey() {
+    return {DenseMapInfo<mlir::Operation *>::getTombstoneKey(), 0};
+  }
+  static unsigned getHashValue(
+      const onnx_mlir::spatial::GraphBatchPublicationKey &key) {
+    return hash_combine(key.graphBatch, key.resultIndex);
+  }
+  static bool isEqual(
+      const onnx_mlir::spatial::GraphBatchPublicationKey &lhs,
+      const onnx_mlir::spatial::GraphBatchPublicationKey &rhs) {
+    return lhs == rhs;
+  }
+};
+} // namespace llvm

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

@@ -0,0 +1,306 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/IRMapping.h"
+
+#include "DeferredResultRealization.hpp"
+#include "DeferredBoundaryRealization.hpp"
+#include "DeferredProjectionAnalysis.hpp"
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
+
+namespace onnx_mlir::spatial {
+using namespace mlir;
+namespace {
+
+static Value getSequencePosition(
+  const LaneSet& familyLanes, Value lane, Operation* anchor, DeferredEmissionContext& context, Location loc) {
+  unsigned begin = familyLanes.intervals().front().begin;
+  if (!lane)
+    return context.constants.getIndex(0);
+  return affineAddConst(context.rewriter, loc, lane, -static_cast<int64_t>(begin), anchor);
+}
+
+static OpFoldResult materializeGeometryValue(
+  const StaticIntSequence& sequence, Value lane, Operation* anchor, DeferredEmissionContext& context, Location loc) {
+  if (sequence.getKind() == StaticIntSequenceKind::Uniform)
+    return context.rewriter.getIndexAttr(sequence.valueAt(0));
+  return emitStaticIntLookup(sequence, lane, anchor, context.constants, context.rewriter, loc);
+}
+
+static MixedSliceGeometry materializeGeometry(const DeferredResultPlan::SliceGeometry& geometry,
+                                              Value lane,
+                                              Operation* anchor,
+                                              DeferredEmissionContext& context,
+                                              Location loc) {
+  MixedSliceGeometry result;
+  for (const StaticIntSequence& value : geometry.offsets)
+    result.offsets.push_back(materializeGeometryValue(value, lane, anchor, context, loc));
+  for (const StaticIntSequence& value : geometry.sizes)
+    result.sizes.push_back(materializeGeometryValue(value, lane, anchor, context, loc));
+  for (const StaticIntSequence& value : geometry.strides)
+    result.strides.push_back(materializeGeometryValue(value, lane, anchor, context, loc));
+  return result;
+}
+
+static RequirementFamily*
+findRequirement(const DeferredResultPlan& plan, RequirementCoordinate coordinate, unsigned representativeLane) {
+  RequirementFamily* match = nullptr;
+  for (RequirementFamily* requirement : plan.requirements)
+    if (requirement->coordinate == coordinate && requirement->targetLanes.contains(representativeLane)) {
+      assert(!match && "requirement coordinate is not unique");
+      match = requirement;
+    }
+  return match;
+}
+
+static bool isIdentityGeometry(const DeferredResultPlan::SliceGeometry& geometry, RankedTensorType type) {
+  if (geometry.offsets.size() != static_cast<size_t>(type.getRank()))
+    return false;
+  for (auto [dimension, values] :
+       llvm::enumerate(llvm::zip_equal(geometry.offsets, geometry.sizes, geometry.strides))) {
+    const StaticIntSequence& offset = std::get<0>(values);
+    const StaticIntSequence& size = std::get<1>(values);
+    const StaticIntSequence& stride = std::get<2>(values);
+    if (offset.getKind() != StaticIntSequenceKind::Uniform || size.getKind() != StaticIntSequenceKind::Uniform
+        || stride.getKind() != StaticIntSequenceKind::Uniform || offset.valueAt(0) != 0
+        || size.valueAt(0) != type.getDimSize(dimension) || stride.valueAt(0) != 1)
+      return false;
+  }
+  return true;
+}
+
+static FailureOr<Value> applyInnerGeometry(Value fragment,
+                                           const DeferredProjectionLeafTemplate& leaf,
+                                           const DeferredResultPlan::SliceGeometry& geometry,
+                                           Value lane,
+                                           DeferredEmissionContext& context) {
+  if (leaf.form != DeferredLeafForm::GraphBatchProjection)
+    return fragment;
+  auto fragmentType = dyn_cast<RankedTensorType>(fragment.getType());
+  if (!fragmentType || geometry.offsets.size() != static_cast<size_t>(fragmentType.getRank()))
+    return failure();
+  if (isIdentityGeometry(geometry, fragmentType))
+    return fragment;
+  SmallVector<int64_t> shape(leaf.reconstructedType.getShape().drop_front());
+  auto resultType = RankedTensorType::get(shape, fragmentType.getElementType());
+  Value sourceRoot = leaf.sourceRoot;
+  MixedSliceGeometry mixed =
+    materializeGeometry(geometry, lane, sourceRoot.getParentBlock()->getParentOp(), context, sourceRoot.getLoc());
+  return extractMixedSliceOrIdentity(context.rewriter, leaf.sourceRoot.getLoc(), fragment, resultType, mixed);
+}
+
+static FailureOr<Value> reconstructLeaf(const DeferredResultPlan& plan,
+                                        unsigned leafIndex,
+                                        const LaneSet& activeLanes,
+                                        Value lane,
+                                        DeferredEmissionContext& context) {
+  DeferredExchangePlan& exchange = *plan.exchange;
+  const DeferredProjectionLeafTemplate& leaf = exchange.program.leaves[leafIndex];
+  unsigned representative = activeLanes.intervals().front().begin;
+  if (Value assembled = context.projectionAssemblies.lookup({&exchange, leafIndex})) {
+    for (RequirementFamily* requirement : plan.requirements) {
+      if (requirement->coordinate.leafIndex != leafIndex || !requirement->targetLanes.contains(representative))
+        continue;
+      Value local = context.receives.lookup(requirement);
+      if (!local)
+        continue;
+      auto shaped = applyInnerGeometry(local, leaf, plan.innerGeometry[leafIndex], lane, context);
+      auto source = succeeded(shaped)
+                    ? addLeadingUnitTensorDimension(context.rewriter, exchange.deferred.getLoc(), *shaped)
+                    : FailureOr<Value>(failure());
+      if (failed(source))
+        return failure();
+      auto sourceType = cast<RankedTensorType>(source->getType());
+      MixedSliceGeometry geometry;
+      geometry.offsets.assign(sourceType.getRank(), context.rewriter.getIndexAttr(0));
+      geometry.offsets.front() = context.rewriter.getIndexAttr(requirement->coordinate.selectedPosition);
+      for (int64_t dimension : sourceType.getShape())
+        geometry.sizes.push_back(context.rewriter.getIndexAttr(dimension));
+      geometry.strides.assign(sourceType.getRank(), context.rewriter.getIndexAttr(1));
+      assembled = insertMixedSlice(context.rewriter, exchange.deferred.getLoc(), *source, assembled, geometry);
+    }
+    return assembled;
+  }
+  unsigned positionCount = 0;
+  for (RequirementFamily* requirement : plan.requirements)
+    if (requirement->coordinate.leafIndex == leafIndex && requirement->targetLanes.contains(representative))
+      positionCount = std::max(positionCount, requirement->coordinate.selectedPosition + 1);
+  if (positionCount == 0)
+    return failure();
+  SmallVector<Value> expanded;
+  for (unsigned position = 0; position < positionCount; ++position) {
+    RequirementFamily* requirement = findRequirement(plan, {leafIndex, position}, representative);
+    if (!requirement)
+      return failure();
+    auto fragment = materializeDeferredRequirement(*requirement, activeLanes, lane, context);
+    if (failed(fragment))
+      return failure();
+    auto shaped = applyInnerGeometry(*fragment, leaf, plan.innerGeometry[leafIndex], lane, context);
+    if (failed(shaped))
+      return failure();
+    if (positionCount == 1 && shaped->getType() == leaf.reconstructedType)
+      return *shaped;
+    auto value = addLeadingUnitTensorDimension(context.rewriter, exchange.deferred.getLoc(), *shaped);
+    if (failed(value))
+      return failure();
+    expanded.push_back(*value);
+  }
+  if (expanded.size() == 1 && expanded.front().getType() == leaf.reconstructedType)
+    return expanded.front();
+  constexpr size_t maxConcatInputs = 64;
+  while (expanded.size() > 1) {
+    SmallVector<Value> next;
+    next.reserve((expanded.size() + maxConcatInputs - 1) / maxConcatInputs);
+    for (size_t index = 0; index < expanded.size(); index += maxConcatInputs) {
+      ValueRange inputs = ValueRange(expanded).slice(index, std::min(maxConcatInputs, expanded.size() - index));
+      if (inputs.size() == 1) {
+        next.push_back(inputs.front());
+        continue;
+      }
+      auto first = cast<RankedTensorType>(inputs.front().getType());
+      SmallVector<int64_t> shape(first.getShape());
+      shape.front() = 0;
+      for (Value input : inputs)
+        shape.front() += cast<RankedTensorType>(input.getType()).getDimSize(0);
+      auto resultType = RankedTensorType::get(shape, first.getElementType());
+      next.push_back(
+        tensor::ConcatOp::create(context.rewriter, exchange.deferred.getLoc(), resultType, 0, inputs).getResult());
+    }
+    expanded = std::move(next);
+  }
+  return expanded.front().getType() == leaf.reconstructedType ? FailureOr<Value>(expanded.front())
+                                                              : FailureOr<Value>(failure());
+}
+
+} // namespace
+
+FailureOr<DeferredResultPlan>
+buildDeferredResultPlan(DeferredExchangePlan& exchange) {
+  DeferredResultPlan result;
+  result.exchange = &exchange;
+  for (RequirementFamily& requirement : exchange.requirements)
+    result.requirements.push_back(&requirement);
+  DeferredLaneValueEvaluator evaluator(
+    exchange.program, exchange.targetLaneCount);
+  auto buildGeometry = [&](const DeferredSliceTemplate& source,
+                           DeferredResultPlan::SliceGeometry& target) {
+    auto append = [&](ArrayRef<OpFoldResult> values,
+                      SmallVectorImpl<StaticIntSequence>& sequences) {
+      for (OpFoldResult value : values) {
+        auto sequence = evaluator.evaluate(value);
+        if (failed(sequence))
+          return failure();
+        sequences.push_back(*sequence);
+      }
+      return success();
+    };
+    return success(succeeded(append(source.offsets, target.offsets))
+                   && succeeded(append(source.sizes, target.sizes))
+                   && succeeded(append(source.strides, target.strides)));
+  };
+  for (const DeferredProjectionLeafTemplate& leaf : exchange.program.leaves) {
+    DeferredResultPlan::SliceGeometry geometry;
+    if (failed(buildGeometry(leaf.innerGeometry, geometry)))
+      return failure();
+    result.innerGeometry.push_back(std::move(geometry));
+  }
+  if (exchange.program.insertAssembly)
+    for (const DeferredInsertAssemblyEntryTemplate& entry :
+         exchange.program.insertAssembly->entries) {
+      DeferredResultPlan::SliceGeometry geometry;
+      if (failed(buildGeometry(entry.targetGeometry, geometry)))
+        return failure();
+      result.assemblyGeometry.push_back(std::move(geometry));
+    }
+  for (Operation* op : exchange.program.residualOps)
+    op->walk([&](Operation* nested) {
+      for (Value operand : nested->getOperands()) {
+        if ((!operand.getType().isIndex()
+             && !isa<IntegerType>(operand.getType()))
+            || result.residualValues.count(operand))
+          continue;
+        if (auto sequence = evaluator.evaluate(operand); succeeded(sequence))
+          result.residualValues.try_emplace(operand, *sequence);
+      }
+    });
+  return result;
+}
+
+FailureOr<Value> materializeDeferredRequirement(RequirementFamily& requirement,
+                                                const LaneSet& activeLanes,
+                                                Value lane,
+                                                DeferredEmissionContext& context) {
+  if (Value received = context.receives.lookup(&requirement))
+    return received;
+  ProducedValue& producer = *requirement.producer;
+  Value payload = producer.payload;
+  if (payload.getType() == requirement.publicationFragmentType)
+    return payload;
+  auto payloadType = dyn_cast<RankedTensorType>(payload.getType());
+  auto fragmentType = dyn_cast<RankedTensorType>(requirement.publicationFragmentType);
+  if (!payloadType || !fragmentType || !requirement.producerLocalOffsets
+      || payloadType.getRank() != fragmentType.getRank() + 1)
+    return failure();
+  Location loc = requirement.exchange->deferred.getLoc();
+  Value position = getSequencePosition(requirement.targetLanes, lane, requirement.exchange->deferred, context, loc);
+  Value offset = emitStaticIntLookup(*requirement.producerLocalOffsets,
+                                     position,
+                                     requirement.exchange->deferred,
+                                     context.constants,
+                                     context.rewriter,
+                                     loc);
+  MixedSliceGeometry geometry;
+  geometry.offsets.assign(payloadType.getRank(), context.rewriter.getIndexAttr(0));
+  geometry.sizes.push_back(context.rewriter.getIndexAttr(1));
+  geometry.strides.assign(payloadType.getRank(), context.rewriter.getIndexAttr(1));
+  geometry.offsets.front() = offset;
+  for (int64_t dimension : fragmentType.getShape())
+    geometry.sizes.push_back(context.rewriter.getIndexAttr(dimension));
+  SmallVector<int64_t> unitShape {1};
+  llvm::append_range(unitShape, fragmentType.getShape());
+  auto unitType = RankedTensorType::get(unitShape, fragmentType.getElementType());
+  Value unit = extractMixedSliceOrIdentity(context.rewriter, loc, payload, unitType, geometry);
+  return removeLeadingUnitTensorDimension(context.rewriter, loc, unit, fragmentType);
+}
+
+FailureOr<Value> realizeDeferredResult(const DeferredResultPlan& plan,
+                                       const LaneSet& activeLanes,
+                                       Value lane,
+                                       DeferredEmissionContext& context) {
+  DeferredExchangePlan& exchange = *plan.exchange;
+  if (Value assembled = context.assemblies.lookup(&exchange))
+    return assembled;
+  IRMapping mapping;
+  if (exchange.program.scheduledLane)
+    mapping.map(exchange.program.scheduledLane, lane);
+  for (auto [value, sequence] : plan.residualValues) {
+    if (mapping.contains(value))
+      continue;
+    Value selected = emitStaticIntLookup(sequence,
+                                         lane ? lane : context.constants.getIndex(0),
+                                         exchange.deferred,
+                                         context.constants,
+                                         context.rewriter,
+                                         exchange.deferred.getLoc());
+    if (!value.getType().isIndex())
+      selected = arith::IndexCastOp::create(context.rewriter, exchange.deferred.getLoc(), value.getType(), selected);
+    mapping.map(value, selected);
+  }
+  for (unsigned leaf = 0; leaf < exchange.program.leaves.size(); ++leaf) {
+    auto value = reconstructLeaf(plan, leaf, activeLanes, lane, context);
+    if (failed(value))
+      return failure();
+    mapping.map(exchange.program.leaves[leaf].replacementRoot, *value);
+  }
+  for (Operation* op : exchange.program.residualOps) {
+    Operation* copy = context.rewriter.clone(*op, mapping);
+    for (auto [oldValue, newValue] : llvm::zip(op->getResults(), copy->getResults()))
+      mapping.map(oldValue, newValue);
+  }
+  Value result = mapping.lookupOrDefault(exchange.program.yieldedValue);
+  return result && result.getType() == exchange.deferred.getOutput().getType() ? FailureOr<Value>(result)
+                                                                               : FailureOr<Value>(failure());
+}
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,35 @@
+#pragma once
+
+#include "DeferredCommunicationModel.hpp"
+
+namespace onnx_mlir::spatial {
+
+struct DeferredEmissionContext;
+
+struct DeferredResultPlan {
+  DeferredExchangePlan* exchange = nullptr;
+  llvm::SmallVector<RequirementFamily*> requirements;
+  struct SliceGeometry {
+    llvm::SmallVector<StaticIntSequence> offsets;
+    llvm::SmallVector<StaticIntSequence> sizes;
+    llvm::SmallVector<StaticIntSequence> strides;
+  };
+  llvm::SmallVector<SliceGeometry, 0> innerGeometry;
+  llvm::SmallVector<SliceGeometry, 0> assemblyGeometry;
+  llvm::DenseMap<mlir::Value, StaticIntSequence> residualValues;
+};
+
+mlir::FailureOr<DeferredResultPlan>
+buildDeferredResultPlan(DeferredExchangePlan& exchange);
+
+mlir::FailureOr<mlir::Value> realizeDeferredResult(const DeferredResultPlan& plan,
+                                                   const LaneSet& activeLanes,
+                                                   mlir::Value lane,
+                                                   DeferredEmissionContext& context);
+
+mlir::FailureOr<mlir::Value> materializeDeferredRequirement(RequirementFamily& requirement,
+                                                            const LaneSet& activeLanes,
+                                                            mlir::Value lane,
+                                                            DeferredEmissionContext& context);
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,600 @@
+#include "mlir/Dialect/SCF/IR/SCF.h"
+
+#include "llvm/ADT/MapVector.h"
+
+#include "DeferredProjectionAnalysis.hpp"
+#include "DeferredTransferPlanning.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+
+namespace onnx_mlir::spatial {
+using namespace mlir;
+namespace {
+
+static FailureOr<SmallVector<int64_t>> getI64Array(Operation* op, StringRef name) {
+  auto attr = op->getAttrOfType<DenseI64ArrayAttr>(name);
+  if (!attr)
+    return op->emitOpError() << "phase 2 requires '" << name << "' metadata";
+  return SmallVector<int64_t>(attr.asArrayRef());
+}
+
+static FailureOr<SmallVector<int64_t>> getLaneTable(Operation* op, StringRef name, size_t expected) {
+  if (auto array = op->getAttrOfType<DenseI64ArrayAttr>(name)) {
+    if (array.size() != static_cast<int64_t>(expected))
+      return op->emitOpError() << "phase 2 metadata '" << name << "' has the wrong size";
+    return SmallVector<int64_t>(array.asArrayRef());
+  }
+  auto elements = op->getAttrOfType<DenseIntElementsAttr>(name);
+  if (!elements || elements.getNumElements() != static_cast<int64_t>(expected))
+    return op->emitOpError() << "phase 2 requires a correctly-sized '" << name << "' lane table";
+  SmallVector<int64_t> values;
+  for (APInt value : elements.getValues<APInt>())
+    values.push_back(value.getSExtValue());
+  return values;
+}
+
+static Block* getScheduledBlock(SpatDeferredCommunicationOp deferred, Operation* scheduled) {
+  Block* block = deferred->getBlock();
+  while (block && block->getParentOp() != scheduled) {
+    Operation* parent = block->getParentOp();
+    block = parent ? parent->getBlock() : nullptr;
+  }
+  return block;
+}
+
+static FailureOr<unsigned> getStepIndex(ScheduledInfo& info, Block* block) {
+  auto it = llvm::find(info.blocks, block);
+  return it == info.blocks.end() ? FailureOr<unsigned>(failure())
+                                 : FailureOr<unsigned>(std::distance(info.blocks.begin(), it));
+}
+
+static FailureOr<Value>
+getScalarStepResult(SpatScheduledCompute scheduled, Block& block, unsigned resultIndex, unsigned resultCount) {
+  auto yield = dyn_cast<SpatBlockYieldOp>(block.getTerminator());
+  if (!yield || resultIndex >= resultCount || yield.getOutputs().size() < resultCount)
+    return scheduled.emitOpError("phase 2 cannot recover a scalar scheduled step result"), failure();
+  return yield.getOutputs()[yield.getOutputs().size() - resultCount + resultIndex];
+}
+
+struct BatchPublication {
+  Value payload;
+  tensor::ParallelInsertSliceOp insertion;
+};
+
+struct ProducedValueGeometry {
+  int64_t laneStart = 0;
+  int64_t laneCount = 1;
+  int64_t publishedSlotStart = 0;
+  int64_t publishedSlotCount = 1;
+  Value payload;
+  Value published;
+};
+
+static FailureOr<BatchPublication>
+getBatchStepResult(SpatScheduledComputeBatch scheduled, Block& block, unsigned globalResultIndex) {
+  auto parallel = dyn_cast<SpatInParallelOp>(block.getTerminator());
+  if (!parallel)
+    return scheduled.emitOpError("phase 2 cannot recover a batched scheduled step result"), failure();
+  unsigned resultBase = 1 + scheduled.getWeights().size() + scheduled.getInputs().size();
+  for (Operation& op : parallel.getRegion().front()) {
+    auto insert = dyn_cast<tensor::ParallelInsertSliceOp>(op);
+    auto destination = insert ? dyn_cast<BlockArgument>(insert.getDest()) : BlockArgument();
+    if (destination && destination.getOwner() == &block && destination.getArgNumber() == resultBase + globalResultIndex)
+      return BatchPublication {insert.getSource(), insert};
+  }
+  return scheduled.emitOpError("phase 2 cannot find the batched scheduled result insertion"), failure();
+}
+
+static FailureOr<ProducedValueGeometry> verifyScheduledPublicationGeometry(
+  ScheduledInfo& info, unsigned step, unsigned globalResult, unsigned lane,
+  int64_t laneStart, int64_t laneCount, Value payload,
+  tensor::ParallelInsertSliceOp insertion = {}) {
+  ProducedValueGeometry result;
+  result.laneStart = laneStart;
+  result.laneCount = info.isBatch() ? laneCount : std::max<int64_t>(laneCount, 1);
+  result.payload = payload;
+  result.published = info.op->getResult(globalResult);
+  auto payloadType = dyn_cast<RankedTensorType>(payload.getType());
+  auto publishedType = dyn_cast<RankedTensorType>(result.published.getType());
+  if (laneStart < 0 || result.laneCount <= 0)
+    return info.op->emitOpError("phase 2 scheduled publication has an invalid source-lane range"), failure();
+
+  if (!info.isBatch()) {
+    result.publishedSlotCount = result.laneCount;
+    if (payload.getType() != result.published.getType()
+        || (result.laneCount > 1
+            && (!publishedType || !publishedType.hasStaticShape()
+                || publishedType.getRank() == 0
+                || publishedType.getDimSize(0) != result.laneCount)))
+      return info.op->emitOpError("phase 2 scalar scheduled publication is incompatible with its result"), failure();
+    return result;
+  }
+
+  if (!insertion || !payloadType || !payloadType.hasStaticShape()
+      || !publishedType || !publishedType.hasStaticShape()
+      || insertion.getDest() != info.blocks[step]->getArgument(
+           1 + info.op->getNumOperands() + globalResult))
+    return info.op->emitOpError("phase 2 batched publication is malformed"), failure();
+  unsigned rank = publishedType.getRank();
+  if (rank == 0 || payloadType.getRank() != rank
+      || insertion.getMixedOffsets().size() != rank
+      || insertion.getMixedSizes().size() != rank
+      || insertion.getMixedStrides().size() != rank)
+    return info.op->emitOpError("phase 2 batched publication geometry rank is invalid"), failure();
+
+  StaticIndexEnvironment environment;
+  environment.bindings[info.blocks[step]->getArgument(0)] = lane;
+  SmallVector<int64_t> offsets, sizes, strides;
+  for (OpFoldResult value : insertion.getMixedOffsets()) {
+    auto evaluated = evaluateDeferredIndex(value, environment);
+    if (failed(evaluated))
+      return info.op->emitOpError("phase 2 batched publication offset is not static"), failure();
+    offsets.push_back(*evaluated);
+  }
+  for (OpFoldResult value : insertion.getMixedSizes()) {
+    auto evaluated = evaluateDeferredIndex(value, environment);
+    if (failed(evaluated))
+      return info.op->emitOpError("phase 2 batched publication size is not static"), failure();
+    sizes.push_back(*evaluated);
+  }
+  for (OpFoldResult value : insertion.getMixedStrides()) {
+    auto evaluated = evaluateDeferredIndex(value, environment);
+    if (failed(evaluated))
+      return info.op->emitOpError("phase 2 batched publication stride is not static"), failure();
+    strides.push_back(*evaluated);
+  }
+  if (offsets.front() < 0 || sizes.front() <= 0 || sizes.front() != result.laneCount
+      || strides.front() != 1
+      || offsets.front() + sizes.front() > publishedType.getDimSize(0))
+    return info.op->emitOpError("phase 2 batched publication leading geometry is invalid"), failure();
+  for (unsigned dimension = 1; dimension < rank; ++dimension)
+    if (offsets[dimension] != 0 || strides[dimension] != 1)
+      return info.op->emitOpError("phase 2 batched publication inner geometry is invalid"), failure();
+  for (unsigned dimension = 0; dimension < rank; ++dimension)
+    if (sizes[dimension] != payloadType.getDimSize(dimension))
+      return info.op->emitOpError("phase 2 batched publication payload shape is incompatible"), failure();
+  result.publishedSlotStart = offsets.front();
+  result.publishedSlotCount = sizes.front();
+  return result;
+}
+
+static LogicalResult collectScheduledOperations(func::FuncOp funcOp, DeferredTransferPlan& plan) {
+  unsigned nextStream = 0;
+  for (Operation& op : funcOp.getOps()) {
+    if (!isa<SpatScheduledCompute, SpatScheduledComputeBatch>(op))
+      continue;
+    ScheduledInfo info;
+    info.op = &op;
+    Region& body = isa<SpatScheduledCompute>(op) ? cast<SpatScheduledCompute>(op).getBody()
+                                                 : cast<SpatScheduledComputeBatch>(op).getBody();
+    for (Block& block : body) {
+      info.blocks.push_back(&block);
+      info.stepAnchors.push_back(&block.front());
+    }
+    auto sourceIds = getI64Array(&op, "scheduled.step_source_ids");
+    auto offsets = getI64Array(&op, "scheduled.step_result_offsets");
+    auto counts = getI64Array(&op, "scheduled.step_result_counts");
+    if (failed(sourceIds) || failed(offsets) || failed(counts))
+      return failure();
+    info.stepSourceIds = std::move(*sourceIds);
+    info.resultOffsets = std::move(*offsets);
+    info.resultCounts = std::move(*counts);
+    if (info.blocks.size() != info.stepSourceIds.size() || info.blocks.size() != info.resultOffsets.size()
+        || info.blocks.size() != info.resultCounts.size())
+      return op.emitOpError("phase 2 scheduled metadata does not match its block count");
+    if (auto scalar = dyn_cast<SpatScheduledCompute>(op)) {
+      auto core = scalar->getAttrOfType<IntegerAttr>(kCoreIdAttrName);
+      if (!core)
+        return scalar.emitOpError("phase 2 requires scalar coreId metadata");
+      info.cores.push_back(core.getInt());
+    }
+    else {
+      auto cores = op.getAttrOfType<DenseI32ArrayAttr>(kCoreIdsAttrName);
+      if (!cores)
+        return op.emitOpError("phase 2 requires batch coreIds metadata");
+      for (int32_t core : cores.asArrayRef())
+        info.cores.push_back(core);
+    }
+    for (size_t lane = 0; lane < info.cores.size(); ++lane)
+      info.streamIds.push_back(nextStream++);
+    plan.scheduled.push_back(std::move(info));
+  }
+  plan.stepCounts.resize(nextStream);
+  for (ScheduledInfo& info : plan.scheduled)
+    for (unsigned stream : info.streamIds)
+      plan.stepCounts[stream] = info.blocks.size();
+  return success();
+}
+
+static LogicalResult collectProducedValues(DeferredTransferPlan& plan) {
+  for (ScheduledInfo& info : plan.scheduled) {
+    SmallVector<int64_t> laneStarts, laneCounts;
+    size_t tableSize = info.blocks.size() * info.cores.size();
+    auto starts = info.isBatch() ? getLaneTable(info.op, "scheduled.source_lane_starts", tableSize)
+                                 : getI64Array(info.op, "scheduled.source_lane_starts");
+    auto counts = info.isBatch() ? getLaneTable(info.op, "scheduled.source_lane_counts", tableSize)
+                                 : getI64Array(info.op, "scheduled.source_lane_counts");
+    if (failed(starts) || failed(counts))
+      return failure();
+    laneStarts = std::move(*starts);
+    laneCounts = std::move(*counts);
+    for (unsigned step = 0; step < info.blocks.size(); ++step) {
+      if (info.resultOffsets[step] < 0 || info.resultCounts[step] < 0)
+        return info.op->emitOpError("phase 2 scheduled result metadata must be non-negative");
+      for (unsigned result = 0; result < static_cast<unsigned>(info.resultCounts[step]); ++result) {
+        unsigned globalResult = info.resultOffsets[step] + result;
+        if (!info.isBatch()) {
+          auto payload = getScalarStepResult(
+            cast<SpatScheduledCompute>(info.op), *info.blocks[step], result, info.resultCounts[step]);
+          if (failed(payload))
+            return failure();
+          auto geometry = verifyScheduledPublicationGeometry(
+            info, step, globalResult, 0, laneStarts[step], laneCounts[step], *payload);
+          if (failed(geometry))
+            return failure();
+          auto produced = std::make_unique<ProducedValue>(ProducedValue {
+            &info, step, result, info.stepSourceIds[step], info.cores.front(),
+            geometry->laneStart, geometry->laneCount, 0,
+            geometry->publishedSlotStart, geometry->publishedSlotCount,
+            geometry->payload, geometry->published});
+          info.produced.push_back(produced.get());
+          plan.producedByGraph[produced->graphId].push_back(produced.get());
+          plan.producedStorage.push_back(std::move(produced));
+          continue;
+        }
+        auto publication =
+          getBatchStepResult(cast<SpatScheduledComputeBatch>(info.op), *info.blocks[step], globalResult);
+        if (failed(publication))
+          return failure();
+        for (unsigned lane = 0; lane < info.cores.size(); ++lane) {
+          size_t laneIndex = step * info.cores.size() + lane;
+          auto geometry = verifyScheduledPublicationGeometry(
+            info, step, globalResult, lane, laneStarts[laneIndex],
+            laneCounts[laneIndex], publication->payload, publication->insertion);
+          if (failed(geometry))
+            return failure();
+          auto produced = std::make_unique<ProducedValue>(ProducedValue {
+            &info, step, result, info.stepSourceIds[step], info.cores[lane],
+            geometry->laneStart, geometry->laneCount, lane,
+            geometry->publishedSlotStart, geometry->publishedSlotCount,
+            geometry->payload, geometry->published});
+          info.produced.push_back(produced.get());
+          plan.producedByGraph[produced->graphId].push_back(produced.get());
+          plan.producedStorage.push_back(std::move(produced));
+        }
+      }
+    }
+  }
+  return success();
+}
+
+static FailureOr<ProducedValue*> findProducer(DeferredTransferPlan& plan,
+                                              Operation* diagnosticOwner,
+                                              int64_t graphId,
+                                              unsigned resultIndex,
+                                              std::optional<int64_t> graphLane) {
+  ProducedValue* match = nullptr;
+  for (ProducedValue* produced : plan.producedByGraph.lookup(graphId)) {
+    if (produced->resultIndex != resultIndex
+        || (graphLane && (*graphLane < produced->laneStart || *graphLane >= produced->laneStart + produced->laneCount)))
+      continue;
+    if (match)
+      return diagnosticOwner->emitOpError("phase 2 cannot uniquely resolve graph publication ownership"), failure();
+    match = produced;
+  }
+  if (!match)
+    return diagnosticOwner->emitOpError("phase 2 cannot map a graph lane to a scheduled producer"), failure();
+  if (graphLane) {
+    int64_t relative = *graphLane - match->laneStart;
+    if (relative < 0 || relative >= match->laneCount
+        || relative >= match->publishedSlotCount)
+      return diagnosticOwner->emitOpError(
+        "phase 2 graph lane is outside its scheduled publication window"), failure();
+  }
+  return match;
+}
+
+struct RequirementPoint {
+  ProducedValue* producer = nullptr;
+  Type fragmentType;
+  std::optional<int64_t> graphLane;
+  std::optional<int64_t> localOffset;
+
+  bool sameFamily(const RequirementPoint& other) const {
+    return producer == other.producer && fragmentType == other.fragmentType;
+  }
+};
+
+static FailureOr<std::optional<RequirementPoint>>
+resolveRequirementPoint(DeferredTransferPlan& plan,
+                        DeferredExchangePlan& exchange,
+                        const DeferredProjectionLeafTemplate& leaf,
+                        DeferredLaneValueEvaluator& evaluator,
+                        unsigned lane,
+                        unsigned position,
+                        GraphBatchPublicationCache& publicationCache) {
+  auto sourceIndices = evaluator.resolveSourceOperandIndices(leaf.sourceRoot);
+  if (failed(sourceIndices))
+    return failure();
+  unsigned sourceIndex = sourceIndices->valueAt(lane);
+  auto source = dyn_cast<OpResult>(exchange.deferred.getSources()[sourceIndex]);
+  if (!source)
+    return exchange.deferred.emitOpError("phase 2 requires graph-result deferred sources"), failure();
+  auto graphId = source.getOwner()->getAttrOfType<IntegerAttr>("scheduled.graph_id");
+  if (!graphId)
+    return exchange.deferred.emitOpError("phase 2 cannot identify graph producer"), failure();
+
+  RequirementPoint point;
+  if (auto batch = dyn_cast<SpatGraphComputeBatch>(source.getOwner())) {
+    auto publication = getGraphBatchPublicationMap(batch, source.getResultNumber(), publicationCache);
+    if (failed(publication))
+      return failure();
+    int64_t physicalSlot = position;
+    if (leaf.form == DeferredLeafForm::GraphBatchProjection) {
+      auto offset = evaluator.evaluate(leaf.leadingGeometry.offsets.front());
+      auto size = evaluator.evaluate(leaf.leadingGeometry.sizes.front());
+      auto stride = evaluator.evaluate(leaf.leadingGeometry.strides.front());
+      if (failed(offset) || failed(size) || failed(stride))
+        return failure();
+      if (position >= static_cast<unsigned>(size->valueAt(lane)))
+        return std::optional<RequirementPoint>();
+      physicalSlot = offset->valueAt(lane) + static_cast<int64_t>(position) * stride->valueAt(lane);
+    }
+    else if (position >= (*publication)->physicalSlotToGraphLane.size()) {
+      return std::optional<RequirementPoint>();
+    }
+    if (physicalSlot < 0 || physicalSlot >= static_cast<int64_t>((*publication)->physicalSlotToGraphLane.size()))
+      return exchange.deferred.emitOpError("projection physical slot is outside publication map"), failure();
+    point.graphLane = (*publication)->physicalSlotToGraphLane[physicalSlot];
+    point.fragmentType = (*publication)->publicationFragmentType;
+    auto producer = findProducer(plan, exchange.deferred, graphId.getInt(), source.getResultNumber(), point.graphLane);
+    if (failed(producer))
+      return failure();
+    point.producer = *producer;
+    point.localOffset = *point.graphLane - point.producer->laneStart;
+  }
+  else {
+    if (position != 0 || !isa<SpatGraphCompute>(source.getOwner()))
+      return std::optional<RequirementPoint>();
+    point.fragmentType = source.getType();
+    auto producer = findProducer(plan, exchange.deferred, graphId.getInt(), source.getResultNumber(), std::nullopt);
+    if (failed(producer))
+      return failure();
+    point.producer = *producer;
+  }
+  return std::optional<RequirementPoint>(point);
+}
+
+static void appendRequirementFamily(DeferredExchangePlan& exchange,
+                                    RequirementCoordinate coordinate,
+                                    unsigned begin,
+                                    ArrayRef<RequirementPoint> points) {
+  RequirementFamily family;
+  family.exchange = &exchange;
+  family.coordinate = coordinate;
+  family.targetLanes = LaneSet::range(begin, begin + points.size());
+  family.producer = points.front().producer;
+  family.publicationFragmentType = points.front().fragmentType;
+  auto sequence = [&](auto member) -> std::optional<StaticIntSequence> {
+    if (!(points.front().*member))
+      return std::nullopt;
+    SmallVector<int64_t> values;
+    for (const RequirementPoint& point : points)
+      values.push_back(*(point.*member));
+    return StaticIntSequence::fromValues(values);
+  };
+  family.graphLanes = sequence(&RequirementPoint::graphLane);
+  family.producerLocalOffsets = sequence(&RequirementPoint::localOffset);
+  exchange.requirements.push_back(std::move(family));
+}
+
+static LogicalResult buildRequirementFamilies(DeferredTransferPlan& plan,
+                                              DeferredExchangePlan& exchange,
+                                              GraphBatchPublicationCache& publicationCache) {
+  DeferredLaneValueEvaluator evaluator(exchange.program, exchange.targetLaneCount);
+  for (auto leafItem : llvm::enumerate(exchange.program.leaves)) {
+    unsigned leafIndex = leafItem.index();
+    const DeferredProjectionLeafTemplate& leaf = leafItem.value();
+    unsigned positionCount = 1;
+    if (leaf.form == DeferredLeafForm::GraphBatchProjection) {
+      auto sizes = evaluator.evaluate(leaf.leadingGeometry.sizes.front());
+      if (failed(sizes))
+        return failure();
+      for (unsigned lane = 0; lane < exchange.targetLaneCount; ++lane)
+        positionCount = std::max<unsigned>(positionCount, sizes->valueAt(lane));
+    }
+    else {
+      auto sources = evaluator.resolveSourceOperandIndices(leaf.sourceRoot);
+      if (failed(sources))
+        return failure();
+      for (unsigned lane = 0; lane < exchange.targetLaneCount; ++lane) {
+        auto source = dyn_cast<OpResult>(exchange.deferred.getSources()[sources->valueAt(lane)]);
+        auto type = source ? dyn_cast<RankedTensorType>(source.getType()) : RankedTensorType();
+        if (source && isa<SpatGraphComputeBatch>(source.getOwner()) && type)
+          positionCount = std::max<unsigned>(positionCount, type.getDimSize(0));
+      }
+    }
+    for (unsigned position = 0; position < positionCount; ++position) {
+      unsigned runBegin = 0;
+      SmallVector<RequirementPoint> run;
+      auto flush = [&] {
+        if (!run.empty())
+          appendRequirementFamily(exchange, {static_cast<unsigned>(leafIndex), position}, runBegin, run);
+        run.clear();
+      };
+      for (unsigned lane = 0; lane < exchange.targetLaneCount; ++lane) {
+        auto point = resolveRequirementPoint(plan, exchange, leaf, evaluator, lane, position, publicationCache);
+        if (failed(point))
+          return failure();
+        if (!*point) {
+          flush();
+          continue;
+        }
+        if (!run.empty() && !run.front().sameFamily(**point))
+          flush();
+        if (run.empty())
+          runBegin = lane;
+        run.push_back(**point);
+      }
+      flush();
+    }
+  }
+  return success();
+}
+
+static void buildAvailabilityFamilies(DeferredExchangePlan& exchange, uint64_t& nextChannel) {
+  for (RequirementFamily& requirement : exchange.requirements) {
+    for (LaneInterval interval : requirement.targetLanes.intervals()) {
+      unsigned runBegin = interval.begin;
+      bool runLocal = false;
+      bool haveRun = false;
+      auto flush = [&](unsigned end) {
+        if (!haveRun || runBegin == end)
+          return;
+        LaneSet lanes = LaneSet::range(runBegin, end);
+        if (runLocal) {
+          exchange.local.push_back({&requirement, lanes});
+        }
+        else {
+          size_t count = end - runBegin;
+          unsigned sourceStream = requirement.producer->scheduled->streamIds[requirement.producer->scheduledLane];
+          SmallVector<int64_t> targetStreams, targetCores;
+          for (unsigned lane = runBegin; lane < end; ++lane) {
+            targetStreams.push_back(exchange.target->streamIds[lane]);
+            targetCores.push_back(exchange.target->cores[lane]);
+          }
+          ExternalTransferFamily family;
+          family.requirement = &requirement;
+          family.targetLanes = lanes;
+          family.sourceScheduled = requirement.producer->scheduled;
+          family.targetScheduled = exchange.target;
+          family.sourceStreams = StaticIntSequence::uniform(sourceStream, count);
+          family.targetStreams = StaticIntSequence::fromValues(targetStreams);
+          family.sourceCores = StaticIntSequence::uniform(requirement.producer->core, count);
+          family.targetCores = StaticIntSequence::fromValues(targetCores);
+          family.channelIds = StaticIntSequence::affine(nextChannel, 1, count);
+          nextChannel += count;
+          exchange.externalTransferCount += count;
+          exchange.external.push_back(std::move(family));
+        }
+      };
+      for (unsigned lane = interval.begin; lane < interval.end; ++lane) {
+        unsigned sourceStream = requirement.producer->scheduled->streamIds[requirement.producer->scheduledLane];
+        bool local =
+          sourceStream == exchange.target->streamIds[lane] && requirement.producer->step < exchange.consumerStep;
+        if (haveRun && local != runLocal) {
+          flush(lane);
+          runBegin = lane;
+        }
+        runLocal = local;
+        haveRun = true;
+      }
+      flush(interval.end);
+    }
+  }
+}
+
+static LogicalResult buildExchanges(func::FuncOp funcOp, DeferredTransferPlan& plan) {
+  DenseMap<Operation*, ScheduledInfo*> scheduledByOp;
+  for (ScheduledInfo& scheduled : plan.scheduled)
+    scheduledByOp[scheduled.op] = &scheduled;
+  SmallVector<SpatDeferredCommunicationOp> deferredOps;
+  funcOp.walk([&](SpatDeferredCommunicationOp op) { deferredOps.push_back(op); });
+  GraphBatchPublicationCache publicationCache;
+  uint64_t nextChannel = 0;
+  for (SpatDeferredCommunicationOp deferred : deferredOps) {
+    Operation* targetOp = deferred->getParentOfType<SpatScheduledCompute>();
+    if (!targetOp)
+      targetOp = deferred->getParentOfType<SpatScheduledComputeBatch>();
+    ScheduledInfo* target = scheduledByOp.lookup(targetOp);
+    auto step = target ? getStepIndex(*target, getScheduledBlock(deferred, targetOp)) : FailureOr<unsigned>(failure());
+    auto program = analyzeDeferredProgramTemplate(deferred);
+    if (!target || failed(step) || failed(program))
+      return deferred.emitOpError("phase 2 cannot normalize deferred communication");
+    auto exchange = std::make_unique<DeferredExchangePlan>();
+    exchange->exchangeId = plan.exchanges.size();
+    exchange->deferred = deferred;
+    exchange->target = target;
+    exchange->consumerStep = *step;
+    exchange->targetLaneCount = target->cores.size();
+    exchange->program = std::move(*program);
+    if (failed(buildRequirementFamilies(plan, *exchange, publicationCache)))
+      return failure();
+    buildAvailabilityFamilies(*exchange, nextChannel);
+    plan.exchanges.push_back(std::move(exchange));
+  }
+  return success();
+}
+
+static LogicalResult
+retargetBlueprint(DeferredTransferPlan& plan, SpatBlueprintOp blueprint, GraphBatchPublicationCache& publicationCache) {
+  if (blueprint.getMode() != "fragment_assembly")
+    return success();
+  auto operandIndices = blueprint.getFragmentOperandIndices();
+  auto sourceSlots = blueprint.getFragmentSourceSlots();
+  if (!operandIndices || !sourceSlots)
+    return blueprint.emitOpError("phase 2 requires explicit Blueprint fragment ownership");
+  SmallVector<Value> oldOperands {blueprint.getInput()};
+  llvm::append_range(oldOperands, blueprint.getFragments());
+  SmallVector<Value> publications;
+  SmallVector<int64_t> newOperands, newSlots;
+  for (auto [fragmentIndex, operandIndex] : llvm::enumerate(*operandIndices)) {
+    Value source = oldOperands[operandIndex];
+    auto result = dyn_cast<OpResult>(source);
+    auto batch = result ? dyn_cast<SpatGraphComputeBatch>(result.getOwner()) : SpatGraphComputeBatch();
+    int64_t slot = (*sourceSlots)[fragmentIndex];
+    if (batch) {
+      auto graphId = batch->getAttrOfType<IntegerAttr>("scheduled.graph_id");
+      auto map = getGraphBatchPublicationMap(batch, result.getResultNumber(), publicationCache);
+      if (!graphId || failed(map) || slot < 0 || slot >= static_cast<int64_t>((*map)->physicalSlotToGraphLane.size()))
+        return blueprint.emitOpError("phase 2 cannot resolve Blueprint fragment ownership");
+      int64_t graphLane = (*map)->physicalSlotToGraphLane[slot];
+      auto producer = findProducer(plan, blueprint, graphId.getInt(), result.getResultNumber(), graphLane);
+      if (failed(producer))
+        return failure();
+      source = (*producer)->published;
+      slot = (*producer)->publishedSlotStart + graphLane - (*producer)->laneStart;
+      if (slot < (*producer)->publishedSlotStart
+          || slot >= (*producer)->publishedSlotStart
+                     + (*producer)->publishedSlotCount)
+        return blueprint.emitOpError(
+          "phase 2 Blueprint slot is outside its scheduled publication window"), failure();
+    }
+    if (Operation* producer = source.getDefiningOp();
+        producer && blueprint->getBlock() == producer->getBlock() && blueprint->isBeforeInBlock(producer))
+      blueprint->moveAfter(producer);
+    auto it = llvm::find(publications, source);
+    if (it == publications.end()) {
+      publications.push_back(source);
+      it = std::prev(publications.end());
+    }
+    newOperands.push_back(std::distance(publications.begin(), it));
+    newSlots.push_back(slot);
+  }
+  blueprint->setOperands(publications);
+  OpBuilder builder(blueprint);
+  blueprint->setAttr("fragmentOperandIndices", builder.getDenseI64ArrayAttr(newOperands));
+  blueprint->setAttr("fragmentSourceSlots", builder.getDenseI64ArrayAttr(newSlots));
+  return success();
+}
+
+} // namespace
+
+FailureOr<DeferredTransferPlan> buildDeferredTransferPlan(func::FuncOp funcOp) {
+  DeferredTransferPlan plan;
+  if (failed(collectScheduledOperations(funcOp, plan)) || failed(collectProducedValues(plan))
+      || failed(buildExchanges(funcOp, plan)))
+    return failure();
+  return std::move(plan);
+}
+
+LogicalResult retargetDeferredPublications(func::FuncOp funcOp, DeferredTransferPlan& plan) {
+  GraphBatchPublicationCache publicationCache;
+  SmallVector<SpatBlueprintOp> blueprints;
+  funcOp.walk([&](SpatBlueprintOp blueprint) { blueprints.push_back(blueprint); });
+  for (SpatBlueprintOp blueprint : blueprints)
+    if (failed(retargetBlueprint(plan, blueprint, publicationCache)))
+      return failure();
+  return success();
+}
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,21 @@
+#pragma once
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+
+#include "DeferredCommunicationModel.hpp"
+
+namespace onnx_mlir::spatial {
+
+struct DeferredTransferPlan {
+  llvm::SmallVector<ScheduledInfo, 0> scheduled;
+  llvm::SmallVector<std::unique_ptr<ProducedValue>> producedStorage;
+  llvm::DenseMap<int64_t, llvm::SmallVector<ProducedValue*>> producedByGraph;
+  llvm::SmallVector<std::unique_ptr<DeferredExchangePlan>> exchanges;
+  llvm::SmallVector<unsigned> stepCounts;
+};
+
+mlir::FailureOr<DeferredTransferPlan> buildDeferredTransferPlan(mlir::func::FuncOp funcOp);
+
+mlir::LogicalResult retargetDeferredPublications(mlir::func::FuncOp funcOp, DeferredTransferPlan& plan);
+
+} // namespace onnx_mlir::spatial

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

@@ -1,134 +0,0 @@
-#include "HostOutputFinalization.hpp"
-
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/IR/BuiltinTypes.h"
-#include "mlir/IR/PatternMatch.h"
-
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/DenseSet.h"
-#include "llvm/ADT/STLExtras.h"
-
-#include "MaterializedClassState.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-
-using namespace mlir;
-
-namespace onnx_mlir::spatial {
-
-LogicalResult finalizeProjectedHostOutputFragments(MaterializerState& state) {
-  if (state.pendingProjectedHostOutputFragments.empty())
-    return success();
-
-  DenseMap<Value, SmallVector<PendingProjectedHostOutputFragment*, 16>> byOutput;
-  for (PendingProjectedHostOutputFragment& fragment : state.pendingProjectedHostOutputFragments)
-    byOutput[fragment.originalOutput].push_back(&fragment);
-
-  SmallVector<Value, 8> outputs;
-  outputs.reserve(byOutput.size());
-
-  auto returnOp = dyn_cast<func::ReturnOp>(state.func.getBody().front().getTerminator());
-  if (!returnOp)
-    return state.func.emitError("expected func.return terminator while finalizing projected host output fragments");
-
-  DenseSet<Value> seenOutputs;
-  for (Value returned : returnOp.getOperands()) {
-    if (!byOutput.contains(returned) || !seenOutputs.insert(returned).second)
-      continue;
-    outputs.push_back(returned);
-  }
-  if (outputs.size() != byOutput.size())
-    return state.func.emitError("projected host output fragments must be keyed by returned logical host outputs");
-
-  for (Value originalOutput : outputs) {
-    if (isa_and_present<SpatScheduledCompute, SpatScheduledComputeBatch>(originalOutput.getDefiningOp())) {
-      return state.func.emitError(
-        "projected host output assembly must be keyed by the original logical host output, not by a materialized scheduled result");
-    }
-
-    auto resultType = dyn_cast<RankedTensorType>(originalOutput.getType());
-    if (!resultType || !resultType.hasStaticShape())
-      return state.func.emitError("projected host output must have static ranked tensor type");
-
-    SmallVector<PendingProjectedHostOutputFragment*, 16>& fragments = byOutput[originalOutput];
-    llvm::sort(fragments, [](const PendingProjectedHostOutputFragment* lhs,
-                             const PendingProjectedHostOutputFragment* rhs) {
-      if (lhs->sourceClass != rhs->sourceClass)
-        return lhs->sourceClass < rhs->sourceClass;
-      if (lhs->publicationResultIndex != rhs->publicationResultIndex)
-        return lhs->publicationResultIndex < rhs->publicationResultIndex;
-      if (lhs->sourceFragmentOrdinal != rhs->sourceFragmentOrdinal)
-        return lhs->sourceFragmentOrdinal < rhs->sourceFragmentOrdinal;
-      return std::lexicographical_compare(lhs->offsets.begin(),
-                                          lhs->offsets.end(),
-                                          rhs->offsets.begin(),
-                                          rhs->offsets.end());
-    });
-
-    state.rewriter.setInsertionPoint(returnOp);
-    Location loc = fragments.front()->loc;
-    SmallVector<Value, 16> blueprintOperands;
-    SmallVector<int64_t, 16> fragmentOperandIndices;
-    SmallVector<int64_t, 16> fragmentSourceOffsets;
-    SmallVector<int64_t, 64> flatOffsets;
-    SmallVector<int64_t, 64> flatSizes;
-    SmallVector<int64_t, 64> flatStrides;
-    DenseMap<Value, int64_t> operandIndicesByValue;
-
-    for (PendingProjectedHostOutputFragment* fragmentRecord : fragments) {
-      if (fragmentRecord->sourceClass >= state.classes.size())
-        return state.func.emitError("projected host output fragment references an invalid source class");
-
-      MaterializedClass& sourceClass = state.classes[fragmentRecord->sourceClass];
-      if (fragmentRecord->publicationResultIndex >= sourceClass.op->getNumResults()) {
-        return sourceClass.op->emitError("projected host output fragment references an invalid publication result")
-               << " sourceClass=" << sourceClass.id
-               << " resultIndex=" << fragmentRecord->publicationResultIndex
-               << " resultCount=" << sourceClass.op->getNumResults();
-      }
-
-      Value operand = sourceClass.op->getResult(fragmentRecord->publicationResultIndex);
-      auto [operandIt, inserted] =
-        operandIndicesByValue.try_emplace(operand, static_cast<int64_t>(blueprintOperands.size()));
-      if (inserted)
-        blueprintOperands.push_back(operand);
-      fragmentOperandIndices.push_back(operandIt->second);
-      fragmentSourceOffsets.push_back(fragmentRecord->sourceElementOffset);
-      llvm::append_range(flatOffsets, fragmentRecord->offsets);
-      llvm::append_range(flatSizes, fragmentRecord->sizes);
-      llvm::append_range(flatStrides, fragmentRecord->strides);
-
-      auto operandType = dyn_cast<RankedTensorType>(operand.getType());
-      if (!operandType || !operandType.hasStaticShape())
-        return state.func.emitError("projected host output assembly requires static ranked tensor operands");
-    }
-
-    if (blueprintOperands.empty())
-      return state.func.emitError("missing projected host output fragments");
-
-    Value input = blueprintOperands.front();
-    ValueRange extraFragments = ValueRange(blueprintOperands).drop_front();
-    auto blueprint = SpatBlueprintOp::create(
-      state.rewriter,
-      loc,
-      resultType,
-      input,
-      extraFragments,
-      state.rewriter.getStringAttr("nchw"),
-      state.rewriter.getStringAttr("fragmented"),
-      state.rewriter.getDenseI64ArrayAttr(flatOffsets),
-      state.rewriter.getDenseI64ArrayAttr(flatSizes),
-      state.rewriter.getStringAttr("identity"),
-      state.rewriter.getStringAttr("fragment_assembly"),
-      state.rewriter.getDenseI64ArrayAttr(fragmentOperandIndices),
-      state.rewriter.getDenseI64ArrayAttr(fragmentSourceOffsets),
-      state.rewriter.getDenseI64ArrayAttr(flatStrides),
-      state.rewriter.getStringAttr("disjoint"),
-      state.rewriter.getStringAttr("complete"));
-
-    state.hostReplacements[originalOutput] = blueprint.getOutput();
-  }
-
-  return success();
-}
-
-} // namespace onnx_mlir::spatial

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

@@ -1,11 +0,0 @@
-#pragma once
-
-#include "mlir/Support/LogicalResult.h"
-
-namespace onnx_mlir::spatial {
-
-struct MaterializerState;
-
-mlir::LogicalResult finalizeProjectedHostOutputFragments(MaterializerState& state);
-
-} // namespace onnx_mlir::spatial

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/MaterializedClassState.hpp

@@ -1,252 +0,0 @@
-#pragma once
-
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/IR/PatternMatch.h"
-#include "mlir/Transforms/FoldUtils.h"
-
-#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/DenseSet.h"
-#include "llvm/ADT/SmallVector.h"
-
-#include <optional>
-
-#include "MaterializeMergeSchedule.hpp"
-#include "MergeMessages.hpp"
-#include "MergeScheduleKeys.hpp"
-#include "ProjectedFragments.hpp"
-#include "Scheduling/ComputeInstanceUtils.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-
-namespace onnx_mlir::spatial {
-
-struct MaterializedClass {
-  ClassId id = 0;
-  llvm::SmallVector<CpuId, 8> cpus;
-  mlir::Operation* op = nullptr;
-  mlir::Block* body = nullptr;
-  bool isBatch = false;
-
-  llvm::DenseMap<CpuId, unsigned> cpuToLane;
-  llvm::SmallVector<mlir::Value, 8> weights;
-  llvm::SmallVector<mlir::Value, 8> inputs;
-  llvm::SmallVector<mlir::Value, 4> hostOutputs;
-  llvm::DenseMap<mlir::Value, unsigned> publicationOutputToResultIndex;
-  llvm::DenseMap<mlir::Value, mlir::BlockArgument> weightArgs;
-  llvm::DenseMap<mlir::Value, mlir::BlockArgument> inputArgs;
-  llvm::DenseMap<mlir::Value, unsigned> hostOutputToResultIndex;
-};
-
-struct PackedScalarRunSlot {
-  llvm::SmallVector<ProducerKey, 8> keys;
-};
-
-enum class PackedScalarRunKind {
-  Materialized,
-  DeferredReceive,
-  DeferredLocalCompute
-};
-
-struct PackedScalarRunValue {
-  ClassId targetClass = 0;
-  mlir::Operation* sourceOp = nullptr;
-  size_t resultIndex = 0;
-  PackedScalarRunKind kind = PackedScalarRunKind::Materialized;
-
-  mlir::Value packed;
-
-  mlir::RankedTensorType fragmentType;
-  llvm::SmallVector<PackedScalarRunSlot, 8> slots;
-  MessageVector messages;
-};
-
-struct IndexedBatchRunValue {
-  ClassId targetClass = 0;
-  mlir::Operation* sourceOp = nullptr;
-  size_t resultIndex = 0;
-  mlir::Value packed;
-  mlir::RankedTensorType fragmentType;
-  llvm::SmallVector<PackedScalarRunSlot, 8> slots;
-  MessageVector messages;
-};
-
-struct LogicalSlotRange {
-  SlotId start = 0;
-  SlotId count = 0;
-};
-
-struct MaterializationRunSlot {
-  llvm::SmallVector<ComputeInstance, 8> peers;
-};
-
-using MaterializationRun = llvm::SmallVector<MaterializationRunSlot, 8>;
-
-struct OutputDestinationGroup {
-  llvm::SmallVector<size_t, 4> resultIndices;
-  llvm::SmallVector<ClassId, 4> destinationClasses;
-};
-
-struct BatchRunSendPlan {
-  size_t resultIndex = 0;
-  ClassId destinationClass = 0;
-  MessageVector messages;
-};
-
-enum class TensorDemandActionKind {
-  DestinationFanout,
-  SameClassIndexedFragment,
-  TerminalBlueprintPublication,
-  WholeTensorBarrier
-};
-
-enum class WholeTensorBarrierReason {
-  FunctionReturnWithoutBlueprint,
-  DenseLogicalConsumer
-};
-
-struct TensorDemandAction {
-  TensorDemandActionKind kind = TensorDemandActionKind::DestinationFanout;
-  std::optional<ClassId> destinationClass;
-  std::optional<WholeTensorBarrierReason> barrierReason;
-};
-
-struct RunOutputDemand {
-  size_t resultIndex = 0;
-  mlir::Value originalOutput;
-  mlir::RankedTensorType fragmentType;
-  llvm::SmallVector<TensorDemandAction, 4> actions;
-};
-
-struct CompactRunPlan {
-  llvm::SmallVector<RunOutputDemand, 4> outputs;
-};
-
-enum class BatchInputDemandKind {
-  LaneFragment,
-  ProjectedFragment,
-  WholeTensorBarrier
-};
-
-struct BatchInputDemand {
-  BatchInputDemandKind kind = BatchInputDemandKind::LaneFragment;
-  std::optional<ProducerKey> wholeTensorProducer;
-};
-
-struct CloneIndexingContext {
-  std::optional<mlir::Value> runSlotIndex;
-  std::optional<mlir::Value> projectionSlotIndex;
-};
-
-struct MaterializerState;
-
-class AvailableValueStore {
-public:
-  struct ExactBatchFragmentRecord {
-    ProducerKey key;
-    mlir::Value value;
-  };
-
-  void record(ProducerKey key, ClassId classId, mlir::Value value) {
-    exactValues[key][classId] = value;
-
-    auto batch = mlir::dyn_cast_or_null<SpatComputeBatch>(key.instance.op);
-    if (!batch || key.instance.laneCount == 0)
-      return;
-
-    WholeBatchAssemblyLookupKey lookupKey {batch.getOperation(), key.resultIndex, classId};
-    llvm::SmallVector<ExactBatchFragmentRecord, 16>& bucket = exactBatchFragmentsByProducerResultClass[lookupKey];
-    for (ExactBatchFragmentRecord& record : bucket) {
-      if (!(record.key == key))
-        continue;
-      record.value = value;
-      return;
-    }
-    bucket.push_back({key, value});
-  }
-
-  void recordPackedRun(PackedScalarRunValue run) {
-    size_t runIndex = packedScalarRuns.size();
-    packedScalarRuns.push_back(std::move(run));
-    const PackedScalarRunValue& storedRun = packedScalarRuns[runIndex];
-    WholeBatchAssemblyLookupKey lookupKey {storedRun.sourceOp, storedRun.resultIndex, storedRun.targetClass};
-    packedRunsByProducerResultClass[lookupKey].push_back(runIndex);
-  }
-
-  void recordIndexedBatchRun(IndexedBatchRunValue run) { indexedBatchRuns.push_back(std::move(run)); }
-
-  std::optional<mlir::Value> lookupExact(ProducerKey key, ClassId classId) const;
-  std::optional<mlir::Value> lookup(MaterializerState& state, ProducerKey key, ClassId classId);
-  IndexedBatchRunValue* lookupIndexedBatchRun(ProducerKey key, ClassId classId);
-
-  llvm::ArrayRef<size_t> getPackedRunIndicesForWholeBatch(WholeBatchAssemblyLookupKey key) const {
-    auto it = packedRunsByProducerResultClass.find(key);
-    if (it == packedRunsByProducerResultClass.end())
-      return {};
-    return it->second;
-  }
-
-  llvm::ArrayRef<ExactBatchFragmentRecord> getExactFragmentsForWholeBatch(WholeBatchAssemblyLookupKey key) const {
-    auto it = exactBatchFragmentsByProducerResultClass.find(key);
-    if (it == exactBatchFragmentsByProducerResultClass.end())
-      return {};
-    return it->second;
-  }
-
-  PackedScalarRunValue& getPackedRun(size_t index) { return packedScalarRuns[index]; }
-
-private:
-  std::optional<mlir::Value> lookupPackedRun(MaterializerState& state, ProducerKey key, ClassId classId);
-
-  llvm::DenseMap<ProducerKey, llvm::DenseMap<ClassId, mlir::Value>, ProducerKeyInfo> exactValues;
-  llvm::SmallVector<PackedScalarRunValue, 8> packedScalarRuns;
-  llvm::SmallVector<IndexedBatchRunValue, 8> indexedBatchRuns;
-  llvm::DenseMap<WholeBatchAssemblyLookupKey,
-                 llvm::SmallVector<ExactBatchFragmentRecord, 16>,
-                 WholeBatchAssemblyLookupKeyInfo>
-    exactBatchFragmentsByProducerResultClass;
-  llvm::DenseMap<WholeBatchAssemblyLookupKey, llvm::SmallVector<size_t, 16>, WholeBatchAssemblyLookupKeyInfo>
-    packedRunsByProducerResultClass;
-};
-
-struct MaterializerState {
-  mlir::func::FuncOp func;
-  const MergeScheduleResult& schedule;
-  mlir::IRRewriter rewriter;
-  mlir::OperationFolder constantFolder;
-  int64_t& nextChannelId;
-  llvm::SmallVector<MaterializedClass, 8> classes;
-  llvm::DenseMap<CpuId, ClassId> cpuToClass;
-  llvm::DenseMap<CpuId, llvm::SmallVector<ComputeInstance, 32>> logicalInstancesByCpu;
-  llvm::DenseMap<ComputeInstance, LogicalSlotRange> scheduledInstanceToLogicalSlots;
-  llvm::DenseMap<ComputeInstance, ComputeInstance> logicalInstanceToScheduledChunk;
-  llvm::DenseSet<ClassSlotKey> materializedLogicalSlots;
-
-  llvm::DenseMap<ProducerKey, llvm::SmallVector<ClassId, 4>, ProducerKeyInfo> producerDestClasses;
-  llvm::DenseMap<SameClassConsumerLookupKey, llvm::SmallVector<ProducerKey, 4>, SameClassConsumerLookupKeyInfo>
-    sameClassConsumerIndex;
-  llvm::DenseMap<ProjectedBatchInputKey, AffineProjectedInputSliceMatch, ProjectedBatchInputKeyInfo>
-    projectedInputMatches;
-  llvm::DenseSet<ProjectedBatchInputKey, ProjectedBatchInputKeyInfo> nonProjectedInputs;
-  llvm::DenseMap<mlir::Value, bool> liveExternalUseCache;
-  llvm::DenseMap<mlir::Operation*, llvm::SmallVector<mlir::Type, 4>> batchOutputFragmentTypesCache;
-  llvm::DenseMap<ComputeInstance, llvm::SmallVector<mlir::Value, 4>, llvm::DenseMapInfo<ComputeInstance>>
-    computeInstanceOutputsCache;
-  llvm::DenseMap<ProducerKey, llvm::DenseMap<ClassId, ProjectedTransferDescriptor>, ProducerKeyInfo>
-    projectedTransfers;
-  llvm::DenseMap<mlir::Operation*, llvm::DenseMap<ClassId, ProjectedExtractReplacement>>
-    projectedExtractReplacements;
-  AvailableValueStore availableValues;
-  llvm::DenseMap<mlir::Value, mlir::Value> hostReplacements;
-  llvm::DenseMap<mlir::Value, ClassId> hostOutputOwners;
-  llvm::SmallVector<PendingProjectedHostOutputFragment, 32> pendingProjectedHostOutputFragments;
-  llvm::DenseSet<mlir::Operation*> oldComputeOps;
-
-  MaterializerState(mlir::func::FuncOp func, const MergeScheduleResult& schedule, int64_t& nextChannelId)
-  : func(func),
-    schedule(schedule),
-    rewriter(func.getContext()),
-    constantFolder(func.getContext()),
-    nextChannelId(nextChannelId) {}
-};
-
-} // namespace onnx_mlir::spatial

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

@@ -1,391 +1,120 @@
-#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 "SpatialDataflowCsvExporter.hpp"
+#include "DeferredCommunicationRealization.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/BatchCoreUtils.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/Common/Support/DebugDump.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::SpatGraphCompute;
-using SpatComputeBatch = spatial::SpatGraphComputeBatch;
 
-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;
-  };
-
-  //TODO Used for report refactor
-  struct CollectorConcatRow {
-    uint64_t computeId = 0;
-    int32_t coreId = -1;
-    uint64_t operandCount = 0;
-  };
-
-  uint64_t totalComputeOps = 0;
-  uint64_t totalLogicalComputes = 0;
-  uint64_t totalBatchComputeOps = 0;
-  uint64_t totalInstructionCount = 0;
-  uint64_t totalCrossbarCount = 0;
-  uint64_t nextBatchId = 0;
-  //TODO Used for report refactor
-  std::vector<ReportRow> collectedData;
-  //TODO Used for report refactor
-  std::vector<CollectorConcatRow> collectorConcatRows;
-
-  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<spatial::SpatScheduledCompute>(&op)) {
-      uint64_t numInst = spatial::countComputeBodyInstructions(spatCompute.getBody());
-      uint64_t perInstanceCrossbarCount = getPerInstanceCrossbarCount(spatCompute.getOperation());
-      SmallVector<int32_t> coreIds;
-      auto coreId = getOptionalScheduledCoreId(spatCompute, "merge compute core id");
-      if (failed(coreId))
-        return;
-      if (*coreId)
-        coreIds.push_back(**coreId);
-      uint64_t computeId = totalComputeOps++;
-      collectedData.push_back({computeId, 1, perInstanceCrossbarCount, numInst, false, coreIds});
-      uint64_t maxConcatOperands = 0;
-      spatCompute.getBody().walk([&](spatial::SpatConcatOp concatOp) {
-        maxConcatOperands = std::max<uint64_t>(maxConcatOperands, concatOp.getInputs().size());
-      });
-      //TODO 128 is a magic number
-      if (maxConcatOperands >= 128 && !coreIds.empty())
-        collectorConcatRows.push_back({computeId, coreIds.front(), maxConcatOperands});
-      totalLogicalComputes += 1;
-      totalInstructionCount += numInst;
-      totalCrossbarCount += perInstanceCrossbarCount;
-      continue;
-    }
-    if (auto batch = dyn_cast<spatial::SpatScheduledComputeBatch>(&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;
-      auto optionalCoreIds = getOptionalScheduledBatchCoreIds(batch, "merge compute_batch core id");
-      if (failed(optionalCoreIds))
-        return;
-      if (*optionalCoreIds)
-        coreIds = std::move(**optionalCoreIds);
-      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() || !collectorConcatRows.empty())
-    os << "\n";
-
-  if (!collectorConcatRows.empty()) {
-    os << "Collector concat materialization:\n";
-    for (const CollectorConcatRow& row : collectorConcatRows)
-      os << "\tmaterialization_kind = single_collector_concat, compute = " << row.computeId
-         << ", concat_operand_count = " << row.operandCount << ", collector_core = " << row.coreId << "\n";
-    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();
-    if (failed(verifyLogicalSpatialGraphInvariants(func))) {
-      func.emitOpError("logical Spatial graph verification failed at the start of MergeComputeNodes");
-      signalPassFailure();
-      return;
-    }
-    mergeTriviallyConnectedComputes(func);
-    if (failed(verifyLogicalSpatialGraphInvariants(func))) {
-      func.emitOpError("logical Spatial graph verification failed after trivial merge simplification");
+    ModuleOp moduleOp = getOperation();
+    auto entryFunc = getPimEntryFunc(moduleOp);
+    if (failed(entryFunc)) {
+      moduleOp.emitError("failed to locate the PIM entry function during MergeComputeNodes");
       signalPassFailure();
       return;
     }
 
-    const spatial::MergeScheduleResult* analysisResult = nullptr;
-    analysisResult = &getAnalysis<spatial::MergeSchedulingAnalysis>().getResult();
-    if (failed(spatial::MergeScheduleMaterializer().run(func, *analysisResult, nextChannelId))) {
+    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 (!sortTopologically(&func.getBody().front())) {
-      func.emitOpError("failed to topologically order merged Spatial IR");
+    // Phase 1 is intentionally dumped before its verifier: malformed deferred
+    // payloads must be diagnosed from the producer-owned body.
+    dumpModule(moduleOp, "spatial2_scheduled_no_comm", /*assumeVerified=*/true);
+
+    if (failed(verifyMaterializedScheduleMapping(funcOp,
+                                                 schedule,
+                                                 materialization->peftClassPlans,
+                                                 materialization->graphComputeToBlockMap,
+                                                 materialization->materializedSchedules))) {
+      moduleOp.emitError("scheduled Spatial materialization mapping verification failed");
       signalPassFailure();
       return;
     }
-    if (failed(verifyScheduledSpatialInvariants(func))) {
-      func.emitOpError("scheduled Spatial verification failed after merge materialization");
+    if (failed(verifyDeferredTransferPhase1Invariants(funcOp))) {
+      moduleOp.emitError("scheduled Spatial deferred communication verification failed");
       signalPassFailure();
       return;
     }
-    dumpModule(cast<ModuleOp>(func->getParentOp()), "spatial1_merged");
-    generateReport(func, "spatial_merge_report", analysisResult->cpuToLastComputeMap.size());
+    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;
+    }
+
+    SpatialDataflowExportStage exportMode = getSpatialDataflowExportStage();
+    if (shouldExportSpatialDataflowStage(exportMode, SpatialDataflowExportStage::Spatial2)
+        && failed(exportSpatialDataflowCsvScheduled(funcOp, "spatial2_scheduled_no_comm", "spatial2"))) {
+      signalPassFailure();
+      return;
+    }
+
+    dumpScheduledComputeReportAndModule(moduleOp,
+                                        funcOp,
+                                        schedule,
+                                        materialization->peftClassPlans,
+                                        materialization->materializedSchedules);
+    if (failed(realizeDeferredCommunication(funcOp))) {
+      moduleOp.emitError("MergeComputeNodes phase 2 communication realization failed");
+      signalPassFailure();
+      return;
+    }
+    dumpModule(moduleOp, "spatial3_scheduled", /*assumeVerified=*/true);
+    if (failed(verifyScheduledSpatialInvariants(funcOp))) {
+      moduleOp.emitError("scheduled Spatial phase 2 verification failed");
+      signalPassFailure();
+      return;
+    }
+    if (shouldExportSpatialDataflowStage(exportMode, SpatialDataflowExportStage::Spatial3)
+        && failed(exportSpatialDataflowCsvScheduled(funcOp, "spatial3_scheduled", "spatial3"))) {
+      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/MergeMessages.hpp

@@ -1,67 +0,0 @@
-#pragma once
-
-#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/SmallVector.h"
-#include "llvm/ADT/StringRef.h"
-
-#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
-
-namespace onnx_mlir::spatial {
-
-using CpuId = size_t;
-
-inline mlir::FailureOr<int32_t> getCheckedCoreId(mlir::Operation* anchor, CpuId cpu, llvm::StringRef fieldName) {
-  return pim::checkedI32(static_cast<uint64_t>(cpu), anchor, fieldName);
-}
-
-inline mlir::FailureOr<llvm::SmallVector<int32_t, 8>>
-getCheckedCoreIds(mlir::Operation* anchor, llvm::ArrayRef<CpuId> cpus, llvm::StringRef fieldName) {
-  llvm::SmallVector<int32_t, 8> coreIds;
-  coreIds.reserve(cpus.size());
-  for (CpuId cpu : cpus) {
-    auto checkedCoreId = getCheckedCoreId(anchor, cpu, fieldName);
-    if (mlir::failed(checkedCoreId))
-      return mlir::failure();
-    coreIds.push_back(*checkedCoreId);
-  }
-  return coreIds;
-}
-
-struct MessageVector {
-  llvm::SmallVector<int64_t, 16> channelIds;
-  llvm::SmallVector<int32_t, 16> sourceCoreIds;
-  llvm::SmallVector<int32_t, 16> targetCoreIds;
-
-  size_t size() const { return channelIds.size(); }
-  bool empty() const { return channelIds.empty(); }
-
-  mlir::LogicalResult verify(mlir::Operation* anchor) const {
-    if (channelIds.size() != sourceCoreIds.size() || channelIds.size() != targetCoreIds.size())
-      return anchor->emitError("message metadata is inconsistent");
-    return mlir::success();
-  }
-
-  void append(int64_t channelId, int32_t sourceCoreId, int32_t targetCoreId) {
-    channelIds.push_back(channelId);
-    sourceCoreIds.push_back(sourceCoreId);
-    targetCoreIds.push_back(targetCoreId);
-  }
-
-  void append(llvm::ArrayRef<int64_t> channels, llvm::ArrayRef<int32_t> sources, llvm::ArrayRef<int32_t> targets) {
-    assert(channels.size() == sources.size() && "channel/source count mismatch");
-    assert(channels.size() == targets.size() && "channel/target count mismatch");
-    llvm::append_range(channelIds, channels);
-    llvm::append_range(sourceCoreIds, sources);
-    llvm::append_range(targetCoreIds, targets);
-  }
-
-  MessageVector slice(size_t offset, size_t count) const {
-    MessageVector result;
-    result.append(llvm::ArrayRef<int64_t>(channelIds).slice(offset, count),
-                  llvm::ArrayRef<int32_t>(sourceCoreIds).slice(offset, count),
-                  llvm::ArrayRef<int32_t>(targetCoreIds).slice(offset, count));
-    return result;
-  }
-};
-
-} // namespace onnx_mlir::spatial

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

@@ -1,134 +0,0 @@
-#pragma once
-
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/SmallVector.h"
-
-#include <cstddef>
-#include <cstdint>
-#include <limits>
-#include <utility>
-
-#include "Scheduling/ComputeInstanceUtils.hpp"
-
-namespace onnx_mlir::spatial {
-
-using ClassId = size_t;
-using SlotId = size_t;
-
-struct ProducerKey {
-  ComputeInstance instance;
-  size_t resultIndex = 0;
-
-  bool operator==(const ProducerKey& other) const {
-    return instance == other.instance && resultIndex == other.resultIndex;
-  }
-};
-
-struct ProducerKeyInfo {
-  static ProducerKey getEmptyKey() {
-    return {llvm::DenseMapInfo<ComputeInstance>::getEmptyKey(), std::numeric_limits<size_t>::max()};
-  }
-
-  static ProducerKey getTombstoneKey() {
-    return {llvm::DenseMapInfo<ComputeInstance>::getTombstoneKey(), std::numeric_limits<size_t>::max()};
-  }
-
-  static unsigned getHashValue(const ProducerKey& key) {
-    return llvm::hash_combine(llvm::DenseMapInfo<ComputeInstance>::getHashValue(key.instance), key.resultIndex);
-  }
-
-  static bool isEqual(const ProducerKey& lhs, const ProducerKey& rhs) { return lhs == rhs; }
-};
-
-struct SameClassConsumerLookupKey {
-  mlir::Operation* sourceOp = nullptr;
-  size_t resultIndex = 0;
-  ClassId classId = 0;
-
-  bool operator==(const SameClassConsumerLookupKey& other) const {
-    return sourceOp == other.sourceOp && resultIndex == other.resultIndex && classId == other.classId;
-  }
-};
-
-struct SameClassConsumerLookupKeyInfo {
-  static SameClassConsumerLookupKey getEmptyKey() {
-    return {llvm::DenseMapInfo<mlir::Operation*>::getEmptyKey(), std::numeric_limits<size_t>::max(),
-            std::numeric_limits<ClassId>::max()};
-  }
-
-  static SameClassConsumerLookupKey getTombstoneKey() {
-    return {llvm::DenseMapInfo<mlir::Operation*>::getTombstoneKey(), std::numeric_limits<size_t>::max(),
-            std::numeric_limits<ClassId>::max()};
-  }
-
-  static unsigned getHashValue(const SameClassConsumerLookupKey& key) {
-    return llvm::hash_combine(llvm::DenseMapInfo<mlir::Operation*>::getHashValue(key.sourceOp),
-                              key.resultIndex,
-                              key.classId);
-  }
-
-  static bool isEqual(const SameClassConsumerLookupKey& lhs, const SameClassConsumerLookupKey& rhs) {
-    return lhs == rhs;
-  }
-};
-
-struct WholeBatchAssemblyLookupKey {
-  mlir::Operation* sourceOp = nullptr;
-  size_t resultIndex = 0;
-  ClassId classId = 0;
-
-  bool operator==(const WholeBatchAssemblyLookupKey& other) const {
-    return sourceOp == other.sourceOp && resultIndex == other.resultIndex && classId == other.classId;
-  }
-};
-
-struct WholeBatchAssemblyLookupKeyInfo {
-  static WholeBatchAssemblyLookupKey getEmptyKey() {
-    return {llvm::DenseMapInfo<mlir::Operation*>::getEmptyKey(), std::numeric_limits<size_t>::max(),
-            std::numeric_limits<ClassId>::max()};
-  }
-
-  static WholeBatchAssemblyLookupKey getTombstoneKey() {
-    return {llvm::DenseMapInfo<mlir::Operation*>::getTombstoneKey(), std::numeric_limits<size_t>::max(),
-            std::numeric_limits<ClassId>::max()};
-  }
-
-  static unsigned getHashValue(const WholeBatchAssemblyLookupKey& key) {
-    return llvm::hash_combine(llvm::DenseMapInfo<mlir::Operation*>::getHashValue(key.sourceOp),
-                              key.resultIndex,
-                              key.classId);
-  }
-
-  static bool isEqual(const WholeBatchAssemblyLookupKey& lhs, const WholeBatchAssemblyLookupKey& rhs) {
-    return lhs == rhs;
-  }
-};
-
-using ClassSlotKey = std::pair<ClassId, SlotId>;
-
-struct ProjectedBatchInputKey {
-  mlir::Operation* consumerOp = nullptr;
-  unsigned inputIndex = 0;
-
-  bool operator==(const ProjectedBatchInputKey& other) const {
-    return consumerOp == other.consumerOp && inputIndex == other.inputIndex;
-  }
-};
-
-struct ProjectedBatchInputKeyInfo {
-  static ProjectedBatchInputKey getEmptyKey() {
-    return {llvm::DenseMapInfo<mlir::Operation*>::getEmptyKey(), std::numeric_limits<unsigned>::max()};
-  }
-
-  static ProjectedBatchInputKey getTombstoneKey() {
-    return {llvm::DenseMapInfo<mlir::Operation*>::getTombstoneKey(), std::numeric_limits<unsigned>::max()};
-  }
-
-  static unsigned getHashValue(const ProjectedBatchInputKey& key) {
-    return llvm::hash_combine(key.consumerOp, key.inputIndex);
-  }
-
-  static bool isEqual(const ProjectedBatchInputKey& lhs, const ProjectedBatchInputKey& rhs) { return lhs == rhs; }
-};
-
-} // namespace onnx_mlir::spatial

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

@@ -1,104 +0,0 @@
-#include "ProjectedFragments.hpp"
-
-#include "mlir/IR/BuiltinTypes.h"
-
-namespace onnx_mlir::spatial {
-
-static mlir::FailureOr<mlir::RankedTensorType> getPackedBatchTensorType(mlir::Type laneType, size_t laneCount) {
-  auto tensorType = mlir::dyn_cast<mlir::RankedTensorType>(laneType);
-  if (!tensorType || !tensorType.hasStaticShape() || tensorType.getRank() == 0)
-    return mlir::failure();
-
-  llvm::SmallVector<int64_t, 4> shape(tensorType.getShape());
-  shape[0] *= static_cast<int64_t>(laneCount);
-  return mlir::RankedTensorType::get(shape, tensorType.getElementType(), tensorType.getEncoding());
-}
-
-unsigned getProjectedFragmentsPerLogicalSlot(llvm::ArrayRef<int64_t> loopTripCounts) {
-  unsigned fragmentsPerLogicalSlot = 1;
-  for (int64_t tripCount : loopTripCounts) {
-    assert(tripCount > 0 && "projected loop trip counts must be positive");
-    fragmentsPerLogicalSlot *= static_cast<unsigned>(tripCount);
-  }
-  return fragmentsPerLogicalSlot;
-}
-
-mlir::LogicalResult verifyProjectedFragmentLayout(mlir::Operation* anchor, const ProjectedFragmentLayout& layout) {
-  if (!layout.fragmentType || layout.fragmentShape.empty())
-    return anchor->emitError("projected fragment layout is missing fragment type metadata");
-  if (layout.fragmentShape.size() != static_cast<size_t>(layout.fragmentType.getRank()))
-    return anchor->emitError("projected fragment layout rank does not match fragment type");
-  if (layout.payloadFragmentCount == 0 || layout.fragmentsPerLogicalSlot == 0)
-    return anchor->emitError("projected fragment layout has an invalid fragment count");
-  if (layout.payloadFragmentCount % layout.fragmentsPerLogicalSlot != 0)
-    return anchor->emitError("projected fragment layout payload fragment count is incompatible with logical slots");
-  return mlir::success();
-}
-
-mlir::FailureOr<mlir::RankedTensorType>
-getProjectedPayloadType(mlir::Operation* anchor, mlir::RankedTensorType fragmentType, unsigned payloadFragmentCount) {
-  auto packedType = getPackedBatchTensorType(fragmentType, payloadFragmentCount);
-  if (mlir::failed(packedType)) {
-    anchor->emitError("cannot create projected payload type");
-    return mlir::failure();
-  }
-  return *packedType;
-}
-
-llvm::SmallVector<llvm::SmallVector<int64_t, 16>, 4>
-buildProjectedFragmentOffsetsByDim(llvm::ArrayRef<llvm::SmallVector<int64_t, 4>> fragmentOffsets, size_t rank) {
-  llvm::SmallVector<llvm::SmallVector<int64_t, 16>, 4> fragmentOffsetsByDim(rank);
-  for (llvm::ArrayRef<int64_t> offsets : fragmentOffsets) {
-    assert(offsets.size() == rank && "projected offset rank mismatch");
-    for (size_t dim = 0; dim < rank; ++dim)
-      fragmentOffsetsByDim[dim].push_back(offsets[dim]);
-  }
-  return fragmentOffsetsByDim;
-}
-
-mlir::LogicalResult verifyProjectedTransferDescriptor(mlir::Operation* anchor,
-                                                      const ProjectedTransferDescriptor& descriptor) {
-  if (mlir::failed(verifyProjectedFragmentLayout(anchor, descriptor.layout)))
-    return mlir::failure();
-  if (!descriptor.payloadType)
-    return anchor->emitError("projected transfer descriptor is missing payload type");
-  if (descriptor.fragmentOffsets.empty())
-    return anchor->emitError("projected transfer descriptor expected at least one fragment offset");
-  if (descriptor.fragmentOffsetsByDim.size() != descriptor.layout.fragmentShape.size())
-    return anchor->emitError("projected transfer descriptor dimension-major offsets are inconsistent");
-  for (llvm::ArrayRef<int64_t> dimOffsets : descriptor.fragmentOffsetsByDim)
-    if (dimOffsets.size() != descriptor.fragmentOffsets.size())
-      return anchor->emitError("projected transfer descriptor dimension-major offsets are inconsistent");
-  for (llvm::ArrayRef<int64_t> offsets : descriptor.fragmentOffsets)
-    if (offsets.size() != descriptor.layout.fragmentShape.size())
-      return anchor->emitError("projected transfer offset rank does not match fragment rank");
-  return mlir::success();
-}
-
-mlir::LogicalResult verifyProjectedSendDescriptor(mlir::Operation* anchor,
-                                                  const ProjectedTransferDescriptor& descriptor,
-                                                  const MessageVector& messages) {
-  if (mlir::failed(verifyProjectedTransferDescriptor(anchor, descriptor)))
-    return mlir::failure();
-  if (messages.size() * descriptor.layout.payloadFragmentCount != descriptor.fragmentOffsets.size())
-    return anchor->emitError("projected send descriptor metadata is inconsistent");
-  return mlir::success();
-}
-
-mlir::LogicalResult finalizeProjectedTransferDescriptor(mlir::Operation* anchor,
-                                                        ProjectedTransferDescriptor& descriptor) {
-  descriptor.fragmentOffsetsByDim =
-    buildProjectedFragmentOffsetsByDim(descriptor.fragmentOffsets, descriptor.layout.fragmentShape.size());
-
-  auto payloadType =
-    getProjectedPayloadType(anchor, descriptor.layout.fragmentType, descriptor.layout.payloadFragmentCount);
-  if (mlir::failed(payloadType))
-    return mlir::failure();
-  if (descriptor.payloadType && descriptor.payloadType != *payloadType)
-    return anchor->emitError("projected transfer descriptor payload type does not match projected layout");
-  descriptor.payloadType = *payloadType;
-
-  return verifyProjectedTransferDescriptor(anchor, descriptor);
-}
-
-} // namespace onnx_mlir::spatial

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

@@ -1,87 +0,0 @@
-#pragma once
-
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/IR/BuiltinAttributes.h"
-#include "mlir/IR/BuiltinTypes.h"
-#include "mlir/IR/Value.h"
-#include "mlir/IR/ValueRange.h"
-
-#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/SmallVector.h"
-
-#include <cstdint>
-
-#include "MergeMessages.hpp"
-#include "MergeScheduleKeys.hpp"
-
-namespace onnx_mlir::spatial {
-
-struct ProjectedFragmentLayout {
-  mlir::RankedTensorType fragmentType;
-  llvm::SmallVector<int64_t, 4> fragmentShape;
-  unsigned fragmentsPerLogicalSlot = 1;
-  unsigned payloadFragmentCount = 1;
-  llvm::SmallVector<int64_t, 4> loopLowerBounds;
-  llvm::SmallVector<int64_t, 4> loopSteps;
-  llvm::SmallVector<int64_t, 4> loopTripCounts;
-};
-
-struct StaticProjectedLoopInfo {
-  mlir::BlockArgument iv;
-  int64_t lowerBound = 0;
-  int64_t step = 1;
-  int64_t tripCount = 1;
-};
-
-struct ProjectedTransferDescriptor {
-  ProjectedBatchInputKey inputKey;
-  mlir::Operation* extractOp = nullptr;
-  ProjectedFragmentLayout layout;
-  mlir::RankedTensorType payloadType;
-  llvm::SmallVector<llvm::SmallVector<int64_t, 4>, 16> fragmentOffsets;
-  llvm::SmallVector<llvm::SmallVector<int64_t, 16>, 4> fragmentOffsetsByDim;
-};
-
-struct ProjectedExtractReplacement {
-  mlir::Value payload;
-  ProjectedFragmentLayout layout;
-};
-
-struct PendingProjectedHostOutputFragment {
-  mlir::Value originalOutput;
-  ClassId sourceClass = 0;
-  ProducerKey producerKey;
-  unsigned publicationResultIndex = 0;
-  int64_t sourceFragmentOrdinal = 0;
-  int64_t sourceElementOffset = 0;
-  llvm::SmallVector<int64_t, 4> offsets;
-  llvm::SmallVector<int64_t, 4> sizes;
-  llvm::SmallVector<int64_t, 4> strides;
-  uint32_t sourceLane = 0;
-  mlir::Location loc;
-};
-
-struct AffineProjectedInputSliceMatch {
-  mlir::tensor::ExtractSliceOp extract;
-  mlir::RankedTensorType sourceType;
-  mlir::RankedTensorType fragmentType;
-  llvm::SmallVector<int64_t, 4> fragmentShape;
-  llvm::SmallVector<mlir::OpFoldResult, 4> offsets;
-  llvm::SmallVector<StaticProjectedLoopInfo, 4> loops;
-};
-
-unsigned getProjectedFragmentsPerLogicalSlot(llvm::ArrayRef<int64_t> loopTripCounts);
-mlir::LogicalResult verifyProjectedFragmentLayout(mlir::Operation* anchor, const ProjectedFragmentLayout& layout);
-mlir::FailureOr<mlir::RankedTensorType>
-getProjectedPayloadType(mlir::Operation* anchor, mlir::RankedTensorType fragmentType, unsigned payloadFragmentCount);
-llvm::SmallVector<llvm::SmallVector<int64_t, 16>, 4>
-buildProjectedFragmentOffsetsByDim(llvm::ArrayRef<llvm::SmallVector<int64_t, 4>> fragmentOffsets, size_t rank);
-mlir::LogicalResult verifyProjectedTransferDescriptor(mlir::Operation* anchor,
-                                                      const ProjectedTransferDescriptor& descriptor);
-mlir::LogicalResult verifyProjectedSendDescriptor(mlir::Operation* anchor,
-                                                  const ProjectedTransferDescriptor& descriptor,
-                                                  const MessageVector& messages);
-mlir::LogicalResult finalizeProjectedTransferDescriptor(mlir::Operation* anchor,
-                                                        ProjectedTransferDescriptor& descriptor);
-
-} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,997 @@
+#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>
+
+#include "llvm/ADT/SmallPtrSet.h"
+
+namespace onnx_mlir {
+namespace spatial {
+
+using namespace mlir;
+namespace {
+
+struct BatchFragmentSpec {
+  RankedTensorType resultType;
+  RankedTensorType sourceSliceType;
+};
+
+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() != sourceType.getRank() + 1 || destType.getDimSize(0) != batch.getLaneCount()
+        || destType.getElementType() != sourceType.getElementType())
+      return batch.emitOpError("graph_compute_batch result must be a leading physical-slot dimension followed by its fragment");
+    if (!llvm::equal(destType.getShape().drop_front(), sourceType.getShape()))
+      return batch.emitOpError("graph_compute_batch result trailing shape must match its published fragment");
+    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");
+    if (!isa<Value>(offsets.front()) || !valueTransitivelyDependsOn(cast<Value>(offsets.front()), *batch.getLaneArgument()))
+      return batch.emitOpError("graph_compute_batch publication must select its physical slot in dimension zero");
+    for (unsigned dim = 1; dim < offsets.size(); ++dim) {
+      auto offset = dyn_cast<Attribute>(offsets[dim]);
+      auto integer = dyn_cast_or_null<IntegerAttr>(offset);
+      if (!integer || integer.getInt() != 0)
+        return batch.emitOpError("graph_compute_batch publication must have zero trailing offsets");
+    }
+    auto staticIndex = [](OpFoldResult value) -> std::optional<int64_t> {
+      auto attr = dyn_cast<Attribute>(value);
+      auto integer = dyn_cast_or_null<IntegerAttr>(attr);
+      return integer ? std::optional<int64_t>(integer.getInt()) : std::nullopt;
+    };
+    if (staticIndex(sizes.front()) != 1)
+      return batch.emitOpError("graph_compute_batch publication sizes must be [1] plus the fragment shape");
+    for (auto [size, dim] : llvm::zip_equal(ArrayRef<OpFoldResult>(sizes).drop_front(), sourceType.getShape()))
+      if (staticIndex(size) != dim)
+        return batch.emitOpError("graph_compute_batch publication sizes must be [1] plus the fragment shape");
+    return BatchFragmentSpec {spatial::getGraphBatchPhysicalResultType(fragmentLaneCount, sourceType), sourceType};
+  }
+
+  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) {
+  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 createOrFoldAffineApply(builder, loc, expr, ValueRange {scheduledLane}, constantAnchor);
+  }
+
+  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) &&
+            !isDeferredFragmentAssemblyInput(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) &&
+            !isDeferredFragmentAssemblyInput(input, instance))
+          appendUnique(peftClassPlan.inputs, input);
+    }
+  }
+
+  return success();
+}
+
+static void cloneComputeBody(OpBuilder &builder, Block &source, IRMapping &mapper,
+                             SmallVectorImpl<Value> &yieldedValues,
+                             const llvm::SmallPtrSetImpl<Operation *> &absorbed) {
+  for (Operation &op : source.without_terminator())
+    if (!absorbed.contains(&op))
+      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,
+                                                            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 requires a physical graph fragment publication");
+    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);
+      Value localLane = arith::SubIOp::create(builder,
+                                              bodyLoc,
+                                              originalLane,
+                                              getOrCreateIndexConstant(rewriter, batch.getOperation(), instance.laneStart))
+                          .getResult();
+      for (auto [index, weight] : llvm::enumerate(batch.getWeights()))
+        mapper.map(*batch.getWeightArgument(index), getBlockOperand(block, scheduledWeights, weight));
+      SmallVector<DeferredInputPlan> inputPlans;
+      for (auto [index, input] : llvm::enumerate(batch.getInputs())) {
+        DeferredInputPlan plan;
+        if (failed(prepareSingleCpuInput(builder,
+                                     input.getLoc(),
+                                     input,
+                                     *batch.getInputArgument(index),
+                                     instance,
+                                     schedule,
+                                     scheduledInputs,
+                                     block,
+                                     scheduledWeights.size(),
+                                     ArrayRef<ProducerValueKey> {},
+                                     *batch.getLaneArgument(),
+                                     originalLane,
+                                     plan)))
+          return failure();
+        plan.scalarizedLocalLane = localLane;
+        plan.scalarizedGraphLaneBase = lower;
+        plan.scalarizedLaneCount = instance.laneCount;
+        plan.scalarizedHoistBlock = &block;
+        inputPlans.push_back(std::move(plan));
+      }
+      for (auto [index, outputArg] : llvm::enumerate(batch.getOutputs()))
+        (void)outputArg, mapper.map(*batch.getOutputArgument(index), iterArgs[index]);
+
+      Block &source = batch.getBody().front();
+      llvm::SmallPtrSet<Operation *, 32> absorbed;
+      if (failed(materializeDeferredPayloadDemands(builder, bodyLoc, source, inputPlans, mapper, absorbed)))
+        return failure();
+      for (Operation &op : source.without_terminator())
+        if (!absorbed.contains(&op))
+          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();
+      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 requires a physical graph fragment publication");
+        auto oldDest = dyn_cast<BlockArgument>(insert.getDest());
+        if (!oldDest || !outputIndexByArg.count(oldDest))
+          return batch.emitOpError("scheduled materialization requires a physical graph fragment publication"), failure();
+        size_t resultIndex = outputIndexByArg.lookup(oldDest);
+        SmallVector<OpFoldResult> offsets = remapMixedOffsets(insert.getMixedOffsets(), mapper);
+        offsets.front() = 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) {
+  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));
+      SmallVector<DeferredInputPlan> inputPlans;
+      for (auto [index, input] : llvm::enumerate(compute.getInputs())) {
+        DeferredInputPlan plan;
+        if (failed(prepareSingleCpuInput(rewriter,
+                                     input.getLoc(),
+                                     input,
+                                     *compute.getInputArgument(index),
+                                     instance,
+                                     schedule,
+                                     scheduled.getInputs(),
+                                     *block,
+                                     scheduled.getWeights().size(),
+                                     carriedKeys,
+                                     {},
+                                     {},
+                                     plan)))
+          return failure();
+        inputPlans.push_back(std::move(plan));
+      }
+      llvm::SmallPtrSet<Operation *, 32> absorbed;
+      if (failed(materializeDeferredPayloadDemands(rewriter, compute.getLoc(), compute.getBody().front(), inputPlans, mapper, absorbed)))
+        return failure();
+      cloneComputeBody(rewriter, compute.getBody().front(), mapper, yieldedValues, absorbed);
+    } else {
+      auto batch = cast<SpatComputeBatch>(instance.op);
+      if (failed(materializeResultfulBatchChunkAsScalar(rewriter,
+                                                        batch,
+                                                        instance,
+                                                        scheduled.getWeights(),
+                                                        scheduled.getInputs(),
+                                                        *block,
+                                                        schedule,
+                                                        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,
+                                                                   Operation *constantAnchor) {
+  SmallVector<OpFoldResult> offsets;
+  Value scheduledOutputLane = scheduledLane;
+  if (lanesPerScheduledLane != 1) {
+    scheduledOutputLane = affineMulConst(
+      builder, loc, scheduledLane, lanesPerScheduledLane, constantAnchor);
+  }
+  offsets.push_back(scheduledOutputLane);
+  offsets.append(localFragmentType.getRank() - 1, 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) {
+  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;
+
+    if (auto compute = dyn_cast<SpatCompute>(representative.op)) {
+      IRMapping mapper;
+      for (auto [index, weight] : llvm::enumerate(compute.getWeights()))
+        mapper.map(*compute.getWeightArgument(index),
+                   getBlockOperand(*block, scheduled.getWeights(), weight, 1));
+      unsigned firstInputArg = 1 + scheduled.getWeights().size();
+      SmallVector<DeferredInputPlan> inputPlans;
+      for (auto [index, input] : llvm::enumerate(compute.getInputs())) {
+        DeferredInputPlan plan;
+        if (failed(prepareMultiCpuTupleInput(rewriter,
+                                         input.getLoc(),
+                                         input,
+                                         *compute.getInputArgument(index),
+                                         stepTuple,
+                                         peftClassPlan,
+                                         schedule,
+                                         scheduled.getInputs(),
+                                         *block,
+                                         firstInputArg,
+                                         {},
+                                         {},
+                                         scheduledLane,
+                                         plan)))
+          return failure();
+        inputPlans.push_back(std::move(plan));
+      }
+      SmallVector<Value> yieldedValues;
+      llvm::SmallPtrSet<Operation *, 32> absorbed;
+      if (failed(materializeDeferredPayloadDemands(rewriter, compute.getLoc(), compute.getBody().front(), inputPlans, mapper, absorbed)))
+        return failure();
+      cloneComputeBody(rewriter, compute.getBody().front(), mapper, yieldedValues, absorbed);
+      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());
+      }
+    } 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);
+      FailureOr<Value> sourceLaneStart =
+        buildSourceLaneStartForScheduledLane(rewriter, batch.getLoc(), scheduledLane, sourceLaneSelector, scheduled.getOperation());
+      if (failed(sourceLaneStart))
+        return failure();
+      auto loop = buildNormalizedScfFor(
+        rewriter,
+        batch.getLoc(),
+        lower,
+        upper,
+        step,
+        localFragments,
+        [&](OpBuilder &builder, Location bodyLoc, Value innerLane, ValueRange iterArgs, SmallVectorImpl<Value> &yielded) -> LogicalResult {
+
+          IRMapping mapper;
+          Value sourceLane = createOrFoldAffineApply(
+            builder,
+            bodyLoc,
+            builder.getAffineDimExpr(0) + builder.getAffineDimExpr(1),
+            ValueRange {*sourceLaneStart, innerLane},
+            scheduled.getOperation());
+          mapper.map(*batch.getLaneArgument(), sourceLane);
+          for (auto [index, weight] : llvm::enumerate(batch.getWeights()))
+            mapper.map(*batch.getWeightArgument(index),
+                       getBlockOperand(*block, scheduled.getWeights(), weight, 1));
+          unsigned firstInputArg = 1 + scheduled.getWeights().size();
+          SmallVector<DeferredInputPlan> inputPlans;
+          for (auto [index, input] : llvm::enumerate(batch.getInputs())) {
+            DeferredInputPlan plan;
+            if (failed(prepareMultiCpuTupleInput(builder,
+                                             input.getLoc(),
+                                             input,
+                                             *batch.getInputArgument(index),
+                                             stepTuple,
+                                             peftClassPlan,
+                                             schedule,
+                                             scheduled.getInputs(),
+                                             *block,
+                                             firstInputArg,
+                                             *batch.getLaneArgument(),
+                                             sourceLane,
+                                             scheduledLane,
+                                             plan)))
+              return failure();
+            plan.scalarizedLocalLane = innerLane;
+            plan.scalarizedGraphLaneBase = *sourceLaneStart;
+            plan.scalarizedLaneCount = representative.laneCount;
+            plan.scalarizedHoistBlock = block;
+            inputPlans.push_back(std::move(plan));
+          }
+          for (unsigned index = 0; index < batch.getNumResults(); ++index)
+            mapper.map(*batch.getOutputArgument(index), iterArgs[index]);
+          llvm::SmallPtrSet<Operation *, 32> absorbed;
+          if (failed(materializeDeferredPayloadDemands(builder, bodyLoc, batch.getBody().front(), inputPlans, mapper, absorbed)))
+            return failure();
+          for (Operation &op : batch.getBody().front().without_terminator())
+            if (!absorbed.contains(&op))
+              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 requires a physical graph fragment publication");
+            auto oldDest = dyn_cast<BlockArgument>(insert.getDest());
+            if (!oldDest || !outputIndexByArg.count(oldDest))
+              return batch.emitOpError("scheduled materialization requires a physical graph fragment publication"), failure();
+            size_t resultIndex = outputIndexByArg.lookup(oldDest);
+            SmallVector<OpFoldResult> offsets = remapMixedOffsets(insert.getMixedOffsets(), mapper);
+            offsets.front() = 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());
+    }
+
+    struct Publication {
+      Value fragment;
+      SmallVector<OpFoldResult> offsets;
+      SmallVector<OpFoldResult> sizes;
+      SmallVector<OpFoldResult> strides;
+    };
+    SmallVector<Publication> publications;
+    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,
+        scheduled.getOperation());
+      SmallVector<OpFoldResult> sizes;
+      SmallVector<OpFoldResult> strides;
+      for (int64_t dim : localFragmentType.getShape()) {
+        sizes.push_back(rewriter.getIndexAttr(dim));
+        strides.push_back(rewriter.getIndexAttr(1));
+      }
+      publications.push_back(
+        {localFragment, std::move(offsets), std::move(sizes),
+         std::move(strides)});
+    }
+    auto inParallel = SpatInParallelOp::create(rewriter, scheduled.getLoc());
+    rewriter.setInsertionPointToStart(&inParallel.getRegion().front());
+    for (auto [resultIndex, publication] : llvm::enumerate(publications))
+      tensor::ParallelInsertSliceOp::create(
+        rewriter,
+        scheduled.getLoc(),
+        publication.fragment,
+        block->getArgument(getScheduledBatchResultArgBase(scheduled) + stepPlan.resultOffset + resultIndex),
+        publication.offsets,
+        publication.sizes,
+        publication.strides);
+  }
+  return success();
+}
+
+
+} // namespace
+
+
+FailureOr<ScheduledComputeMaterializationResult>
+materializeScheduledCompute(func::FuncOp funcOp,
+                            const MergeScheduleResult &schedule,
+                            PatternRewriter &rewriter) {
+  DenseMap<Operation *, int64_t> graphIds;
+  int64_t nextGraphId = 0;
+  for (Operation &op : funcOp.getOps())
+    if (isa<SpatGraphCompute, SpatGraphComputeBatch>(op)) {
+      graphIds[&op] = nextGraphId;
+      op.setAttr("scheduled.graph_id", rewriter.getI64IntegerAttr(nextGraphId++));
+    }
+
+  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 std::tie(graphIds.find(lhs.op)->second,
+                        schedule.computeToCpuSlotMap.find(lhs)->second) <
+               std::tie(graphIds.find(rhs.op)->second,
+                        schedule.computeToCpuSlotMap.find(rhs)->second);
+      });
+    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;
+
+  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;
+      SmallVector<int64_t> stepSourceIds;
+      size_t resultOffset = 0;
+      for (const ComputeInstance &instance : peftClassPlan.instancesByCpu.lookup(peftClassPlan.cpus.front())) {
+        stepSources.push_back(rewriter.getStringAttr(getInstanceName(instance)));
+        stepSourceIds.push_back(graphIds.lookup(instance.op));
+        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_source_ids", rewriter.getDenseI64ArrayAttr(stepSourceIds));
+      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;
+      SmallVector<int64_t> stepSourceIds;
+      for (const ScheduledStepPlan &stepPlan : record.stepPlans) {
+        stepSources.push_back(rewriter.getStringAttr(getInstanceName(stepPlan.stepTuple.instances.front())));
+        stepSourceIds.push_back(graphIds.lookup(stepPlan.stepTuple.instances.front().op));
+        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_source_ids", rewriter.getDenseI64ArrayAttr(stepSourceIds));
+      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)))
+        return failure();
+    } else {
+      if (failed(materializeMultiCpuPeftClass(rewriter,
+                                              scheduledComputeBatches.lookup(peftClassPlan.canonicalClassId),
+                                              peftClassPlan,
+                                              schedule,
+                                              graphComputeToBlockMap,
+                                              record)))
+        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,320 @@
+#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"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.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 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) {
+  RankedTensorType tableType = RankedTensorType::get({static_cast<int64_t>(values.size())}, builder.getI64Type());
+  DenseElementsAttr tableAttr = DenseElementsAttr::get(tableType, values);
+  return getOrCreateConstant(builder, constantAnchor, tableAttr, tableType);
+}
+
+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,304 @@
+#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 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 (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 << "  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,296 @@
+#include "ScheduledComputeVerification.hpp"
+#include "DeferredProjectionAnalysis.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;
+  GraphBatchPublicationCache publicationCache;
+  funcOp.walk([&](SpatDeferredCommunicationOp transfer) {
+    bool ownershipValid = true;
+    for (Value source : transfer.getSources()) {
+      auto result = dyn_cast<OpResult>(source);
+      if (!result || !isa<SpatGraphCompute, SpatGraphComputeBatch>(result.getOwner())) {
+        ownershipValid = false;
+        diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) {
+          illegalOp->emitOpError("phase-check deferred communication source operand must be an original graph SSA result");
+        });
+      }
+    }
+    if (!transfer->getParentOfType<SpatScheduledCompute>() &&
+        !transfer->getParentOfType<SpatScheduledComputeBatch>()) {
+      ownershipValid = false;
+      diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) {
+        illegalOp->emitOpError("phase-check deferred communication must be inside a scheduled compute");
+        });
+    }
+    if (!ownershipValid)
+      return;
+    if (failed(verifyDeferredProgramContract(transfer))) {
+      diagnostics.report(transfer.getOperation(), [&](Operation *) {});
+      return;
+    }
+    for (Value source : transfer.getSources()) {
+      auto result = dyn_cast<OpResult>(source);
+      auto batch = result
+        ? dyn_cast<SpatGraphComputeBatch>(result.getOwner())
+        : SpatGraphComputeBatch();
+      if (batch && failed(getGraphBatchPublicationMap(
+                     batch, result.getResultNumber(), publicationCache)))
+        diagnostics.report(transfer.getOperation(), [&](Operation *) {});
+    }
+  });
+
+  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/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/SpatialDataflowCsvExporter.cpp

@@ -0,0 +1,898 @@
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/AsmState.h"
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/Value.h"
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <cstdint>
+#include <fstream>
+#include <optional>
+#include <string>
+#include <utility>
+
+#include "SpatialDataflowCsvExporter.hpp"
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
+#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace spatial {
+
+namespace {
+
+struct TopLevelOpInfo {
+  Operation* op = nullptr;
+  size_t opId = 0;
+  bool isScheduled = false;
+  std::optional<int32_t> scalarCore;
+};
+
+struct ExpandedNodeInfo {
+  std::string id;
+  std::optional<int32_t> core;
+  std::optional<uint32_t> lane;
+};
+
+struct ChannelSendRecord {
+  std::string sourceId;
+  std::optional<uint32_t> sourceLane;
+};
+
+enum class LogicalNodeSelector {
+  Scalar,
+  Lane,
+  RangeRepresentative,
+};
+
+struct ResolvedProducer {
+  Operation* op = nullptr;
+  size_t resultIndex = 0;
+  LogicalNodeSelector selector = LogicalNodeSelector::Scalar;
+  uint32_t lane = 0;
+  uint32_t laneStart = 0;
+  uint32_t laneCount = 1;
+};
+
+struct EdgeSource {
+  std::string id;
+  std::optional<uint32_t> sourceLane;
+};
+
+using ScheduledNodeByGraphLane = DenseMap<std::pair<int64_t, uint32_t>, ExpandedNodeInfo>;
+
+void emitEdgeRow(std::fstream& edgesFile,
+                 StringRef sourceId,
+                 StringRef targetId,
+                 std::optional<uint64_t> byteSize,
+                 Type propagatedType,
+                 StringRef stage,
+                 std::optional<uint32_t> sourceLane,
+                 std::optional<uint32_t> targetLane,
+                 std::optional<int64_t> channelId);
+
+std::string csvEscape(StringRef field) {
+  bool needsQuotes = field.contains(',') || field.contains('"') || field.contains('\n') || field.contains('\r');
+  if (!needsQuotes)
+    return field.str();
+
+  std::string escaped;
+  escaped.reserve(field.size() + 2);
+  escaped.push_back('"');
+  for (char ch : field)
+    if (ch == '"')
+      escaped += "\"\"";
+    else
+      escaped.push_back(ch);
+  escaped.push_back('"');
+  return escaped;
+}
+
+void writeCsvRow(std::fstream& file, ArrayRef<std::string> fields) {
+  for (size_t i = 0; i < fields.size(); ++i) {
+    if (i != 0)
+      file << ",";
+    file << csvEscape(fields[i]);
+  }
+  file << "\n";
+}
+
+template <typename NumberT>
+std::string maybeNumber(std::optional<NumberT> value) {
+  if (!value)
+    return "";
+  return std::to_string(*value);
+}
+
+std::string stringifyType(Type type) {
+  std::string storage;
+  llvm::raw_string_ostream os(storage);
+  type.print(os);
+  return os.str();
+}
+
+std::string stringifyValueAsOperand(Value value, AsmState& asmState) {
+  std::string storage;
+  llvm::raw_string_ostream os(storage);
+  value.printAsOperand(os, asmState);
+  return os.str();
+}
+
+std::string stringifyResultSsaNames(Operation* op, AsmState* asmState) {
+  if (!asmState || op->getNumResults() == 0)
+    return "";
+
+  std::string storage;
+  llvm::raw_string_ostream os(storage);
+  llvm::interleave(
+    op->getResults(), [&](Value result) { os << stringifyValueAsOperand(result, *asmState); }, [&]() { os << ";"; });
+  return os.str();
+}
+
+std::optional<uint64_t> getTypeSizeBytes(Type type) {
+  if (auto shapedType = dyn_cast<ShapedType>(type)) {
+    if (!shapedType.hasStaticShape() || !hasByteSizedElementType(shapedType.getElementType()))
+      return std::nullopt;
+    return static_cast<uint64_t>(getShapedTypeSizeInBytes(shapedType));
+  }
+
+  if (isa<IndexType>(type))
+    return static_cast<uint64_t>(getElementTypeSizeInBytes(type));
+  if (auto intType = dyn_cast<IntegerType>(type)) {
+    if (intType.getWidth() <= 0 || intType.getWidth() % 8 != 0)
+      return std::nullopt;
+    return static_cast<uint64_t>(getElementTypeSizeInBytes(type));
+  }
+  if (auto floatType = dyn_cast<FloatType>(type)) {
+    if (floatType.getWidth() <= 0 || floatType.getWidth() % 8 != 0)
+      return std::nullopt;
+    return static_cast<uint64_t>(getElementTypeSizeInBytes(type));
+  }
+  return std::nullopt;
+}
+
+std::string getScalarId(bool isScheduled, size_t opId) { return (isScheduled ? "sc:" : "gc:") + std::to_string(opId); }
+
+std::string getBatchLaneId(bool isScheduled, size_t opId, uint32_t lane) {
+  return (isScheduled ? "scb:" : "gcb:") + std::to_string(opId) + ":" + std::to_string(lane);
+}
+
+template <typename ComputeOpTy, typename BatchOpTy>
+bool isTopLevelRelevantCompute(Operation& op) {
+  return isa<ComputeOpTy, BatchOpTy>(&op);
+}
+
+template <typename ComputeOpTy, typename BatchOpTy>
+FailureOr<TopLevelOpInfo> buildTopLevelOpInfo(Operation& op, bool isScheduled, size_t opId) {
+  TopLevelOpInfo info;
+  info.op = &op;
+  info.opId = opId;
+  info.isScheduled = isScheduled;
+
+  if constexpr (std::is_same_v<ComputeOpTy, SpatScheduledCompute>) {
+    if (auto compute = dyn_cast<ComputeOpTy>(&op)) {
+      auto coreId = getOptionalScheduledCoreId(compute, "spatial dataflow export core id");
+      if (failed(coreId))
+        return failure();
+      if (*coreId)
+        info.scalarCore = **coreId;
+    }
+  }
+
+  return info;
+}
+
+template <typename BatchOpTy>
+FailureOr<SmallVector<int32_t, 8>> getBatchLaneCoreIds(BatchOpTy batch) {
+  if constexpr (std::is_same_v<BatchOpTy, SpatScheduledComputeBatch>) {
+    auto coreIds = getOptionalScheduledBatchCoreIds(batch, "spatial dataflow export core ids");
+    if (failed(coreIds))
+      return failure();
+    if (!*coreIds)
+      return SmallVector<int32_t, 8> {};
+    return SmallVector<int32_t, 8>((**coreIds).begin(), (**coreIds).end());
+  }
+  return SmallVector<int32_t, 8> {};
+}
+
+std::string getExpandedNodeId(const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
+                              Operation* op,
+                              uint32_t lane) {
+  auto it = expandedNodes.find({op, lane});
+  if (it == expandedNodes.end())
+    return "";
+  return it->second.id;
+}
+
+void addScalarNodeRow(std::fstream& nodesFile,
+                      DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
+                      const TopLevelOpInfo& info,
+                      AsmState* asmState = nullptr) {
+  std::string id = getScalarId(info.isScheduled, info.opId);
+  SmallVector<std::string, 5> row {id, std::to_string(info.opId), "", maybeNumber<int32_t>(info.scalarCore)};
+  if (asmState)
+    row.push_back(stringifyResultSsaNames(info.op, asmState));
+  writeCsvRow(nodesFile, row);
+  expandedNodes[{info.op, 0}] = {id, info.scalarCore, std::nullopt};
+}
+
+template <typename BatchOpTy>
+void addBatchNodeRows(std::fstream& nodesFile,
+                      DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
+                      const TopLevelOpInfo& info,
+                      BatchOpTy batch,
+                      ArrayRef<std::optional<int32_t>> laneCoreIds,
+                      AsmState* asmState = nullptr) {
+  for (uint32_t lane = 0; lane < static_cast<uint32_t>(batch.getLaneCount()); ++lane) {
+    std::string id = getBatchLaneId(info.isScheduled, info.opId, lane);
+    SmallVector<std::string, 5> row {
+      id, std::to_string(info.opId), std::to_string(lane), maybeNumber<int32_t>(laneCoreIds[lane])};
+    if (asmState)
+      row.push_back(stringifyResultSsaNames(info.op, asmState));
+    writeCsvRow(nodesFile, row);
+    expandedNodes[{info.op, lane}] = {id, laneCoreIds[lane], lane};
+  }
+}
+
+std::optional<int64_t> evaluateIndexLike(Value value, Value laneArg, uint32_t lane);
+
+std::optional<int64_t> evaluateIndexLike(Value value, Value laneArg, uint32_t lane) {
+  if (value == laneArg)
+    return static_cast<int64_t>(lane);
+
+  if (std::optional<int64_t> constant = matchConstantIndexValue(value))
+    return *constant;
+
+  if (auto constant = value.getDefiningOp<arith::ConstantOp>()) {
+    if (auto intAttr = dyn_cast<IntegerAttr>(constant.getValue()))
+      return intAttr.getInt();
+  }
+
+  if (auto extract = value.getDefiningOp<tensor::ExtractOp>()) {
+    auto constant = extract.getTensor().getDefiningOp<arith::ConstantOp>();
+    auto elements = constant ? dyn_cast<ElementsAttr>(constant.getValue()) : nullptr;
+    auto shapedType = elements ? dyn_cast<ShapedType>(elements.getType()) : nullptr;
+    if (!elements || !shapedType || shapedType.getRank() != 1 || extract.getIndices().size() != 1)
+      return std::nullopt;
+
+    std::optional<int64_t> index = evaluateIndexLike(extract.getIndices().front(), laneArg, lane);
+    if (!index || *index < 0 || *index >= static_cast<int64_t>(elements.getNumElements()))
+      return std::nullopt;
+
+    if (auto denseInts = dyn_cast<DenseIntElementsAttr>(elements))
+      return (*(denseInts.value_begin<APInt>() + *index)).getSExtValue();
+    return std::nullopt;
+  }
+
+  if (auto affineApply = value.getDefiningOp<affine::AffineApplyOp>())
+    if (FailureOr<int64_t> folded = evaluateAffineApply(affineApply,
+                                                        [&](Value operand) -> FailureOr<int64_t> {
+                                                          if (std::optional<int64_t> resolved =
+                                                                evaluateIndexLike(operand, laneArg, lane))
+                                                            return *resolved;
+                                                          return failure();
+                                                        });
+        succeeded(folded)) {
+      return *folded;
+    }
+
+  return std::nullopt;
+}
+
+SmallVector<int64_t, 8> collectPossibleIntValues(Value value, Value laneArg, uint32_t lane) {
+  if (std::optional<int64_t> exact = evaluateIndexLike(value, laneArg, lane))
+    return {*exact};
+
+  auto extract = value.getDefiningOp<tensor::ExtractOp>();
+  auto constant = extract ? extract.getTensor().getDefiningOp<arith::ConstantOp>() : nullptr;
+  auto elements = constant ? dyn_cast<ElementsAttr>(constant.getValue()) : nullptr;
+  if (!elements)
+    return {};
+
+  SmallVector<int64_t, 8> values;
+  if (auto denseInts = dyn_cast<DenseIntElementsAttr>(elements)) {
+    values.reserve(elements.getNumElements());
+    for (APInt element : denseInts.getValues<APInt>())
+      if (!llvm::is_contained(values, element.getSExtValue()))
+        values.push_back(element.getSExtValue());
+  }
+  return values;
+}
+
+template <typename BatchOpTy>
+std::optional<Value> getBatchLaneInput(BatchOpTy batch, uint32_t lane, unsigned inputIndex) {
+  if (batch.getNumResults() != 0)
+    return batch.getInputs()[inputIndex];
+
+  size_t laneCount = static_cast<size_t>(batch.getLaneCount());
+  if (laneCount == 0 || batch.getInputs().size() % laneCount != 0)
+    return std::nullopt;
+
+  size_t inputsPerLane = batch.getInputs().size() / laneCount;
+  size_t flatIndex = static_cast<size_t>(lane) * inputsPerLane + inputIndex;
+  if (flatIndex >= batch.getInputs().size())
+    return std::nullopt;
+  return batch.getInputs()[flatIndex];
+}
+
+template <typename BatchOpTy>
+unsigned getBatchLaneInputCount(BatchOpTy batch) {
+  if (batch.getNumResults() != 0)
+    return batch.getInputs().size();
+
+  size_t laneCount = static_cast<size_t>(batch.getLaneCount());
+  if (laneCount == 0 || batch.getInputs().size() % laneCount != 0)
+    return 0;
+  return static_cast<unsigned>(batch.getInputs().size() / laneCount);
+}
+
+template <typename ComputeOpTy, typename BatchOpTy>
+std::optional<ResolvedProducer> resolveProducerForValue(Value value, std::optional<uint32_t> consumerLane) {
+  Operation* op = value.getDefiningOp();
+  if (!op)
+    return std::nullopt;
+
+  while (auto extract = dyn_cast<tensor::ExtractSliceOp>(op)) {
+    Value source = extract.getSource();
+    Operation* sourceOp = source.getDefiningOp();
+    auto sourceBatch = dyn_cast_or_null<BatchOpTy>(sourceOp);
+    if (sourceBatch && sourceBatch.getNumResults() != 0) {
+      auto staticOffsets = extract.getStaticOffsets();
+      if (!staticOffsets.empty() && staticOffsets.front() != ShapedType::kDynamic) {
+        uint32_t lane = static_cast<uint32_t>(staticOffsets.front());
+        return ResolvedProducer {sourceOp, 0, LogicalNodeSelector::Lane, lane, lane, 1};
+      }
+      if (consumerLane)
+        return ResolvedProducer {sourceOp, 0, LogicalNodeSelector::Lane, *consumerLane, *consumerLane, 1};
+      return ResolvedProducer {
+        sourceOp, 0, LogicalNodeSelector::RangeRepresentative, 0, 0, static_cast<uint32_t>(sourceBatch.getLaneCount())};
+    }
+    value = source;
+    op = sourceOp;
+    if (!op)
+      return std::nullopt;
+  }
+
+  if (auto compute = dyn_cast<ComputeOpTy>(op))
+    return ResolvedProducer {compute.getOperation(),
+                             static_cast<size_t>(cast<OpResult>(value).getResultNumber()),
+                             LogicalNodeSelector::Scalar,
+                             0,
+                             0,
+                             1};
+
+  if (auto batch = dyn_cast<BatchOpTy>(op)) {
+    if (batch.getNumResults() != 0) {
+      if (consumerLane)
+        return ResolvedProducer {op, 0, LogicalNodeSelector::Lane, *consumerLane, *consumerLane, 1};
+      return ResolvedProducer {
+        op, 0, LogicalNodeSelector::RangeRepresentative, 0, 0, static_cast<uint32_t>(batch.getLaneCount())};
+    }
+
+    uint32_t lane = static_cast<uint32_t>(cast<OpResult>(value).getResultNumber());
+    return ResolvedProducer {op, static_cast<size_t>(lane), LogicalNodeSelector::Lane, lane, lane, 1};
+  }
+
+  return std::nullopt;
+}
+
+SmallVector<EdgeSource, 8>
+resolveProducerSourcesForCsv(const ResolvedProducer& producer,
+                             const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes) {
+  SmallVector<EdgeSource, 8> sources;
+
+  if (producer.selector == LogicalNodeSelector::Scalar) {
+    std::string id = getExpandedNodeId(expandedNodes, producer.op, 0);
+    if (!id.empty())
+      sources.push_back({id, std::nullopt});
+    return sources;
+  }
+
+  if (producer.selector == LogicalNodeSelector::Lane) {
+    std::string id = getExpandedNodeId(expandedNodes, producer.op, producer.lane);
+    if (!id.empty())
+      sources.push_back({id, producer.lane});
+    return sources;
+  }
+
+  for (uint32_t lane = producer.laneStart; lane < producer.laneStart + producer.laneCount; ++lane) {
+    std::string id = getExpandedNodeId(expandedNodes, producer.op, lane);
+    if (!id.empty())
+      sources.push_back({id, lane});
+  }
+  return sources;
+}
+
+FailureOr<SmallVector<int64_t>> getIntegerValues(Operation* op, StringRef name) {
+  Attribute attr = op->getAttr(name);
+  if (auto array = dyn_cast_or_null<DenseI64ArrayAttr>(attr))
+    return SmallVector<int64_t>(array.asArrayRef());
+  if (auto elements = dyn_cast_or_null<DenseIntElementsAttr>(attr))
+    return SmallVector<int64_t>(elements.getValues<int64_t>());
+  return op->emitOpError() << "expected " << name << " integer array for Spatial dataflow report";
+}
+
+FailureOr<ScheduledNodeByGraphLane>
+buildScheduledNodesByGraphLane(const DenseMap<Operation*, TopLevelOpInfo>& topLevelInfo,
+                               const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
+                               const DenseMap<int64_t, Operation*>& graphOpsById) {
+  ScheduledNodeByGraphLane nodesByGraphLane;
+  for (const auto& entry : topLevelInfo) {
+    Operation* scheduledOp = entry.first;
+    auto sourceIds = getIntegerValues(scheduledOp, "scheduled.step_source_ids");
+    auto sourceStarts = getIntegerValues(scheduledOp, "scheduled.source_lane_starts");
+    auto sourceCounts = getIntegerValues(scheduledOp, "scheduled.source_lane_counts");
+    if (failed(sourceIds) || failed(sourceStarts) || failed(sourceCounts))
+      return failure();
+
+    uint32_t scheduledLaneCount = 1;
+    if (auto batch = dyn_cast<SpatScheduledComputeBatch>(scheduledOp))
+      scheduledLaneCount = static_cast<uint32_t>(batch.getLaneCount());
+    size_t expectedEntries = sourceIds->size() * scheduledLaneCount;
+    if (sourceStarts->size() != expectedEntries || sourceCounts->size() != expectedEntries)
+      return scheduledOp->emitOpError("inconsistent scheduling provenance arrays for Spatial dataflow report");
+
+    for (auto [step, graphId] : llvm::enumerate(*sourceIds)) {
+      auto graphIt = graphOpsById.find(graphId);
+      if (graphIt == graphOpsById.end())
+        return scheduledOp->emitOpError() << "references unknown scheduled graph id " << graphId;
+      bool graphIsBatch = isa<SpatGraphComputeBatch>(graphIt->second);
+      for (uint32_t scheduledLane = 0; scheduledLane < scheduledLaneCount; ++scheduledLane) {
+        auto nodeIt = expandedNodes.find({scheduledOp, scheduledLane});
+        if (nodeIt == expandedNodes.end())
+          continue;
+        size_t index = step * scheduledLaneCount + scheduledLane;
+        int64_t start = graphIsBatch ? (*sourceStarts)[index] : 0;
+        int64_t count = graphIsBatch ? (*sourceCounts)[index] : 1;
+        if (start < 0 || count < 0)
+          return scheduledOp->emitOpError("negative scheduling provenance range for Spatial dataflow report");
+        for (int64_t lane = start; lane < start + count; ++lane)
+          nodesByGraphLane[{graphId, static_cast<uint32_t>(lane)}] = nodeIt->second;
+      }
+    }
+  }
+  return nodesByGraphLane;
+}
+
+SmallVector<ExpandedNodeInfo, 8> resolveScheduledProducerNodes(const ResolvedProducer& producer,
+                                                               const ScheduledNodeByGraphLane& nodesByGraphLane) {
+  SmallVector<ExpandedNodeInfo, 8> nodes;
+  auto graphId = producer.op->getAttrOfType<IntegerAttr>("scheduled.graph_id");
+  if (!graphId)
+    return nodes;
+
+  uint32_t laneStart = producer.selector == LogicalNodeSelector::Scalar ? 0 : producer.laneStart;
+  uint32_t laneCount = producer.selector == LogicalNodeSelector::RangeRepresentative ? producer.laneCount : 1;
+  for (uint32_t lane = laneStart; lane < laneStart + laneCount; ++lane)
+    if (auto it = nodesByGraphLane.find({graphId.getInt(), lane}); it != nodesByGraphLane.end())
+      nodes.push_back(it->second);
+  return nodes;
+}
+
+LogicalResult
+emitScheduledPlanningEdges(std::fstream& edgesFile,
+                           func::FuncOp func,
+                           const DenseMap<Operation*, TopLevelOpInfo>& topLevelInfo,
+                           const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
+                           StringRef stage) {
+  DenseMap<int64_t, Operation*> graphOpsById;
+  for (Operation& op : func.getBody().front())
+    if (auto graphId = op.getAttrOfType<IntegerAttr>("scheduled.graph_id"))
+      graphOpsById[graphId.getInt()] = &op;
+
+  auto nodesByGraphLane = buildScheduledNodesByGraphLane(topLevelInfo, expandedNodes, graphOpsById);
+  if (failed(nodesByGraphLane))
+    return failure();
+
+  auto emitMappedEdge =
+    [&](const ResolvedProducer& producer, int64_t targetGraphId, uint32_t targetGraphLane, Type type) {
+      auto targetIt = nodesByGraphLane->find({targetGraphId, targetGraphLane});
+      if (targetIt == nodesByGraphLane->end())
+        return;
+      for (const ExpandedNodeInfo& source : resolveScheduledProducerNodes(producer, *nodesByGraphLane)) {
+        if (source.id == targetIt->second.id)
+          continue;
+        emitEdgeRow(edgesFile,
+                    source.id,
+                    targetIt->second.id,
+                    getTypeSizeBytes(type),
+                    type,
+                    stage,
+                    source.lane,
+                    targetIt->second.lane,
+                    std::nullopt);
+      }
+    };
+
+  for (Operation& op : func.getBody().front()) {
+    auto graphId = op.getAttrOfType<IntegerAttr>("scheduled.graph_id");
+    if (!graphId)
+      continue;
+    if (auto compute = dyn_cast<SpatGraphCompute>(&op)) {
+      for (Value input : compute.getInputs())
+        if (auto producer = resolveProducerForValue<SpatGraphCompute, SpatGraphComputeBatch>(input, std::nullopt))
+          emitMappedEdge(*producer, graphId.getInt(), 0, input.getType());
+      continue;
+    }
+    auto batch = dyn_cast<SpatGraphComputeBatch>(&op);
+    if (!batch)
+      continue;
+    unsigned inputCount = getBatchLaneInputCount(batch);
+    for (uint32_t lane = 0; lane < static_cast<uint32_t>(batch.getLaneCount()); ++lane)
+      for (unsigned inputIndex = 0; inputIndex < inputCount; ++inputIndex)
+        if (std::optional<Value> input = getBatchLaneInput(batch, lane, inputIndex))
+          if (auto producer = resolveProducerForValue<SpatGraphCompute, SpatGraphComputeBatch>(*input, lane))
+            emitMappedEdge(*producer, graphId.getInt(), lane, input->getType());
+  }
+  return success();
+}
+
+void emitEdgeRow(std::fstream& edgesFile,
+                 StringRef sourceId,
+                 StringRef targetId,
+                 std::optional<uint64_t> byteSize,
+                 Type propagatedType,
+                 StringRef stage,
+                 std::optional<uint32_t> sourceLane,
+                 std::optional<uint32_t> targetLane,
+                 std::optional<int64_t> channelId) {
+  writeCsvRow(edgesFile,
+              {sourceId.str(),
+               targetId.str(),
+               maybeNumber<uint64_t>(byteSize),
+               stringifyType(propagatedType),
+               stage.str(),
+               maybeNumber<uint32_t>(sourceLane),
+               maybeNumber<uint32_t>(targetLane),
+               maybeNumber<int64_t>(channelId)});
+}
+
+template <typename ComputeOpTy, typename BatchOpTy>
+LogicalResult emitDataEdges(std::fstream& edgesFile,
+                            const DenseMap<Operation*, TopLevelOpInfo>& topLevelInfo,
+                            const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
+                            StringRef stage) {
+  for (const auto& entry : topLevelInfo) {
+    Operation* op = entry.first;
+    const TopLevelOpInfo& info = entry.second;
+
+    if (auto compute = dyn_cast<ComputeOpTy>(op)) {
+      for (Value input : compute.getInputs()) {
+        if (isa_and_nonnull<SpatChannelReceiveOp>(input.getDefiningOp()))
+          continue;
+
+        auto producer = resolveProducerForValue<ComputeOpTy, BatchOpTy>(input, std::nullopt);
+        if (!producer)
+          continue;
+
+        SmallVector<EdgeSource, 8> sources = resolveProducerSourcesForCsv(*producer, expandedNodes);
+        std::optional<uint64_t> byteSize = getTypeSizeBytes(input.getType());
+        std::string targetId = getScalarId(info.isScheduled, info.opId);
+        for (const EdgeSource& source : sources)
+          emitEdgeRow(edgesFile,
+                      source.id,
+                      targetId,
+                      byteSize,
+                      input.getType(),
+                      stage,
+                      source.sourceLane,
+                      std::nullopt,
+                      std::nullopt);
+      }
+      continue;
+    }
+
+    auto batch = dyn_cast<BatchOpTy>(op);
+    if (!batch)
+      continue;
+
+    unsigned inputCount = getBatchLaneInputCount(batch);
+    for (uint32_t lane = 0; lane < static_cast<uint32_t>(batch.getLaneCount()); ++lane) {
+      std::string targetId = getBatchLaneId(info.isScheduled, info.opId, lane);
+      for (unsigned inputIndex = 0; inputIndex < inputCount; ++inputIndex) {
+        std::optional<Value> input = getBatchLaneInput(batch, lane, inputIndex);
+        if (!input || isa_and_nonnull<SpatChannelReceiveOp>((*input).getDefiningOp()))
+          continue;
+
+        auto producer = resolveProducerForValue<ComputeOpTy, BatchOpTy>(*input, lane);
+        if (!producer)
+          continue;
+
+        SmallVector<EdgeSource, 8> sources = resolveProducerSourcesForCsv(*producer, expandedNodes);
+        std::optional<uint64_t> byteSize = getTypeSizeBytes((*input).getType());
+        for (const EdgeSource& source : sources)
+          emitEdgeRow(
+            edgesFile, source.id, targetId, byteSize, (*input).getType(), stage, source.sourceLane, lane, std::nullopt);
+      }
+    }
+  }
+
+  return success();
+}
+
+template <typename BatchOpTy>
+void collectChannelSends(DenseMap<int64_t, SmallVector<ChannelSendRecord, 4>>& sendsByChannelId,
+                         const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
+                         BatchOpTy batch) {
+  std::optional<BlockArgument> laneArg = batch.getLaneArgument();
+  if (!laneArg)
+    return;
+
+  for (uint32_t lane = 0; lane < static_cast<uint32_t>(batch.getLaneCount()); ++lane) {
+    std::string sourceId = getExpandedNodeId(expandedNodes, batch.getOperation(), lane);
+    if (sourceId.empty())
+      continue;
+    batch.getBody().walk([&](SpatChannelSendOp send) {
+      std::optional<int64_t> channelId = evaluateIndexLike(send.getChannelId(), *laneArg, lane);
+      if (!channelId)
+        return;
+      sendsByChannelId[*channelId].push_back({sourceId, lane});
+    });
+  }
+}
+
+void collectChannelSends(DenseMap<int64_t, SmallVector<ChannelSendRecord, 4>>& sendsByChannelId,
+                         const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
+                         SpatScheduledCompute compute) {
+  std::string sourceId = getExpandedNodeId(expandedNodes, compute.getOperation(), 0);
+  if (sourceId.empty())
+    return;
+  compute.getBody().walk([&](SpatChannelSendOp send) {
+    std::optional<int64_t> channelId = evaluateIndexLike(send.getChannelId(), Value(), 0);
+    if (!channelId)
+      return;
+    sendsByChannelId[*channelId].push_back({sourceId, std::nullopt});
+  });
+}
+
+DenseMap<int32_t, SmallVector<ChannelSendRecord, 4>>
+buildNodesByCore(const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes) {
+  DenseMap<int32_t, SmallVector<ChannelSendRecord, 4>> nodesByCore;
+  for (const auto& entry : expandedNodes) {
+    const ExpandedNodeInfo& node = entry.second;
+    if (!node.core)
+      continue;
+    nodesByCore[*node.core].push_back({node.id, node.lane});
+  }
+  return nodesByCore;
+}
+
+template <typename ComputeOpTy, typename BatchOpTy, typename ResolveChannelSourcesFn>
+LogicalResult emitExplicitChannelEdges(std::fstream& edgesFile,
+                                       const DenseMap<Operation*, TopLevelOpInfo>& topLevelInfo,
+                                       ResolveChannelSourcesFn&& resolveChannelSources,
+                                       StringRef stage) {
+  for (const auto& entry : topLevelInfo) {
+    Operation* op = entry.first;
+    const TopLevelOpInfo& info = entry.second;
+
+    if (auto compute = dyn_cast<ComputeOpTy>(op)) {
+      compute.getBody().walk([&](SpatChannelReceiveOp receive) {
+        SmallVector<ChannelSendRecord, 4> sources = resolveChannelSources(receive, 0);
+        if (sources.empty())
+          return;
+        std::optional<int64_t> channelId = evaluateIndexLike(receive.getChannelId(), Value(), 0);
+        std::string targetId = getScalarId(info.isScheduled, info.opId);
+        std::optional<uint64_t> byteSize = getTypeSizeBytes(receive.getType());
+        for (const ChannelSendRecord& source : sources)
+          emitEdgeRow(edgesFile,
+                      source.sourceId,
+                      targetId,
+                      byteSize,
+                      receive.getType(),
+                      stage,
+                      source.sourceLane,
+                      std::nullopt,
+                      channelId);
+      });
+      continue;
+    }
+
+    auto batch = dyn_cast<BatchOpTy>(op);
+    if (!batch)
+      continue;
+    auto laneArg = batch.getLaneArgument();
+    if (!laneArg)
+      continue;
+    for (uint32_t lane = 0; lane < static_cast<uint32_t>(batch.getLaneCount()); ++lane) {
+      std::string targetId = getBatchLaneId(info.isScheduled, info.opId, lane);
+      batch.getBody().walk([&](SpatChannelReceiveOp receive) {
+        SmallVector<ChannelSendRecord, 4> sources = resolveChannelSources(receive, lane);
+        if (sources.empty())
+          return;
+        std::optional<int64_t> channelId = evaluateIndexLike(receive.getChannelId(), *laneArg, lane);
+        std::optional<uint64_t> byteSize = getTypeSizeBytes(receive.getType());
+        for (const ChannelSendRecord& source : sources)
+          emitEdgeRow(edgesFile,
+                      source.sourceId,
+                      targetId,
+                      byteSize,
+                      receive.getType(),
+                      stage,
+                      source.sourceLane,
+                      lane,
+                      channelId);
+      });
+    }
+  }
+
+  return success();
+}
+
+LogicalResult exportGraph(func::FuncOp func, StringRef reportName) {
+  std::fstream nodesFile = openDialectDumpFileWithExtension((reportName + ".nodes").str(), "/reports", "csv");
+  std::fstream edgesFile = openDialectDumpFileWithExtension((reportName + ".edges").str(), "/reports", "csv");
+  if (!nodesFile.is_open() || !edgesFile.is_open())
+    return success();
+
+  writeCsvRow(nodesFile, {"Id", "op_id", "lane", "core", "ssa_name"});
+  writeCsvRow(edgesFile, {"Source", "Target", "Weight", "Type", "stage", "source_lane", "target_lane", "channel_id"});
+
+  Operation* asmRoot = func.getOperation();
+  if (auto moduleOp = func->getParentOfType<ModuleOp>())
+    asmRoot = moduleOp.getOperation();
+  OpPrintingFlags flags;
+  flags.elideLargeElementsAttrs().enableDebugInfo(true, false);
+  AsmState asmState(asmRoot, flags);
+
+  DenseMap<Operation*, TopLevelOpInfo> topLevelInfo;
+  DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo> expandedNodes;
+
+  size_t opId = 0;
+  for (Operation& op : func.getBody().front()) {
+    if (!isTopLevelRelevantCompute<SpatGraphCompute, SpatGraphComputeBatch>(op))
+      continue;
+    FailureOr<TopLevelOpInfo> info = buildTopLevelOpInfo<SpatGraphCompute, SpatGraphComputeBatch>(op, false, opId++);
+    if (failed(info))
+      return failure();
+    topLevelInfo[&op] = *info;
+
+    if (auto compute = dyn_cast<SpatGraphCompute>(&op)) {
+      addScalarNodeRow(nodesFile, expandedNodes, *info, &asmState);
+      continue;
+    }
+
+    auto batch = cast<SpatGraphComputeBatch>(&op);
+    SmallVector<std::optional<int32_t>, 8> laneCoreIds(batch.getLaneCount());
+    addBatchNodeRows(nodesFile, expandedNodes, *info, batch, laneCoreIds, &asmState);
+  }
+
+  return emitDataEdges<SpatGraphCompute, SpatGraphComputeBatch>(edgesFile, topLevelInfo, expandedNodes, "spatial1");
+}
+
+LogicalResult exportScheduled(func::FuncOp func, StringRef reportName, StringRef stage) {
+  std::fstream nodesFile = openDialectDumpFileWithExtension((reportName + ".nodes").str(), "/reports", "csv");
+  std::fstream edgesFile = openDialectDumpFileWithExtension((reportName + ".edges").str(), "/reports", "csv");
+  if (!nodesFile.is_open() || !edgesFile.is_open())
+    return success();
+
+  writeCsvRow(nodesFile, {"Id", "op_id", "lane", "core", "ssa_name"});
+  writeCsvRow(edgesFile, {"Source", "Target", "Weight", "Type", "stage", "source_lane", "target_lane", "channel_id"});
+
+  Operation* asmRoot = func.getOperation();
+  if (auto moduleOp = func->getParentOfType<ModuleOp>())
+    asmRoot = moduleOp.getOperation();
+  OpPrintingFlags flags;
+  flags.elideLargeElementsAttrs().enableDebugInfo(true, false);
+  AsmState asmState(asmRoot, flags);
+
+  DenseMap<Operation*, TopLevelOpInfo> topLevelInfo;
+  DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo> expandedNodes;
+
+  size_t opId = 0;
+  for (Operation& op : func.getBody().front()) {
+    if (!isTopLevelRelevantCompute<SpatScheduledCompute, SpatScheduledComputeBatch>(op))
+      continue;
+    FailureOr<TopLevelOpInfo> info =
+      buildTopLevelOpInfo<SpatScheduledCompute, SpatScheduledComputeBatch>(op, true, opId++);
+    if (failed(info))
+      return failure();
+    topLevelInfo[&op] = *info;
+
+    if (isa<SpatScheduledCompute>(&op)) {
+      addScalarNodeRow(nodesFile, expandedNodes, *info, &asmState);
+      continue;
+    }
+
+    auto batch = cast<SpatScheduledComputeBatch>(&op);
+    auto coreIds = getBatchLaneCoreIds(batch);
+    if (failed(coreIds))
+      return failure();
+    SmallVector<std::optional<int32_t>, 8> laneCoreIds(batch.getLaneCount());
+    for (uint32_t lane = 0; lane < static_cast<uint32_t>(batch.getLaneCount()); ++lane)
+      if (lane < coreIds->size())
+        laneCoreIds[lane] = (*coreIds)[lane];
+    addBatchNodeRows(nodesFile, expandedNodes, *info, batch, laneCoreIds, &asmState);
+  }
+
+  if (stage == "spatial2")
+    return emitScheduledPlanningEdges(edgesFile, func, topLevelInfo, expandedNodes, stage);
+  if (failed(
+        emitDataEdges<SpatScheduledCompute, SpatScheduledComputeBatch>(edgesFile, topLevelInfo, expandedNodes, stage)))
+    return failure();
+
+  DenseMap<int64_t, SmallVector<ChannelSendRecord, 4>> sendsByChannelId;
+  for (const auto& entry : topLevelInfo) {
+    Operation* op = entry.first;
+    if (auto compute = dyn_cast<SpatScheduledCompute>(op))
+      collectChannelSends(sendsByChannelId, expandedNodes, compute);
+    else if (auto batch = dyn_cast<SpatScheduledComputeBatch>(op))
+      collectChannelSends(sendsByChannelId, expandedNodes, batch);
+  }
+
+  DenseMap<int32_t, SmallVector<ChannelSendRecord, 4>> nodesByCore = buildNodesByCore(expandedNodes);
+  auto resolveChannelSources = [&](SpatChannelReceiveOp receive, uint32_t lane) {
+    SmallVector<ChannelSendRecord, 4> sources;
+
+    Value laneArg;
+    if (auto owner = receive->getParentOfType<SpatScheduledComputeBatch>())
+      if (auto maybeLaneArg = owner.getLaneArgument())
+        laneArg = *maybeLaneArg;
+
+    if (std::optional<int64_t> channelId = evaluateIndexLike(receive.getChannelId(), laneArg, lane)) {
+      if (auto it = sendsByChannelId.find(*channelId); it != sendsByChannelId.end())
+        return it->second;
+    }
+
+    for (int64_t sourceCore : collectPossibleIntValues(receive.getSourceCoreId(), laneArg, lane)) {
+      auto it = nodesByCore.find(static_cast<int32_t>(sourceCore));
+      if (it == nodesByCore.end())
+        continue;
+      llvm::append_range(sources, it->second);
+    }
+    return sources;
+  };
+
+  return emitExplicitChannelEdges<SpatScheduledCompute, SpatScheduledComputeBatch>(
+    edgesFile, topLevelInfo, resolveChannelSources, stage);
+}
+
+} // namespace
+
+SpatialDataflowExportStage getSpatialDataflowExportStage() {
+  switch (pimExportSpatialDataflow.getValue()) {
+  case SpatialDataflowExportNone:     return SpatialDataflowExportStage::None;
+  case SpatialDataflowExportSpatial1: return SpatialDataflowExportStage::Spatial1;
+  case SpatialDataflowExportSpatial2: return SpatialDataflowExportStage::Spatial2;
+  case SpatialDataflowExportSpatial3: return SpatialDataflowExportStage::Spatial3;
+  case SpatialDataflowExportAll:      return SpatialDataflowExportStage::All;
+  }
+  llvm_unreachable("unknown spatial dataflow export mode");
+}
+
+bool shouldExportSpatialDataflowStage(SpatialDataflowExportStage mode, SpatialDataflowExportStage stage) {
+  switch (mode) {
+  case SpatialDataflowExportStage::None:     return false;
+  case SpatialDataflowExportStage::Spatial1: return stage == SpatialDataflowExportStage::Spatial1;
+  case SpatialDataflowExportStage::Spatial2: return stage == SpatialDataflowExportStage::Spatial2;
+  case SpatialDataflowExportStage::Spatial3: return stage == SpatialDataflowExportStage::Spatial3;
+  case SpatialDataflowExportStage::All:      return stage != SpatialDataflowExportStage::None;
+  }
+  return false;
+}
+
+LogicalResult exportSpatialDataflowCsvGraph(func::FuncOp func, StringRef reportName) {
+  return exportGraph(func, reportName);
+}
+
+LogicalResult exportSpatialDataflowCsvScheduled(func::FuncOp func, StringRef reportName, StringRef stage) {
+  return exportScheduled(func, reportName, stage);
+}
+
+} // namespace spatial
+} // namespace onnx_mlir

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

@@ -0,0 +1,28 @@
+#pragma once
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Support/LogicalResult.h"
+
+#include "llvm/ADT/StringRef.h"
+
+namespace onnx_mlir {
+namespace spatial {
+
+enum class SpatialDataflowExportStage {
+  None,
+  Spatial1,
+  Spatial2,
+  Spatial3,
+  All,
+};
+
+SpatialDataflowExportStage getSpatialDataflowExportStage();
+
+mlir::LogicalResult exportSpatialDataflowCsvGraph(mlir::func::FuncOp func, llvm::StringRef reportName);
+mlir::LogicalResult
+exportSpatialDataflowCsvScheduled(mlir::func::FuncOp func, llvm::StringRef reportName, llvm::StringRef stage);
+
+bool shouldExportSpatialDataflowStage(SpatialDataflowExportStage mode, SpatialDataflowExportStage stage);
+
+} // namespace spatial
+} // namespace onnx_mlir

src/PIM/Pass/PIMPasses.h

@@ -11,8 +11,6 @@ std::unique_ptr<mlir::Pass> createONNXToSpatialPass();
 std::unique_ptr<mlir::Pass> createSpatialLayoutPlanningPass();
 std::unique_ptr<mlir::Pass> createLowerSpatialPlansPass();
 
-std::unique_ptr<mlir::Pass> createSpatialToGraphvizPass();
-
 std::unique_ptr<mlir::Pass> createSpatialToPimPass();
 
 std::unique_ptr<mlir::Pass> createPimBufferizationPass();