blazingly faster

.gitignore

@@ -11,4 +11,4 @@ build_*
 compile.sh
 pimcomp_utils/*
 
-**/__*
+**/__pycache__/

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

@@ -99,12 +99,13 @@ Pass these to `onnx-mlir` when compiling for PIM:
 - `--core-count=<N>` - required positive core count for PIM compilation.
 - `--crossbar-size=<N>` - crossbar width/height. Default in code is `2`.
 - `--crossbar-count=<N>` - crossbars per core. Default in code is `256`.
-- `--pim-merge-scheduler=peft` - merge scheduler. `peft` is the only accepted
-  value in the current code.
 - `--pim-only-codegen` - assume input is already bufferized PIM IR and only run
   the codegen tail.
 - `--pim-emit-json` - also emit `core_*.json` instruction files alongside
   `core_*.pim`.
+- `--pim-export-spatial-dataflow=<none|spatial1|spatial2|spatial3|spatial4|all>` - control Spatial
+  dataflow CSV reports for the graph, trivially merged 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 +168,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_trivial_merged.mlir`,
+`spatial3_scheduled_no_comm.mlir`, `spatial4_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.

requirements.txt

@@ -0,0 +1,4 @@
+colorama>=0.4.6,<1
+numpy>=1.26.4,<3
+onnx>=1.17,<2
+-e ./tools/raptor_graph_explorer[dev]

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,9 @@ add_pim_library(OMPimCommon
   IR/IndexingUtils.cpp
   IR/LoopUtils.cpp
   IR/ShapeUtils.cpp
+  IR/ShapingUtils.cpp
+  IR/StaticIntSequence.cpp
+  IR/StaticIntGrid.cpp
   IR/SubviewUtils.cpp
   IR/TensorSliceUtils.cpp
   IR/WeightUtils.cpp

src/PIM/Common/IR/AddressAnalysis.cpp

@@ -1,3 +1,4 @@
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
@@ -7,6 +8,7 @@
 #include <limits>
 
 #include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 
@@ -391,6 +393,11 @@ llvm::FailureOr<int64_t> resolveIndexValueImpl(mlir::Value value, const StaticVa
   if (!definingOp)
     return mlir::failure();
 
+  if (auto affineApplyOp = mlir::dyn_cast<mlir::affine::AffineApplyOp>(definingOp))
+    return evaluateAffineApply(affineApplyOp, [&](mlir::Value operand) {
+      return resolveIndexValueImpl(operand, knowledge);
+    });
+
   if (auto indexCastOp = mlir::dyn_cast<mlir::arith::IndexCastOp>(definingOp))
     return resolveIndexValueImpl(indexCastOp.getIn(), knowledge);
 

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

@@ -1,3 +1,4 @@
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
@@ -10,7 +11,8 @@
 namespace onnx_mlir {
 
 bool isCoreStaticAddressOp(mlir::Operation* op) {
-  if (mlir::isa<mlir::arith::ConstantOp,
+  if (mlir::isa<mlir::affine::AffineApplyOp,
+                mlir::arith::ConstantOp,
                 mlir::arith::AddIOp,
                 mlir::arith::SubIOp,
                 mlir::arith::MulIOp,
@@ -36,9 +38,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 +77,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 +121,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 +168,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/StaticIntGrid.cpp

@@ -0,0 +1,223 @@
+#include "StaticIntGrid.hpp"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+
+#include "AffineUtils.hpp"
+#include "ConstantUtils.hpp"
+#include <algorithm>
+#include <limits>
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+static std::optional<size_t> cellCount(size_t rows, size_t columns) {
+  if (!rows || !columns ||
+      rows > static_cast<size_t>(std::numeric_limits<int64_t>::max()) / columns)
+    return std::nullopt;
+  return rows * columns;
+}
+
+static bool checkedFlatIndex(
+    size_t row, size_t columns, size_t column, size_t &flat) {
+  bool overflow;
+  flat = llvm::SaturatingMultiplyAdd(row, columns, column, &overflow);
+  return !overflow &&
+         flat <= static_cast<size_t>(std::numeric_limits<int64_t>::max());
+}
+
+static bool affineValue(int64_t base, int64_t rowStep, int64_t columnStep,
+                        size_t row, size_t column, int64_t &result) {
+  if (row > static_cast<size_t>(std::numeric_limits<int64_t>::max()) ||
+      column > static_cast<size_t>(std::numeric_limits<int64_t>::max()))
+    return false;
+  int64_t rowValue, columnValue;
+  return !llvm::MulOverflow(rowStep, static_cast<int64_t>(row), rowValue)
+      && !llvm::MulOverflow(columnStep, static_cast<int64_t>(column),
+                           columnValue)
+      && !llvm::AddOverflow(base, rowValue, result)
+      && !llvm::AddOverflow(result, columnValue, result);
+}
+
+} // namespace
+
+FailureOr<StaticIntGrid> StaticIntGrid::fromSequences(
+    ArrayRef<StaticIntSequence> input, bool columnsInput,
+    int64_t sparseBase) {
+  if (input.empty() || !input.front().size())
+    return failure();
+  size_t rowCount = columnsInput ? input.front().size() : input.size();
+  size_t columnCount = columnsInput ? input.size() : input.front().size();
+  auto cells = cellCount(rowCount, columnCount);
+  if (!cells ||
+      llvm::any_of(input, [&](const StaticIntSequence &sequence) {
+        return sequence.size() != input.front().size();
+      }))
+    return failure();
+  StaticIntGrid result(rowCount, columnCount, input.front().valueAt(0));
+  if (llvm::all_equal(input)) {
+    if (input.front().getKind() == StaticIntSequenceKind::Uniform)
+      return result;
+    result.kind = columnsInput ? Kind::ActionOnly : Kind::LaneOnly;
+    result.values = input.front();
+    return result;
+  }
+  SmallVector<int64_t> outerBases;
+  for (const StaticIntSequence &sequence : input)
+    outerBases.push_back(sequence.valueAt(0));
+  if (llvm::all_of(input, [](const StaticIntSequence &sequence) {
+        return sequence.getKind() == StaticIntSequenceKind::Uniform;
+      })) {
+    result.kind = columnsInput ? Kind::LaneOnly : Kind::ActionOnly;
+    result.values = StaticIntSequence::fromValues(outerBases);
+    return result;
+  }
+  auto innerStep = input.front().getAffineStep();
+  StaticIntSequence bases = StaticIntSequence::fromValues(outerBases);
+  auto outerStep = bases.getAffineStep();
+  if (innerStep && outerStep &&
+      llvm::all_of(input, [&](const StaticIntSequence &sequence) {
+        return sequence.getAffineStep() == innerStep;
+      }))
+    return affine2D(result.base,
+      columnsInput ? *innerStep : *outerStep,
+      columnsInput ? *outerStep : *innerStep, rowCount, columnCount);
+  SmallVector<int64_t> values;
+  values.reserve(*cells);
+  for (size_t row = 0; row < rowCount; ++row)
+    for (size_t column = 0; column < columnCount; ++column)
+      values.push_back(columnsInput ? input[column].valueAt(row)
+                                    : input[row].valueAt(column));
+  result.values = StaticIntSequence::fromValues(values);
+  for (size_t index = 0; index < *cells; ++index)
+    if (values[index] != sparseBase)
+      result.overrideKeys.push_back(static_cast<int64_t>(index));
+  if (result.overrideKeys.size() <= *cells / 4) {
+    result.kind = Kind::SparseLaneOverrides;
+    result.base = sparseBase;
+  } else {
+    result.kind = Kind::Dense;
+    result.overrideKeys.clear();
+  }
+  return result;
+}
+
+FailureOr<StaticIntGrid> StaticIntGrid::fromRows(
+    ArrayRef<StaticIntSequence> rows) {
+  if (rows.empty() || !rows.front().size())
+    return failure();
+  return fromSequences(rows, false, rows.front().valueAt(0));
+}
+
+FailureOr<StaticIntGrid> StaticIntGrid::fromColumns(
+    size_t rowCount, ArrayRef<StaticIntSequence> columnSequences,
+    int64_t defaultValue) {
+  if (!cellCount(rowCount, columnSequences.size()))
+    return failure();
+  SmallVector<StaticIntSequence> padded;
+  padded.reserve(columnSequences.size());
+  for (const StaticIntSequence &sequence : columnSequences) {
+    if (sequence.size() > rowCount)
+      return failure();
+    if (sequence.size() == rowCount) {
+      padded.push_back(sequence);
+      continue;
+    }
+    SmallVector<int64_t> values(rowCount, defaultValue);
+    for (size_t row = 0; row < sequence.size(); ++row)
+      values[row] = sequence.valueAt(row);
+    padded.push_back(StaticIntSequence::fromValues(values));
+  }
+  return fromSequences(padded, true, defaultValue);
+}
+
+FailureOr<StaticIntGrid> StaticIntGrid::affine2D(
+    int64_t base, int64_t rowStep, int64_t columnStep,
+    size_t rows, size_t columns) {
+  int64_t last;
+  if (!cellCount(rows, columns) ||
+      !affineValue(base, rowStep, columnStep, rows - 1, columns - 1, last))
+    return failure();
+  StaticIntGrid result(rows, columns, base);
+  result.kind = rowStep || columnStep ? Kind::Affine : Kind::Uniform;
+  result.rowStep = rowStep;
+  result.columnStep = columnStep;
+  return result;
+}
+
+FailureOr<StaticIntGrid> StaticIntGrid::laneIntervals(
+    size_t columns, ArrayRef<std::pair<size_t, size_t>> intervals,
+    int64_t insideValue, int64_t outsideValue) {
+  if (!columns)
+    return failure();
+  SmallVector<int64_t> values(columns, outsideValue);
+  for (auto [begin, end] : intervals) {
+    if (begin > end || end > columns)
+      return failure();
+    std::fill(values.begin() + begin, values.begin() + end, insideValue);
+  }
+  StaticIntSequence row = StaticIntSequence::fromValues(values);
+  return fromRows(ArrayRef<StaticIntSequence>(row));
+}
+
+int64_t StaticIntGrid::valueAt(size_t row, size_t column) const {
+  assert(row < rows && column < columns);
+  if (kind == Kind::Uniform)
+    return base;
+  if (kind == Kind::ActionOnly)
+    return values->valueAt(row);
+  if (kind == Kind::LaneOnly)
+    return values->valueAt(column);
+  if (kind == Kind::Affine) {
+    int64_t result;
+    bool valid = affineValue(base, rowStep, columnStep, row, column, result);
+    assert(valid);
+    return result;
+  }
+  size_t flat;
+  bool valid = checkedFlatIndex(row, columns, column, flat);
+  assert(valid);
+  if (kind == Kind::SparseLaneOverrides) {
+    auto found = llvm::lower_bound(overrideKeys, static_cast<int64_t>(flat));
+    if (found == overrideKeys.end() || *found != static_cast<int64_t>(flat))
+      return base;
+  }
+  assert(kind == Kind::Dense || kind == Kind::SparseLaneOverrides);
+  return values->valueAt(flat);
+}
+
+Value StaticIntGrid::emitLookup(Value row, Value column,
+                                Operation *constantAnchor,
+                                ConstantPool &constants, OpBuilder &builder,
+                                Location loc) const {
+  if (kind == Kind::Uniform)
+    return constants.getIndex(base);
+  if (kind == Kind::ActionOnly || kind == Kind::LaneOnly)
+    return emitStaticIntLookup(
+      *values, kind == Kind::ActionOnly ? row : column,
+      constantAnchor, constants, builder, loc);
+  Value flat = affineMulConst(builder, loc, row, columns, constantAnchor);
+  flat = arith::AddIOp::create(builder, loc, flat, column);
+  if (kind == Kind::Affine) {
+    Value rowValue = affineMulConst(
+      builder, loc, row, rowStep, constantAnchor);
+    Value columnValue = affineMulConst(
+      builder, loc, column, columnStep, constantAnchor);
+    Value result = arith::AddIOp::create(builder, loc, rowValue, columnValue);
+    return affineAddConst(builder, loc, result, base, constantAnchor);
+  }
+  return emitStaticIntLookup(
+    *values, flat, constantAnchor, constants, builder, loc);
+}
+
+OpFoldResult StaticIntGrid::emitFoldedLookup(
+    Value row, Value column, Operation *constantAnchor,
+    ConstantPool &constants, OpBuilder &builder, Location loc) const {
+  return kind == Kind::Uniform
+       ? OpFoldResult(builder.getIndexAttr(base))
+       : OpFoldResult(emitLookup(
+           row, column, constantAnchor, constants, builder, loc));
+}
+
+} // namespace onnx_mlir

src/PIM/Common/IR/StaticIntGrid.hpp

@@ -0,0 +1,61 @@
+#pragma once
+
+#include "StaticIntSequence.hpp"
+
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+
+#include <utility>
+
+namespace onnx_mlir {
+
+class ConstantPool;
+
+class StaticIntGrid {
+public:
+  static mlir::FailureOr<StaticIntGrid> fromColumns(
+    size_t rows, llvm::ArrayRef<StaticIntSequence> columns,
+    int64_t defaultValue);
+  static mlir::FailureOr<StaticIntGrid> fromRows(
+    llvm::ArrayRef<StaticIntSequence> rows);
+  static mlir::FailureOr<StaticIntGrid> affine2D(
+    int64_t base, int64_t rowStep, int64_t columnStep,
+    size_t rows, size_t columns);
+  static mlir::FailureOr<StaticIntGrid> laneIntervals(
+    size_t columns,
+    llvm::ArrayRef<std::pair<size_t, size_t>> intervals,
+    int64_t insideValue, int64_t outsideValue);
+
+  int64_t valueAt(size_t row, size_t column) const;
+
+  mlir::Value emitLookup(mlir::Value row, mlir::Value column,
+                         mlir::Operation *constantAnchor,
+                         ConstantPool &constants, mlir::OpBuilder &builder,
+                         mlir::Location loc) const;
+  mlir::OpFoldResult emitFoldedLookup(
+    mlir::Value row, mlir::Value column, mlir::Operation *constantAnchor,
+    ConstantPool &constants, mlir::OpBuilder &builder,
+    mlir::Location loc) const;
+
+private:
+  enum class Kind { Uniform, ActionOnly, LaneOnly, Affine,
+                    SparseLaneOverrides, Dense };
+
+  StaticIntGrid(size_t rows, size_t columns, int64_t base)
+      : rows(rows), columns(columns), base(base) {}
+
+  static mlir::FailureOr<StaticIntGrid> fromSequences(
+    llvm::ArrayRef<StaticIntSequence> sequences, bool columns,
+    int64_t sparseBase);
+
+  Kind kind = Kind::Uniform;
+  size_t rows = 0;
+  size_t columns = 0;
+  int64_t base = 0;
+  int64_t rowStep = 0;
+  int64_t columnStep = 0;
+  llvm::SmallVector<int64_t> overrideKeys;
+  std::optional<StaticIntSequence> values;
+};
+
+} // 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,123 @@
+#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; }
+  std::optional<int64_t> getAffineStep() const {
+    if (kind == StaticIntSequenceKind::Uniform)
+      return 0;
+    return kind == StaticIntSequenceKind::Affine
+         ? std::optional<int64_t>(step) : std::nullopt;
+  }
+
+  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

@@ -22,7 +22,7 @@ Value extractAxisSlice(
     .getResult();
 }
 
-Value extractStaticSliceOrIdentity(RewriterBase& rewriter,
+Value extractStaticSliceOrIdentity(OpBuilder& rewriter,
                                    Location loc,
                                    Value source,
                                    RankedTensorType resultType,
@@ -52,7 +52,8 @@ Value extractStaticSliceOrIdentity(RewriterBase& rewriter,
   if (isIdentitySlice)
     return source;
 
-  return tensor::ExtractSliceOp::create(rewriter, loc, resultType, source, offsets, sizes, strides).getResult();
+  return rewriter.createOrFold<tensor::ExtractSliceOp>(
+    loc, resultType, source, offsets, sizes, strides);
 }
 
 Value insertStaticSlice(
@@ -68,4 +69,56 @@ Value insertStaticSlice(
     .getResult();
 }
 
+Value extractMixedSliceOrIdentity(OpBuilder &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,10 +8,16 @@
 
 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);
 
-mlir::Value extractStaticSliceOrIdentity(mlir::RewriterBase& rewriter,
+mlir::Value extractStaticSliceOrIdentity(mlir::OpBuilder& rewriter,
                                          mlir::Location loc,
                                          mlir::Value source,
                                          mlir::RankedTensorType resultType,
@@ -25,4 +31,22 @@ mlir::Value insertStaticSlice(mlir::PatternRewriter& rewriter,
                               mlir::Value dest,
                               llvm::ArrayRef<mlir::OpFoldResult> offsets);
 
+mlir::Value extractMixedSliceOrIdentity(mlir::OpBuilder &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

@@ -17,14 +17,16 @@ std::fstream openDialectDumpFileWithExtension(const std::string& name, llvm::Str
   return std::fstream(dialectsDir + "/" + name + "." + extension.str(), std::ios::out);
 }
 
-void dumpModule(mlir::ModuleOp moduleOp, const std::string& name) {
+void dumpModule(mlir::ModuleOp moduleOp, const std::string& name, bool assumeVerified) {
   std::fstream file = openDialectDumpFileWithExtension(name, "/dialects", "mlir");
   if (!file.is_open())
     return;
 
   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

@@ -10,7 +10,7 @@ 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");

src/PIM/Compiler/CMakeLists.txt

@@ -26,6 +26,7 @@ add_pim_library(OMPimCompilerUtils
   ${PIM_COMPILER_INCLUDE_DIRS}
 
   LINK_LIBS PUBLIC
+  MLIRAffineToStandard
   OMPimCompilerOptions
   OMPimCommon
   OMPimBufferization

src/PIM/Compiler/PimCodeGen.cpp

@@ -270,7 +270,7 @@ void PimMemory::allocateHost(ModuleOp moduleOp, func::FuncOp funcOp) {
 }
 
 void PimMemory::allocateCore(Operation* op, std::optional<unsigned> lane) {
-  auto intervals = buildLocalAllocIntervals(op, lane);
+  auto intervals = buildLocalAllocIntervals(op, lane, pimMemoryReport == PimMemoryReportFull);
   SmallVector<PlannedPhysicalSlot> plannedSlots = planPhysicalSlots(intervals);
 
   SmallVector<size_t> slotOrder(plannedSlots.size());
@@ -414,31 +414,35 @@ size_t PimAcceleratorMemory::getValueAddress(mlir::Value value,
                                              const StaticValueKnowledge& knowledge,
                                              std::optional<unsigned> lane) const {
   value = resolveCachedAlias(value, knowledge);
-  auto compiledIt = compiledAddressExprs.find(value);
-  if (compiledIt == compiledAddressExprs.end()) {
-    auto compiledExpr = compileContiguousAddressExpr(value);
-    if (failed(compiledExpr)) {
-      errs() << "Failed to compile contiguous address for value: ";
+
+  FailureOr<ResolvedContiguousAddress> resolvedAddress = resolveContiguousAddress(value, knowledge);
+  if (failed(resolvedAddress)) {
+    auto compiledIt = compiledAddressExprs.find(value);
+    if (compiledIt == compiledAddressExprs.end()) {
+      auto compiledExpr = compileContiguousAddressExpr(value);
+      if (failed(compiledExpr)) {
+        errs() << "Failed to compile contiguous address for value: ";
+        value.print(errs());
+        errs() << " : " << value.getType();
+        errs() << "\n";
+        llvm_unreachable("Failed to compile contiguous address");
+      }
+      compiledIt = compiledAddressExprs.try_emplace(value, *compiledExpr).first;
+    }
+
+    resolvedAddress = compiledIt->second.evaluate(knowledge, lane);
+    if (failed(resolvedAddress)) {
+      errs() << "Failed to evaluate contiguous address for value: ";
       value.print(errs());
       errs() << " : " << value.getType();
       errs() << "\n";
-      llvm_unreachable("Failed to compile contiguous address");
+      if (auto* definingOp = value.getDefiningOp()) {
+        errs() << "Defining op:\n";
+        definingOp->print(errs());
+        errs() << "\n";
+      }
+      llvm_unreachable("Failed to resolve contiguous address");
     }
-    compiledIt = compiledAddressExprs.try_emplace(value, *compiledExpr).first;
-  }
-
-  auto resolvedAddress = compiledIt->second.evaluate(knowledge, lane);
-  if (failed(resolvedAddress)) {
-    errs() << "Failed to evaluate contiguous address for value: ";
-    value.print(errs());
-    errs() << " : " << value.getType();
-    errs() << "\n";
-    if (auto* definingOp = value.getDefiningOp()) {
-      errs() << "Defining op:\n";
-      definingOp->print(errs());
-      errs() << "\n";
-    }
-    llvm_unreachable("Failed to resolve contiguous address");
   }
 
   MemoryValueKey key = getMemoryValueKey(resolvedAddress->base, lane);
@@ -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

@@ -15,13 +15,6 @@ llvm::cl::opt<PimEmissionTargetType> pimEmissionTarget(
   llvm::cl::init(EmitPimCodegen),
   llvm::cl::cat(OnnxMlirOptions));
 
-llvm::cl::opt<PimMergeSchedulerType>
-  pimMergeScheduler("pim-merge-scheduler",
-                    llvm::cl::desc("Scheduler used by the Spatial merge-compute-nodes pass"),
-                    llvm::cl::values(clEnumValN(MergeSchedulerPeft, "peft", "Use PEFT scheduling")),
-                    llvm::cl::init(MergeSchedulerPeft),
-                    llvm::cl::cat(OnnxMlirOptions));
-
 llvm::cl::opt<PimMemoryReportLevel> pimMemoryReport(
   "pim-memory-report",
   llvm::cl::desc("Emit a human-readable PIM memory planning report"),
@@ -59,13 +52,17 @@ llvm::cl::opt<PimConvLoweringType> pimConvLowering(
 
 llvm::cl::opt<PimSpatialDataflowExportType> pimExportSpatialDataflow(
   "pim-export-spatial-dataflow",
-  llvm::cl::desc("Emit Gephi-importable CSV dataflow reports around MergeComputeNodes materialization"),
+  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(SpatialDataflowExportPre, "pre", "Emit pre-materialization Spatial dataflow CSV reports")),
   llvm::cl::values(
-    clEnumValN(SpatialDataflowExportPost, "post", "Emit post-materialization Spatial dataflow CSV reports")),
+    clEnumValN(SpatialDataflowExportSpatial1, "spatial1", "Emit spatial1 graph dataflow CSV reports")),
   llvm::cl::values(
-    clEnumValN(SpatialDataflowExportBoth, "both", "Emit both pre- and post-materialization Spatial dataflow CSV reports")),
+    clEnumValN(SpatialDataflowExportSpatial2, "spatial2", "Emit spatial2 trivially merged graph dataflow CSV reports")),
+  llvm::cl::values(
+    clEnumValN(SpatialDataflowExportSpatial3, "spatial3", "Emit spatial3 scheduled dataflow CSV reports")),
+  llvm::cl::values(
+    clEnumValN(SpatialDataflowExportSpatial4, "spatial4", "Emit spatial4 realized dataflow CSV reports")),
+  llvm::cl::values(clEnumValN(SpatialDataflowExportAll, "all", "Emit all Spatial dataflow CSV reports")),
   llvm::cl::init(SpatialDataflowExportNone),
   llvm::cl::cat(OnnxMlirOptions));
 

src/PIM/Compiler/PimCompilerOptions.hpp

@@ -20,10 +20,6 @@ typedef enum {
   EmitPimCodegen = 3
 } PimEmissionTargetType;
 
-typedef enum {
-  MergeSchedulerPeft = 0,
-} PimMergeSchedulerType;
-
 typedef enum {
   PimMemoryReportNone = 0,
   PimMemoryReportSummary = 1,
@@ -44,14 +40,15 @@ typedef enum {
 
 typedef enum {
   SpatialDataflowExportNone = 0,
-  SpatialDataflowExportPre = 1,
-  SpatialDataflowExportPost = 2,
-  SpatialDataflowExportBoth = 3,
+  SpatialDataflowExportSpatial1 = 1,
+  SpatialDataflowExportSpatial2 = 2,
+  SpatialDataflowExportSpatial3 = 3,
+  SpatialDataflowExportSpatial4 = 4,
+  SpatialDataflowExportAll = 5,
 } 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;

src/PIM/Compiler/PimCompilerUtils.cpp

@@ -1,3 +1,4 @@
+#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Transforms/Passes.h"
 
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
@@ -31,6 +32,7 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
     pm.addPass(createONNXToSpatialPass());
     pm.addPass(createSpatialLayoutPlanningPass());
     pm.addPass(createLowerSpatialPlansPass());
+    pm.addPass(createTrivialGraphComputeMergePass());
     pm.addPass(createMergeComputeNodesPass());
     pm.addPass(createMessagePass("Onnx lowered to Spatial"));
   }
@@ -46,6 +48,7 @@ void addPassesPim(OwningOpRef<ModuleOp>& module,
   }
 
   if (pimEmissionTarget >= EmitPimCodegen) {
+    pm.addPass(mlir::createLowerAffinePass());
     pm.addPass(createPimHostConstantFoldingPass());
     pm.addPass(createMessagePass("Pim host constants folded"));
     if (!pimDisableMemoryCoalescing)

src/PIM/Compiler/PimMemoryLiveness.cpp

@@ -154,7 +154,10 @@ static OperationOrdering buildOperationOrdering(Operation* coreLikeOp) {
 }
 
 static bool isSupportedAliasOp(Operation* op) {
-  return isa<memref::SubViewOp, memref::CastOp, memref::CollapseShapeOp, memref::ExpandShapeOp>(op);
+  return isa<memref::SubViewOp,
+             memref::CastOp,
+             memref::CollapseShapeOp,
+             memref::ExpandShapeOp>(op);
 }
 
 static bool isRuntimeMemoryTouchOp(Operation* op) {
@@ -237,12 +240,19 @@ getEffectiveTouchRange(mlir::Value definingValue, Operation* user, const Operati
 }
 
 static MemoryTouchInterval
-computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ordering, uint64_t fallbackEnd) {
+computeMemoryTouchInterval(memref::AllocOp allocOp,
+                           const OperationOrdering& ordering,
+                           uint64_t fallbackEnd,
+                           bool includeAliasDescriptions) {
   MemoryTouchInterval interval;
   interval.start = ordering.position.lookup(allocOp);
   interval.end = interval.start;
   interval.startOp = allocOp;
   interval.endOp = allocOp;
+  auto recordAlias = [&](mlir::Value value) {
+    if (includeAliasDescriptions)
+      appendAliasDescription(interval.aliasesFollowed, value);
+  };
 
   SmallPtrSet<mlir::Value, 16> visitedValues;
   SmallPtrSet<Operation*, 32> visitedUsers;
@@ -262,7 +272,7 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ord
       if (isSupportedAliasOp(user)) {
         for (mlir::Value result : user->getResults()) {
           pendingValues.push_back(result);
-          appendAliasDescription(interval.aliasesFollowed, result);
+          recordAlias(result);
         }
       }
 
@@ -272,7 +282,7 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ord
           if (!tiedOperand || tiedOperand->get() != value)
             continue;
           pendingValues.push_back(result);
-          appendAliasDescription(interval.aliasesFollowed, result);
+          recordAlias(result);
         }
       }
 
@@ -282,8 +292,8 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ord
             continue;
           pendingValues.push_back(forOp.getRegionIterArgs()[index]);
           pendingValues.push_back(forOp.getResult(index));
-          appendAliasDescription(interval.aliasesFollowed, forOp.getRegionIterArgs()[index]);
-          appendAliasDescription(interval.aliasesFollowed, forOp.getResult(index));
+          recordAlias(forOp.getRegionIterArgs()[index]);
+          recordAlias(forOp.getResult(index));
           if (parentLoop && forOp != parentLoop)
             interval.escapesLoop = true;
         }
@@ -291,7 +301,25 @@ 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));
+            recordAlias(ifOp.getResult(index));
+          }
+        }
+        else if (indexSwitch) {
+          for (auto [index, operand] : llvm::enumerate(yieldOp.getOperands())) {
+            if (operand != value)
+              continue;
+            pendingValues.push_back(indexSwitch.getResult(index));
+            recordAlias(indexSwitch.getResult(index));
+          }
+        }
+        else if (!forOp) {
           addFallbackReason(interval.fallbackReason, "yield without scf.for parent");
         }
         else {
@@ -299,7 +327,7 @@ computeMemoryTouchInterval(memref::AllocOp allocOp, const OperationOrdering& ord
             if (operand != value)
               continue;
             pendingValues.push_back(forOp.getResult(index));
-            appendAliasDescription(interval.aliasesFollowed, forOp.getResult(index));
+            recordAlias(forOp.getResult(index));
             if (parentLoop && forOp == parentLoop)
               interval.escapesLoop = true;
           }
@@ -392,7 +420,8 @@ static uint64_t getSlotLogicalBytes(const PlannedPhysicalSlot& slot, ArrayRef<Lo
 } // namespace
 
 SmallVector<LocalAllocInterval, 0> onnx_mlir::buildLocalAllocIntervals(Operation* coreLikeOp,
-                                                                       std::optional<unsigned> lane) {
+                                                                       std::optional<unsigned> lane,
+                                                                       bool includeAliasDescriptions) {
   SmallVector<LocalAllocInterval, 0> intervals;
   OperationOrdering ordering = buildOperationOrdering(coreLikeOp);
   if (ordering.position.empty())
@@ -409,7 +438,8 @@ SmallVector<LocalAllocInterval, 0> onnx_mlir::buildLocalAllocIntervals(Operation
       llvm_unreachable("Failed to compute local allocation size");
     }
 
-    MemoryTouchInterval touchInterval = computeMemoryTouchInterval(allocOp, ordering, fallbackEnd);
+    MemoryTouchInterval touchInterval =
+      computeMemoryTouchInterval(allocOp, ordering, fallbackEnd, includeAliasDescriptions);
     LocalAllocInterval interval;
     interval.id = nextIntervalId++;
     interval.alloc = allocOp;

src/PIM/Compiler/PimMemoryLiveness.hpp

@@ -49,7 +49,8 @@ struct PlannedPhysicalSlot {
 };
 
 llvm::SmallVector<LocalAllocInterval, 0> buildLocalAllocIntervals(mlir::Operation* coreLikeOp,
-                                                                  std::optional<unsigned> lane);
+                                                                  std::optional<unsigned> lane,
+                                                                  bool includeAliasDescriptions = true);
 
 llvm::SmallVector<PlannedPhysicalSlot, 0> planPhysicalSlots(llvm::MutableArrayRef<LocalAllocInterval> intervals);
 

src/PIM/Conversion/CMakeLists.txt

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

src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt

@@ -31,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

@@ -13,6 +13,7 @@
 #include <utility>
 
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
@@ -60,6 +61,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 +138,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 +173,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 +208,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 +230,15 @@ auto createSpatGraphCompute(RewriterT& rewriter,
   }
 }
 
-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) {
+template <typename RewriterT>
+auto createEmptySpatGraphComputeBatch(RewriterT& rewriter,
+                                      mlir::Location loc,
+                                      mlir::TypeRange resultTypes,
+                                      int64_t laneCount,
+                                      mlir::ValueRange weights,
+                                      mlir::ValueRange inputs,
+                                      mlir::TypeRange blockArgTypes,
+                                      llvm::ArrayRef<mlir::Location> blockArgLocs) {
   if (laneCount <= 0 || laneCount > std::numeric_limits<int32_t>::max())
     return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
 
@@ -186,27 +247,36 @@ auto createSpatGraphComputeBatch(RewriterT& rewriter,
     return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
 
   auto batchOp = spatial::SpatGraphComputeBatch::create(rewriter, loc, resultTypes, *laneCountAttr, weights, inputs);
+  rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), blockArgTypes, blockArgLocs);
+  rewriter.setInsertionPointToStart(&batchOp.getBody().front());
+  return mlir::FailureOr<spatial::SpatGraphComputeBatch>(batchOp);
+}
 
-  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);
-  }
+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);
+}
 
-  auto* block =
-    rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), mlir::TypeRange(blockArgTypes), blockArgLocs);
-  rewriter.setInsertionPointToStart(block);
+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 +287,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 +347,46 @@ 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<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));
+  }
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, physicalBatch, fragmentType, {offsets, sizes, strides});
+}
+
 template <typename BodyFn>
 mlir::Value materializeOrComputeUnary(mlir::Value input,
                                       mlir::RankedTensorType resultType,

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

@@ -3,6 +3,7 @@
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
@@ -38,6 +39,30 @@ Value createPaddedInputCompute(Value input,
   if (inputType == paddedInputType)
     return input;
 
+  auto producer = inputType.getRank() == 2 && paddedInputType.getRank() == 2
+    ? input.getDefiningOp<spatial::SpatGraphComputeBatch>()
+    : spatial::SpatGraphComputeBatch();
+  auto inputFragmentType = producer
+    ? spatial::getGraphBatchFragmentType(inputType, producer.getLaneCount())
+    : FailureOr<RankedTensorType>(failure());
+  auto paddedFragmentType = producer
+    ? spatial::getGraphBatchFragmentType(paddedInputType, producer.getLaneCount())
+    : FailureOr<RankedTensorType>(failure());
+  if (producer && succeeded(inputFragmentType) && succeeded(paddedFragmentType)) {
+    auto batch = createSpatComputeBatch(rewriter, loc, TypeRange {paddedInputType}, producer.getLaneCount(), {}, input,
+      [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
+        auto fragment = extractGraphBatchPhysicalFragment(
+          rewriter, loc, args.inputs.front(), args.lane, *inputFragmentType);
+        if (failed(fragment))
+          return failure();
+        Value padded = createZeroPaddedTensor(*fragment, *paddedFragmentType, rewriter, loc);
+        publishGraphBatchPhysicalFragment(rewriter, loc, padded, args.outputs.front(), args.lane);
+        return success();
+      });
+    if (succeeded(batch))
+      return batch->getResult(0);
+  }
+
   auto computeOp = createSpatCompute<1>(rewriter, loc, TypeRange {paddedInputType}, {}, input, [&](Value computeInput) {
     Value paddedInput = createZeroPaddedTensor(computeInput, paddedInputType, rewriter, loc);
     spatial::SpatYieldOp::create(rewriter, loc, paddedInput);

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,13 +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/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.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"
 
@@ -29,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);
@@ -46,112 +41,93 @@ 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();
+  return createRowStripAssemblyBlueprint(rowStripValue.storage, rowStripValue.logicalType, rewriter, loc);
+}
 
-  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))
+static FailureOr<Value> lowerDenseBatchBiasAdd(Value input, Value bias, RankedTensorType resultType,
+                                               PatternRewriter& rewriter, Location loc) {
+  auto producer = input.getDefiningOp<spatial::SpatGraphComputeBatch>();
+  auto inputType = dyn_cast<RankedTensorType>(input.getType());
+  auto biasType = dyn_cast<RankedTensorType>(bias.getType());
+  if (!producer || !inputType || !biasType || !inputType.hasStaticShape() || !biasType.hasStaticShape()
+      || !resultType.hasStaticShape() || inputType.getDimSize(0) != producer.getLaneCount()
+      || biasType.getDimSize(0) != producer.getLaneCount() || resultType.getDimSize(0) != producer.getLaneCount())
     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)
+  auto inputFragmentType = spatial::getGraphBatchFragmentType(inputType, producer.getLaneCount());
+  auto outputFragmentType = spatial::getGraphBatchFragmentType(resultType, producer.getLaneCount());
+  if (failed(inputFragmentType) || failed(outputFragmentType) || inputFragmentType->getRank() != biasType.getRank()
+      || inputFragmentType->getDimSize(0) != 1 || inputFragmentType->getShape().drop_front() != biasType.getShape().drop_front()
+      || inputFragmentType->getRank() != outputFragmentType->getRank() + 1)
+    return failure();
+  for (auto [inputDim, outputDim] : llvm::zip(inputFragmentType->getShape().drop_front(), outputFragmentType->getShape()))
+    if (outputDim > inputDim)
       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));
+  auto batch = createSpatComputeBatch(rewriter, loc, TypeRange {resultType}, producer.getLaneCount(), {}, ValueRange {input, bias},
+    [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
+      FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, *inputFragmentType);
+      if (failed(fragment))
+        return failure();
+      MixedSliceGeometry biasSlice;
+      for (int64_t dim : inputFragmentType->getShape()) {
+        biasSlice.offsets.push_back(biasSlice.offsets.empty() ? OpFoldResult(args.lane) : rewriter.getIndexAttr(0));
+        biasSlice.sizes.push_back(rewriter.getIndexAttr(dim));
+        biasSlice.strides.push_back(rewriter.getIndexAttr(1));
+      }
+      Value biasFragment = extractMixedSliceOrIdentity(rewriter, loc, args.inputs[1], *inputFragmentType, biasSlice);
+      if (!biasFragment)
+        return failure();
+      Value added = spatial::SpatVAddOp::create(rewriter, loc, *inputFragmentType, *fragment, biasFragment);
+      MixedSliceGeometry outputSlice;
+      outputSlice.offsets.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(0));
+      outputSlice.sizes.push_back(rewriter.getIndexAttr(1));
+      outputSlice.strides.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(1));
+      for (int64_t dim : outputFragmentType->getShape())
+        outputSlice.sizes.push_back(rewriter.getIndexAttr(dim));
+      Value output = extractMixedSliceOrIdentity(rewriter, loc, added, *outputFragmentType, outputSlice);
+      if (!output)
+        return failure();
+      publishGraphBatchPhysicalFragment(rewriter, loc, output, args.outputs.front(), args.lane);
       return success();
     });
-  if (failed(batchOp))
+  if (failed(batch))
     return failure();
-  return batchOp->getResult(0);
+  return batch->getResult(0);
 }
 
 struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
@@ -194,7 +170,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)) {
@@ -266,6 +242,71 @@ 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;
+        }
+        if (planOp.getInput().getDefiningOp<spatial::SpatGraphComputeBatch>()) {
+          FailureOr<Value> lowered = lowerDenseBatchBiasAdd(planOp.getInput(), *denseBias, resultType, rewriter, planOp.getLoc());
+          if (succeeded(lowered)) {
+            rewriter.replaceOp(planOp, *lowered);
+            continue;
+          }
+        }
+        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) {
@@ -295,6 +336,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;
@@ -346,17 +389,25 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
     RewritePatternSet helperPatterns(ctx);
     populateGemmPatterns(helperPatterns, ctx);
     populateTransposePatterns(helperPatterns, ctx);
-    if (failed(applyPartialConversion(moduleOp, helperTarget, std::move(helperPatterns)))) {
-      moduleOp.emitError("failed to lower helper ONNX ops emitted by selected Spatial plan lowering");
-      signalPassFailure();
-      return;
+    FrozenRewritePatternSet 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,
@@ -372,7 +423,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;
@@ -384,22 +436,33 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
     moduleOp.walk([&](Operation* op) {
       if (isa<ONNXEntryPointOp>(op))
         return;
-      if (isa<spatial::SpatConv2DPlanOp,
-              spatial::SpatReluPlanOp,
-              spatial::SpatBlueprintOp,
-              spatial::SpatMaterializeLayoutOp>(op)
+      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::SpatMaterializeLayoutOp>(op)
           || op->getDialect()->getNamespace() == "onnx") {
         op->emitOpError("operation must not remain after LowerSpatialPlans");
         hasIllegalOps = true;
       }
     });
+
+    PassManager canonicalizationPM(ctx);
+    canonicalizationPM.addPass(createCanonicalizerPass());
+    if (failed(canonicalizationPM.run(moduleOp)))
+      moduleOp.emitWarning("failed to run LowerSpatialPlansPass canonicalization; continuing");
+
     if (hasIllegalOps) {
       signalPassFailure();
     } else {
       dumpModule(moduleOp, "spatial1_graph");
       spatial::SpatialDataflowExportStage exportMode = spatial::getSpatialDataflowExportStage();
-      if (spatial::shouldExportSpatialDataflowStage(exportMode, spatial::SpatialDataflowExportStage::Pre)
-          && failed(spatial::exportSpatialDataflowCsvPre(funcOp))) {
+      if (spatial::shouldExportSpatialDataflowStage(exportMode, spatial::SpatialDataflowExportStage::Spatial1)
+          && failed(spatial::exportSpatialDataflowCsvGraph(funcOp, "spatial1_graph"))) {
         signalPassFailure();
         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())});
@@ -174,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))) {
@@ -214,13 +211,18 @@ void ONNXToSpatialPass::runOnOperation() {
 
   populateEmptyFunction(*entryFunc);
 
-    if (failed(verifyLogicalSpatialGraphInvariants(*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/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

@@ -17,6 +17,7 @@
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
@@ -251,10 +252,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 +399,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 +441,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 +458,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 +498,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)};
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, partialPiecesArg, pieceType,
+    {pieceOffsets, pieceSizes, unitStrides});
 }
 
 static Value reducePartialPiecesForHSlice(Value partialPiecesArg,
@@ -730,7 +722,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 +794,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

@@ -9,6 +9,7 @@
 
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
@@ -299,22 +300,14 @@ static Value extractBatchedATile(Value a,
                                  RankedTensorType aTileType,
                                  PatternRewriter& rewriter,
                                  Location loc) {
-  auto aSliceType = RankedTensorType::get({1, 1, aTileType.getDimSize(1)}, aTileType.getElementType());
   Value sourceBatchIndex =
     mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
   SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), row, kOffset};
   SmallVector<OpFoldResult> sizes {
     rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(aTileType.getDimSize(1))};
-  auto slice =
-    tensor::ExtractSliceOp::create(rewriter, loc, aSliceType, a, offsets, sizes, getUnitStrides(rewriter, 3));
-  return tensor::CollapseShapeOp::create(rewriter,
-                                         loc,
-                                         aTileType,
-                                         slice,
-                                         SmallVector<ReassociationIndices> {
-                                           {0, 1},
-                                           {2}
-  });
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, a, aTileType,
+    {offsets, sizes, getUnitStrides(rewriter, 3)});
 }
 
 static Value extractBatchedBTile(Value b,
@@ -326,24 +319,15 @@ static Value extractBatchedBTile(Value b,
                                  RankedTensorType bTileType,
                                  PatternRewriter& rewriter,
                                  Location loc) {
-  auto bSliceType =
-    RankedTensorType::get({1, bTileType.getDimSize(0), bTileType.getDimSize(1)}, bTileType.getElementType());
   Value sourceBatchIndex =
     mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
   SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), kOffset, hOffset};
   SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
                                    rewriter.getIndexAttr(bTileType.getDimSize(0)),
                                    rewriter.getIndexAttr(bTileType.getDimSize(1))};
-  auto slice =
-    tensor::ExtractSliceOp::create(rewriter, loc, bSliceType, b, offsets, sizes, getUnitStrides(rewriter, 3));
-  return tensor::CollapseShapeOp::create(rewriter,
-                                         loc,
-                                         bTileType,
-                                         slice,
-                                         SmallVector<ReassociationIndices> {
-                                           {0, 1},
-                                           {2}
-  });
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, b, bTileType,
+    {offsets, sizes, getUnitStrides(rewriter, 3)});
 }
 
 static Value getBatchLaneIndex(
@@ -398,10 +382,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();
@@ -451,22 +432,14 @@ static Value extractDynamicBatchedRowVector(Value matrix,
                                             RankedTensorType vectorType,
                                             PatternRewriter& rewriter,
                                             Location loc) {
-  auto rowSliceType = RankedTensorType::get({1, 1, vectorType.getDimSize(1)}, vectorType.getElementType());
   Value sourceBatchIndex =
     mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
   SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), row, rewriter.getIndexAttr(0)};
   SmallVector<OpFoldResult> sizes {
     rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(vectorType.getDimSize(1))};
-  auto rowSlice =
-    tensor::ExtractSliceOp::create(rewriter, loc, rowSliceType, matrix, offsets, sizes, getUnitStrides(rewriter, 3));
-  return tensor::CollapseShapeOp::create(rewriter,
-                                         loc,
-                                         vectorType,
-                                         rowSlice,
-                                         SmallVector<ReassociationIndices> {
-                                           {0, 1},
-                                           {2}
-  });
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, matrix, vectorType,
+    {offsets, sizes, getUnitStrides(rewriter, 3)});
 }
 
 static FailureOr<spatial::SpatComputeBatch> createBatchedVvdmulBatch(Value a,
@@ -506,10 +479,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();
@@ -525,7 +495,6 @@ static FailureOr<Value> createBatchedDynamicOutputCompute(Value scalarPieces,
   const int64_t numOutRows = outType.getDimSize(1);
   const int64_t numOutCols = outType.getDimSize(2);
   auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
-  auto outputScalarType = RankedTensorType::get({1, 1, 1}, outType.getElementType());
 
   auto computeOp = createSpatCompute<1>(
     rewriter, loc, TypeRange {outType}, {}, ValueRange {scalarPieces}, [&](Value pieces) -> LogicalResult {
@@ -548,24 +517,15 @@ 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));
-          Value expanded = tensor::ExpandShapeOp::create(rewriter,
-                                                         nestedLoc,
-                                                         outputScalarType,
-                                                         scalar,
-                                                         SmallVector<ReassociationIndices> {
-                                                           {0},
-                                                           {1, 2}
-          });
+          FailureOr<Value> scalar = extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, scalarType);
+          if (failed(scalar))
+            return failure();
           SmallVector<OpFoldResult> outputOffsets {batch, row, column};
           SmallVector<OpFoldResult> outputSizes = {
             rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
           Value next =
             tensor::InsertSliceOp::create(
-              rewriter, nestedLoc, expanded, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
+              rewriter, nestedLoc, *scalar, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
               .getResult();
           yielded.push_back(next);
           return success();
@@ -596,10 +556,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())};
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, partialPiecesArg, pieceType,
+    {offsets, sizes, getUnitStrides(rewriter, 3)});
 }
 
 static Value reduceBatchedPartialPiecesForHSlice(Value partialPiecesArg,
@@ -646,8 +607,6 @@ static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
       const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(2), crossbarSize.getValue());
       auto pieceType = RankedTensorType::get({numOutRows, static_cast<int64_t>(crossbarSize.getValue())},
                                              partialPiecesType.getElementType());
-      auto outputSliceType = RankedTensorType::get({1, numOutRows, static_cast<int64_t>(crossbarSize.getValue())},
-                                                   partialPiecesType.getElementType());
 
       Value outputInit =
         tensor::EmptyOp::create(rewriter, loc, paddedOutType.getShape(), paddedOutType.getElementType()).getResult();
@@ -677,14 +636,6 @@ static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
               Value outputAcc = hIterArgs.front();
               Value reduced = reduceBatchedPartialPiecesForHSlice(
                 partialPiecesArg, batch, hSlice, pieceType, numKSlices, numOutHSlices, numOutRows, rewriter, hLoc);
-              Value expandedReduced = tensor::ExpandShapeOp::create(rewriter,
-                                                                    hLoc,
-                                                                    outputSliceType,
-                                                                    reduced,
-                                                                    SmallVector<ReassociationIndices> {
-                                                                      {0, 1},
-                                                                      {2}
-              });
               Value hOffset = affineMulConst(
                 rewriter, hLoc, hSlice, crossbarSize.getValue(), rewriter.getInsertionBlock()->getParentOp());
               SmallVector<OpFoldResult> outputOffsets {batch, rewriter.getIndexAttr(0), hOffset};
@@ -693,7 +644,7 @@ static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
                                                      rewriter.getIndexAttr(crossbarSize.getValue())};
               Value next =
                 tensor::InsertSliceOp::create(
-                  rewriter, hLoc, expandedReduced, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
+                  rewriter, hLoc, reduced, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
                   .getResult();
               hYielded.push_back(next);
               return success();
@@ -917,9 +868,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 +970,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 +1012,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,
-                                              Location loc) {
-  auto batchType = cast<RankedTensorType>(batchValue.getType());
-  if (batchType == keepdimsType)
-    return batchValue;
+static FailureOr<Value> buildReduceMeanKeepdimsBlueprint(
+    Value batchValue, RankedTensorType keepdimsType,
+    ArrayRef<bool> reducedAxes, ConversionPatternRewriter& rewriter,
+    Location loc) {
+  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);
-
-  SmallVector<ReassociationIndices> expandFlatToCompact(1);
-  for (int64_t axis = 0; axis < compactKeptType.getRank(); ++axis)
-    expandFlatToCompact.front().push_back(axis);
-
-  SmallVector<ReassociationIndices> expandCompactToKeepdims;
-  ReassociationIndices pendingLeadingReducedAxes;
+  int64_t laneCount = 1;
+  SmallVector<int64_t> keptAxes;
+  SmallVector<int64_t> keptAxisStrides;
   for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
-    if (isReduced) {
-      if (expandCompactToKeepdims.empty())
-        pendingLeadingReducedAxes.push_back(axis);
-      else
-        expandCompactToKeepdims.back().push_back(axis);
-      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);
+    if (dim <= 0 || (isReduced && dim != 1))
+      return failure();
+    if (!isReduced)
+      keptAxes.push_back(axis);
   }
-  if (!pendingLeadingReducedAxes.empty())
-    expandCompactToKeepdims.back().append(pendingLeadingReducedAxes.begin(), pendingLeadingReducedAxes.end());
+  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();
 
-  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> 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) {
+        fragmentOffsets.push_back(0);
+        continue;
+      }
+      int64_t dim = keepdimsType.getDimSize(axis);
+      fragmentOffsets.push_back(
+        (lane / keptAxisStrides[keptAxisIndex]) % dim);
+      ++keptAxisIndex;
+    }
+  }
+  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/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())
-    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;
-  }
-
+  auto transposedAttr = transposeDenseElementsAttr(denseAttr, permutation);
+  if (failed(transposedAttr) || transposedAttr->getType() != resultType)
+    return failure();
   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/Conversion/SpatialToPim/Patterns/ChannelLowering.cpp

@@ -1,7 +1,8 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
-
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/SpatialToPim/Patterns.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
@@ -18,6 +19,30 @@ static void copyRaptorDebugAttrs(Operation* source, Operation* target) {
   }
 }
 
+static Value createDestinationByteOffset(PatternRewriter& rewriter,
+                                         tensor::InsertSliceOp insert) {
+  auto destinationType = cast<RankedTensorType>(insert.getDestType());
+  SmallVector<int64_t> strides = computeRowMajorStrides(destinationType.getShape());
+  int64_t elementBytes = getElementTypeSizeInBytes(destinationType.getElementType());
+  Value total = arith::ConstantIndexOp::create(rewriter, insert.getLoc(), 0);
+  for (auto [dimension, offset] : llvm::enumerate(insert.getMixedOffsets())) {
+    int64_t scale = strides[dimension] * elementBytes;
+    Value component;
+    if (auto attribute = dyn_cast<Attribute>(offset)) {
+      component = arith::ConstantIndexOp::create(
+        rewriter, insert.getLoc(), cast<IntegerAttr>(attribute).getInt() * scale);
+    } else {
+      component = cast<Value>(offset);
+      if (scale != 1)
+        component = arith::MulIOp::create(
+          rewriter, insert.getLoc(), component,
+          arith::ConstantIndexOp::create(rewriter, insert.getLoc(), scale));
+    }
+    total = arith::AddIOp::create(rewriter, insert.getLoc(), total, component);
+  }
+  return total;
+}
+
 struct ChannelSendLowering : OpRewritePattern<spatial::SpatChannelSendOp> {
   using OpRewritePattern::OpRewritePattern;
 
@@ -40,7 +65,21 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
       rewriter.eraseOp(op);
       return success();
     }
-    auto outputType = cast<ShapedType>(op.getResult().getType());
+    auto outputType = cast<RankedTensorType>(op.getResult().getType());
+    tensor::InsertSliceOp destinationInsert;
+    if (op->hasOneUse()) {
+      auto insert = dyn_cast<tensor::InsertSliceOp>(*op->getUsers().begin());
+      auto destinationType = insert
+        ? dyn_cast<RankedTensorType>(insert.getDestType()) : RankedTensorType();
+      if (insert && insert.getSource() == op.getOutput()
+          && insert.getSourceType() == outputType
+          && insert->getBlock() == op->getBlock() && destinationType
+          && destinationType.hasStaticShape()
+          && isContiguousSubviewWithDynamicOffsets(
+            destinationType.getShape(), insert.getMixedOffsets(),
+            insert.getStaticSizes(), insert.getStaticStrides()))
+        destinationInsert = insert;
+    }
     Value outputBuffer =
       tensor::EmptyOp::create(rewriter, op.getLoc(), outputType.getShape(), outputType.getElementType()).getResult();
     auto sizeAttr = getTensorSizeInBytesAttr(rewriter, op.getOperation(), op.getResult());
@@ -50,7 +89,19 @@ struct ChannelReceiveLowering : OpRewritePattern<spatial::SpatChannelReceiveOp>
       rewriter, op.getLoc(), op.getResult().getType(), outputBuffer, *sizeAttr, op.getSourceCoreId());
     copyRaptorDebugAttrs(op.getOperation(), receive.getOperation());
     Value received = receive.getOutput();
-    rewriter.replaceOp(op, received);
+    if (!destinationInsert) {
+      rewriter.replaceOp(op, received);
+      return success();
+    }
+
+    rewriter.setInsertionPoint(destinationInsert);
+    Value targetOffset = createDestinationByteOffset(rewriter, destinationInsert);
+    Value zero = arith::ConstantIndexOp::create(rewriter, op.getLoc(), 0);
+    auto copy = pim::PimMemCopyOp::create(
+      rewriter, op.getLoc(), destinationInsert.getDestType(), targetOffset, zero,
+      destinationInsert.getDest(), received, *sizeAttr);
+    rewriter.replaceOp(destinationInsert, copy.getOutput());
+    rewriter.eraseOp(op);
     return success();
   }
 };

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

@@ -1,4 +1,3 @@
-#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
@@ -175,11 +174,11 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
 
   RewritePatternSet coreBodyPatterns(ctx);
   populateCoreBodyPatterns(coreBodyPatterns);
-  populateAffineToStdConversionPatterns(coreBodyPatterns);
   FrozenRewritePatternSet frozenCoreBodyPatterns(std::move(coreBodyPatterns));
 
   ConversionTarget coreBodyTarget(*ctx);
-  coreBodyTarget.addLegalDialect<PimDialect,
+  coreBodyTarget.addLegalDialect<affine::AffineDialect,
+                                 PimDialect,
                                  tensor::TensorDialect,
                                  arith::ArithDialect,
                                  bufferization::BufferizationDialect,
@@ -226,7 +225,8 @@ void onnx_mlir::raptor::SpatialToPimPass::runOnOperation() {
   eraseUnusedTensorPackingOps(funcOp, rewriter);
 
   ConversionTarget communicationTarget(*ctx);
-  communicationTarget.addLegalDialect<PimDialect,
+  communicationTarget.addLegalDialect<affine::AffineDialect,
+                                      PimDialect,
                                       tensor::TensorDialect,
                                       arith::ArithDialect,
                                       bufferization::BufferizationDialect,

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

@@ -3,8 +3,6 @@
 
 #include "src/Accelerators/PIM/Common/IR/AddressAnalysis.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
-#include "src/Accelerators/PIM/Dialect/Pim/PimOps.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/BufferizationUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Pim/Transforms/Bufferization/Common.hpp"
 
@@ -25,25 +23,8 @@ FailureOr<Value> materializeContiguousInputMemRef(Value memrefValue,
   auto shapedType = cast<ShapedType>(memrefValue.getType());
   auto contiguousType = MemRefType::get(shapedType.getShape(), shapedType.getElementType());
   Value contiguousBuffer = memref::AllocOp::create(rewriter, loc, contiguousType);
-  auto sizeInBytes =
-    getCheckedShapedTypeSizeInBytes(shapedType, contiguousBuffer.getDefiningOp(), "contiguous copy byte size");
-  if (failed(sizeInBytes))
-    return failure();
-  Value zeroOffset = getOrCreateIndexConstant(rewriter, contiguousBuffer.getDefiningOp(), 0);
-  auto sizeAttr =
-    getCheckedI32Attr(rewriter, contiguousBuffer.getDefiningOp(), *sizeInBytes, "contiguous copy byte size");
-  if (failed(sizeAttr))
-    return failure();
-
-  if (isHostBackedPimAddress(memrefValue, knowledge)) {
-    return PimMemCopyHostToDevOp::create(
-             rewriter, loc, contiguousType, zeroOffset, zeroOffset, contiguousBuffer, memrefValue, *sizeAttr)
-      .getOutput();
-  }
-
-  return PimMemCopyOp::create(
-           rewriter, loc, contiguousType, zeroOffset, zeroOffset, contiguousBuffer, memrefValue, *sizeAttr)
-    .getOutput();
+  memref::CopyOp::create(rewriter, loc, memrefValue, contiguousBuffer);
+  return contiguousBuffer;
 }
 
 Value allocateContiguousResultMemRefLike(Value memrefValue, Location loc, RewriterBase& rewriter) {

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))
-    return false;
-
-  if (isCoreBatchInputArgument(*base))
-    return true;
-
-  return isa_and_nonnull<memref::GetGlobalOp>(base->getDefiningOp());
+  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;
+    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))
-    return false;
-
-  return isa_and_nonnull<memref::AllocOp>(base->getDefiningOp());
+  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;
+    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

@@ -1,14 +1,19 @@
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotAnalysis.h"
 #include "mlir/Dialect/Bufferization/Transforms/OneShotModuleBufferize.h"
+#include "mlir/Dialect/Bufferization/Transforms/Transforms.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/Dominance.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Rewrite/PatternApplicator.h"
 
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallSet.h"
 #include "llvm/Support/Casting.h"
 
 #include "Common/PimCommon.hpp"
@@ -116,33 +121,57 @@ lowerMemRefCopyToPimCopy(memref::CopyOp copyOp, PatternRewriter& rewriter, const
   return success();
 }
 
-static LogicalResult verifyLoweredPimCopy(pim::PimMemCopyHostToDevOp copyOp, const StaticValueKnowledge& knowledge) {
+static LogicalResult verifyLoweredPimCopy(pim::PimMemCopyHostToDevOp copyOp,
+                                          const StaticValueKnowledge& knowledge,
+                                          bool emitDiagnostic) {
   bool sourceIsHost = isHostBackedPimAddress(copyOp.getHostSource(), knowledge);
   bool targetIsHost = isHostBackedPimAddress(copyOp.getDeviceTarget(), knowledge);
   bool sourceIsDevice = isDeviceLocalPimAddress(copyOp.getHostSource(), knowledge);
   bool targetIsDevice = isDeviceLocalPimAddress(copyOp.getDeviceTarget(), knowledge);
-  if (!sourceIsHost || !targetIsDevice || targetIsHost || sourceIsDevice)
-    return copyOp.emitOpError("pim.memcp_hd requires a host-backed source and a device-local target");
+  if (!sourceIsHost || !targetIsDevice || targetIsHost || sourceIsDevice) {
+    if (emitDiagnostic)
+      copyOp.emitOpError() << "pim.memcp_hd requires a host-backed source and a device-local target: source="
+                           << copyOp.getHostSource() << " host=" << sourceIsHost << " device=" << sourceIsDevice
+                           << ", target=" << copyOp.getDeviceTarget() << " host=" << targetIsHost
+                           << " device=" << targetIsDevice;
+    return failure();
+  }
   return success();
 }
 
-static LogicalResult verifyLoweredPimCopy(pim::PimMemCopyDevToHostOp copyOp, const StaticValueKnowledge& knowledge) {
+static LogicalResult verifyLoweredPimCopy(pim::PimMemCopyDevToHostOp copyOp,
+                                          const StaticValueKnowledge& knowledge,
+                                          bool emitDiagnostic) {
   bool sourceIsHost = isHostBackedPimAddress(copyOp.getDeviceSource(), knowledge);
   bool targetIsHost = isHostBackedPimAddress(copyOp.getHostTarget(), knowledge);
   bool sourceIsDevice = isDeviceLocalPimAddress(copyOp.getDeviceSource(), knowledge);
   bool targetIsDevice = isDeviceLocalPimAddress(copyOp.getHostTarget(), knowledge);
-  if (!targetIsHost || !sourceIsDevice || sourceIsHost || targetIsDevice)
-    return copyOp.emitOpError("pim.memcp_dh requires a device-local source and a host-backed target");
+  if (!targetIsHost || !sourceIsDevice || sourceIsHost || targetIsDevice) {
+    if (emitDiagnostic)
+      copyOp.emitOpError() << "pim.memcp_dh requires a device-local source and a host-backed target: source="
+                           << copyOp.getDeviceSource() << " host=" << sourceIsHost << " device=" << sourceIsDevice
+                           << ", target=" << copyOp.getHostTarget() << " host=" << targetIsHost
+                           << " device=" << targetIsDevice;
+    return failure();
+  }
   return success();
 }
 
-static LogicalResult verifyLoweredPimCopy(pim::PimMemCopyOp copyOp, const StaticValueKnowledge& knowledge) {
+static LogicalResult verifyLoweredPimCopy(pim::PimMemCopyOp copyOp,
+                                          const StaticValueKnowledge& knowledge,
+                                          bool emitDiagnostic) {
   bool sourceIsHost = isHostBackedPimAddress(copyOp.getSource(), knowledge);
   bool targetIsHost = isHostBackedPimAddress(copyOp.getTarget(), knowledge);
   bool sourceIsDevice = isDeviceLocalPimAddress(copyOp.getSource(), knowledge);
   bool targetIsDevice = isDeviceLocalPimAddress(copyOp.getTarget(), knowledge);
-  if (!sourceIsDevice || !targetIsDevice || sourceIsHost || targetIsHost)
-    return copyOp.emitOpError("pim.memcp requires device-local source and target operands");
+  if (!sourceIsDevice || !targetIsDevice || sourceIsHost || targetIsHost) {
+    if (emitDiagnostic)
+      copyOp.emitOpError() << "pim.memcp requires device-local source and target operands: source="
+                           << copyOp.getSource() << " host=" << sourceIsHost << " device=" << sourceIsDevice
+                           << ", target=" << copyOp.getTarget() << " host=" << targetIsHost
+                           << " device=" << targetIsDevice;
+    return failure();
+  }
   return success();
 }
 
@@ -170,6 +199,109 @@ static LogicalResult applyPatternsOnce(Operation* op, PatternApplicator& applica
   return applicator.matchAndRewrite(op, rewriter);
 }
 
+static void materializeWritableConstantDestinations(func::FuncOp funcOp) {
+  SmallVector<OpOperand*> constantBackedRoots;
+  llvm::SmallPtrSet<OpOperand*, 8> seenRoots;
+  auto collect = [&](Operation* coreOp) {
+    coreOp->walk([&](tensor::InsertSliceOp insert) {
+      OpOperand* root = &insert.getDestMutable();
+      Value value = root->get();
+      llvm::SmallDenseSet<Value, 4> visited;
+      while (visited.insert(value).second) {
+        if (value.getDefiningOp<arith::ConstantOp>()) {
+          if (seenRoots.insert(root).second)
+            constantBackedRoots.push_back(root);
+          break;
+        }
+        auto argument = dyn_cast<BlockArgument>(value);
+        auto loop = argument
+          ? dyn_cast_or_null<scf::ForOp>(argument.getOwner()->getParentOp())
+          : scf::ForOp();
+        if (!loop || argument.getArgNumber() == 0)
+          break;
+        auto initArgs = loop.getInitArgsMutable();
+        root = &*(initArgs.begin() + argument.getArgNumber() - 1);
+        value = root->get();
+      }
+    });
+  };
+  funcOp.walk([&](pim::PimCoreOp coreOp) { collect(coreOp); });
+  funcOp.walk([&](pim::PimCoreBatchOp coreOp) { collect(coreOp); });
+
+  for (OpOperand* root : constantBackedRoots) {
+    OpBuilder builder(root->getOwner());
+    auto allocation = bufferization::AllocTensorOp::create(
+      builder, root->getOwner()->getLoc(), cast<RankedTensorType>(root->get().getType()),
+      ValueRange {}, root->get());
+    root->set(allocation.getResult());
+  }
+}
+
+static LogicalResult verifyConflictFreePimCoreWrites(
+  func::FuncOp funcOp, const bufferization::OneShotBufferizationOptions& options) {
+  bufferization::AnalysisState analysisState(options);
+  DominanceInfo dominance(funcOp);
+  size_t violationCount = 0;
+
+  auto verifyCore = [&](Operation* coreOp) {
+    coreOp->walk([&](Operation* writeOp) {
+      for (OpOperand& write : writeOp->getOpOperands()) {
+        if (!isa<TensorType>(write.get().getType()) || !analysisState.bufferizesToMemoryWrite(write))
+          continue;
+
+        SmallVector<Value, 8> worklist {write.get()};
+        llvm::SmallDenseSet<Value, 8> visited;
+        bool hasConflict = false;
+        Value conflictingAlias;
+        Operation* conflictingUse = nullptr;
+        while (!worklist.empty() && !hasConflict) {
+          Value alias = worklist.pop_back_val();
+          if (!visited.insert(alias).second)
+            continue;
+
+          for (OpOperand& use : alias.getUses()) {
+            bool usePrecedesWrite = false;
+            for (Operation* ancestor = use.getOwner(); ancestor; ancestor = ancestor->getParentOp())
+              if (ancestor == writeOp || dominance.properlyDominates(ancestor, writeOp)) {
+                usePrecedesWrite = true;
+                break;
+              }
+            if (usePrecedesWrite || analysisState.insideMutuallyExclusiveRegions(use.getOwner(), writeOp))
+              continue;
+            hasConflict = true;
+            conflictingAlias = alias;
+            conflictingUse = use.getOwner();
+            break;
+          }
+
+          if (alias.getDefiningOp<tensor::ExtractSliceOp>()
+              && llvm::all_of(alias.getUses(), [&](OpOperand& use) {
+                   return use.getOwner() == writeOp;
+                 }))
+            continue;
+          if (isa<OpResult>(alias))
+            for (bufferization::AliasingOpOperand tied : analysisState.getAliasingOpOperands(alias).getAliases())
+              worklist.push_back(tied.opOperand->get());
+        }
+
+        if (!hasConflict)
+          continue;
+        if (violationCount++ == 0)
+          writeOp->emitOpError() << "PIM core tensor write may modify an alias used later: operand #"
+                                 << write.getOperandNumber() << " (" << write.get() << "), alias="
+                                 << conflictingAlias << ", later use=" << *conflictingUse;
+      }
+    });
+  };
+
+  funcOp.walk([&](pim::PimCoreOp coreOp) { verifyCore(coreOp); });
+  funcOp.walk([&](pim::PimCoreBatchOp coreOp) { verifyCore(coreOp); });
+  if (violationCount != 0)
+    funcOp.emitError() << "found " << violationCount
+                       << " non-linear PIM tensor write(s); the first is reported above";
+  return success(violationCount == 0);
+}
+
 } // namespace
 
 void PimBufferizationPass::runOnOperation() {
@@ -180,9 +312,27 @@ void PimBufferizationPass::runOnOperation() {
   options.allowUnknownOps = true;
   options.bufferizeFunctionBoundaries = true;
   options.setFunctionBoundaryTypeConversion(bufferization::LayoutMapOption::IdentityLayoutMap);
-  bufferization::BufferizationState state;
 
-  if (failed(bufferization::runOneShotModuleBufferize(moduleOp, options, state))) {
+  materializeWritableConstantDestinations(funcOp);
+  if (failed(verifyConflictFreePimCoreWrites(funcOp, options))) {
+    signalPassFailure();
+    return;
+  }
+
+  auto hostOptions = options;
+  hostOptions.opFilter.denyOperation([](Operation *op) {
+    return op->getParentOfType<pim::PimCoreOp>()
+        || op->getParentOfType<pim::PimCoreBatchOp>();
+  });
+  bufferization::BufferizationState state;
+  if (failed(bufferization::insertTensorCopies(
+        moduleOp, hostOptions, state))) {
+    moduleOp.emitError("Failed to bufferize PIM and Spatial ops");
+    signalPassFailure();
+    return;
+  }
+  if (failed(bufferization::bufferizeModuleOp(
+        moduleOp, options, state))) {
     moduleOp.emitError("Failed to bufferize PIM and Spatial ops");
     signalPassFailure();
     return;
@@ -302,76 +452,87 @@ void PimBufferizationPass::annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncO
 
 LogicalResult PimBufferizationPass::verifyContiguousRuntimeOperands(ModuleOp moduleOp) const {
   bool hasFailure = false;
-  moduleOp.walk([&](Operation* op) {
-    auto verifyOperand = [&](Value operand, unsigned operandIndex) {
-      if (!isa<BaseMemRefType>(operand.getType()))
-        return;
-      if (succeeded(resolveContiguousAddress(operand)) || succeeded(compileContiguousAddressExpr(operand)))
-        return;
-      op->emitOpError() << "operand #" << operandIndex
-                        << " is not backed by contiguous addressable storage after PIM bufferization";
-      hasFailure = true;
-    };
 
-    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;
-    }
-    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;
-    }
-    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;
-    }
-    if (auto sendOp = dyn_cast<PimSendOp>(op)) {
-      verifyOperand(sendOp.getInput(), 0);
-      return;
-    }
-    if (auto receiveOp = dyn_cast<PimReceiveOp>(op)) {
-      verifyOperand(receiveOp.getOutputBuffer(), 0);
-      return;
-    }
-    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;
-    }
-    if (isa<PimTransposeOp,
-            PimVMMOp,
-            PimVVAddOp,
-            PimVVSubOp,
-            PimVVMulOp,
-            PimVVMaxOp,
-            PimVVDMulOp,
-            PimVAvgOp,
-            PimVReluOp,
-            PimVTanhOp,
-            PimVSigmOp,
-            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());
-      }
-    }
+  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, knowledge)) || succeeded(compileContiguousAddressExpr(operand)))
+            return;
+          op.emitOpError() << "operand #" << operandIndex
+                           << " is not backed by contiguous addressable storage after PIM bufferization";
+          hasFailure = true;
+        };
+
+        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 success();
+        }
+        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 success();
+        }
+        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 success();
+        }
+        if (auto sendOp = dyn_cast<PimSendOp>(&op)) {
+          verifyOperand(sendOp.getInput(), 0);
+          return success();
+        }
+        if (auto receiveOp = dyn_cast<PimReceiveOp>(&op)) {
+          verifyOperand(receiveOp.getOutputBuffer(), 0);
+          return success();
+        }
+        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 success();
+        }
+        if (isa<PimTransposeOp,
+                PimVMMOp,
+                PimVVAddOp,
+                PimVVSubOp,
+                PimVVMulOp,
+                PimVVMaxOp,
+                PimVVDMulOp,
+                PimVAvgOp,
+                PimVReluOp,
+                PimVTanhOp,
+                PimVSigmOp,
+                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) {
@@ -382,18 +543,19 @@ LogicalResult PimBufferizationPass::verifyContiguousRuntimeOperands(ModuleOp mod
 }
 
 LogicalResult PimBufferizationPass::verifyPimCopyAddressSpaces(ModuleOp moduleOp) const {
-  bool hasFailure = false;
+  size_t failureCount = 0;
   auto verifyWithKnowledge = [&](auto coreLikeOp, const StaticValueKnowledge& initialKnowledge) {
     (void) walkPimCoreBlockStructurally(
       coreLikeOp.getBody().front(), initialKnowledge, [&](Operation& op, const StaticValueKnowledge& knowledge) {
-        if (auto copyOp = dyn_cast<pim::PimMemCopyOp>(&op); copyOp && failed(verifyLoweredPimCopy(copyOp, knowledge)))
-          hasFailure = true;
+        if (auto copyOp = dyn_cast<pim::PimMemCopyOp>(&op);
+            copyOp && failed(verifyLoweredPimCopy(copyOp, knowledge, failureCount == 0)))
+          ++failureCount;
         if (auto copyOp = dyn_cast<pim::PimMemCopyHostToDevOp>(&op);
-            copyOp && failed(verifyLoweredPimCopy(copyOp, knowledge)))
-          hasFailure = true;
+            copyOp && failed(verifyLoweredPimCopy(copyOp, knowledge, failureCount == 0)))
+          ++failureCount;
         if (auto copyOp = dyn_cast<pim::PimMemCopyDevToHostOp>(&op);
-            copyOp && failed(verifyLoweredPimCopy(copyOp, knowledge)))
-          hasFailure = true;
+            copyOp && failed(verifyLoweredPimCopy(copyOp, knowledge, failureCount == 0)))
+          ++failureCount;
         return success();
       });
   };
@@ -403,7 +565,10 @@ LogicalResult PimBufferizationPass::verifyPimCopyAddressSpaces(ModuleOp moduleOp
     StaticValueKnowledge knowledge = seedCoreBatchKnowledge(coreBatchOp, 0);
     verifyWithKnowledge(coreBatchOp, knowledge);
   });
-  return success(!hasFailure);
+  if (failureCount != 0)
+    moduleOp.emitError() << "found " << failureCount
+                         << " PIM copy address-space violation(s); the first is reported above";
+  return success(failureCount == 0);
 }
 
 std::unique_ptr<Pass> createPimBufferizationPass() { return std::make_unique<PimBufferizationPass>(); }

src/PIM/Dialect/Spatial/CMakeLists.txt

@@ -7,15 +7,26 @@ 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/SpatialDataflowCsvExporter.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/ScheduledComputePlanning.cpp
+  Transforms/MergeComputeNodes/ScheduledComputeReport.cpp
+  Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp
+  Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp
   Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp
   Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp
+  Transforms/TrivialGraphComputeMergePass.cpp
 
   EXCLUDE_FROM_OM_LIBS
 

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;
@@ -48,6 +49,7 @@ class SpatComputeLikeBase<string mnemonic> : SpatOp<mnemonic,
 }
 
 def SpatGraphCompute : SpatComputeLikeBase<"graph_compute"> {
+  let hasCanonicalizer = 1;
   let extraClassDeclaration = [{
     std::optional<::mlir::BlockArgument> getWeightArgument(unsigned idx);
     std::optional<::mlir::BlockArgument> getInputArgument(unsigned idx);
@@ -76,7 +78,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 +92,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 +116,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 +165,58 @@ 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,
+    OptionalAttr<I64Attr>:$specialization_count
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let regions = (region SizedRegion<1>:$body);
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def SpatDeferredSourceSelectOp : SpatOp<"deferred_source_select", []> {
+  let summary = "Select a deferred tensor source with a statically analyzable index";
+
+  let arguments = (ins
+    Index:$selector,
+    Variadic<SpatTensor>:$sources
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  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 +288,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 +317,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,130 @@ 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 << " : ";
+  printCompressedTypeList(
+    printer, getSources().getTypes(), ListDelimiter::Paren);
+  printer << " -> ";
+  printCompressedTypeSequence(printer, getOperation()->getResultTypes());
+  printer << " ";
+  printer.printRegion(getBody(), /*printEntryBlockArgs=*/false);
+}
+
+ParseResult SpatDeferredCommunicationOp::parse(OpAsmParser& parser, OperationState& result) {
+  SmallVector<OpAsmParser::UnresolvedOperand> sources;
+  SmallVector<Type> sourceTypes, outputTypes;
+
+  if (parseCompressedOperandSequence(parser, sources) || parser.parseOptionalAttrDict(result.attributes)
+      || parser.parseColon()
+      || parseCompressedRepeatedList(
+        parser, ListDelimiter::Paren, sourceTypes,
+        [&](Type& type) { return parser.parseType(type); })
+      || parser.parseArrow()
+      || parseCompressedTypeSequence(
+        parser, outputTypes, /*allowEmpty=*/false))
+    return failure();
+
+  if (sources.size() != sourceTypes.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of sources and source types must match");
+
+  if (parser.resolveOperands(sources, sourceTypes, parser.getCurrentLocation(), result.operands))
+    return failure();
+  result.addTypes(outputTypes);
+
+  Region* body = result.addRegion();
+  SmallVector<OpAsmParser::Argument> bodyArgs;
+  for (Type type : sourceTypes) {
+    OpAsmParser::Argument argument;
+    argument.type = type;
+    bodyArgs.push_back(argument);
+  }
+  if (auto count = dyn_cast_or_null<IntegerAttr>(
+        result.attributes.get("specialization_count"));
+      count && count.getInt() > 1) {
+    OpAsmParser::Argument argument;
+    argument.type = parser.getBuilder().getIndexType();
+    bodyArgs.push_back(argument);
+  }
+  return parser.parseRegion(*body, bodyArgs);
+}
+
+void SpatDeferredSourceSelectOp::print(OpAsmPrinter& printer) {
+  printer << " " << getSelector() << " of ";
+  printCompressedValueSequence(printer, getSources());
+  printer.printOptionalAttrDict((*this)->getAttrs());
+  printer << " : " << getOutput().getType();
+}
+
+ParseResult SpatDeferredSourceSelectOp::parse(
+    OpAsmParser& parser, OperationState& result) {
+  OpAsmParser::UnresolvedOperand selector;
+  SmallVector<OpAsmParser::UnresolvedOperand> sources;
+  Type outputType;
+  if (parser.parseOperand(selector) || parser.parseKeyword("of")
+      || parseCompressedOperandSequence(parser, sources)
+      || parser.parseOptionalAttrDict(result.attributes)
+      || parser.parseColon() || parser.parseType(outputType))
+    return failure();
+  if (parser.resolveOperand(selector, parser.getBuilder().getIndexType(),
+                            result.operands))
+    return failure();
+  SmallVector<Type> sourceTypes(sources.size(), outputType);
+  if (parser.resolveOperands(sources, sourceTypes,
+                             parser.getCurrentLocation(), result.operands))
+    return failure();
+  result.addTypes(outputType);
+  return success();
+}
+
 void SpatExtractRowsOp::print(OpAsmPrinter& printer) {
   printer << " ";
   printer.printOperand(getInput());
@@ -493,6 +617,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 +642,7 @@ void SpatBlueprintOp::print(OpAsmPrinter& printer) {
                                  getIndexMapAttrName().getValue(),
                                  getModeAttrName().getValue(),
                                  getFragmentOperandIndicesAttrName().getValue(),
+                                 getFragmentSourceSlotsAttrName().getValue(),
                                  getFragmentSourceOffsetsAttrName().getValue(),
                                  getFragmentStridesAttrName().getValue(),
                                  getConflictPolicyAttrName().getValue(),
@@ -537,6 +666,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 +684,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 +717,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;
@@ -36,9 +42,208 @@ LogicalResult SpatGraphCompute::fold(FoldAdaptor adaptor, ::llvm::SmallVectorImp
   return foldComputeLike(*this, results);
 }
 
+struct RemoveUnusedGraphComputeInputsPattern : OpRewritePattern<SpatGraphCompute> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(SpatGraphCompute compute, PatternRewriter& rewriter) const override {
+    SmallVector<unsigned> unusedInputs;
+    for (unsigned index = 0; index < compute.getInputs().size(); ++index) {
+      auto argument = compute.getInputArgument(index);
+      if (argument && argument->use_empty())
+        unusedInputs.push_back(index);
+    }
+    if (unusedInputs.empty())
+      return failure();
+
+    rewriter.modifyOpInPlace(compute, [&] {
+      for (unsigned index : llvm::reverse(unusedInputs)) {
+        compute.getBody().front().eraseArgument(compute.getWeights().size() + index);
+        compute.getInputsMutable().erase(index);
+      }
+    });
+    return success();
+  }
+};
+
+void SpatGraphCompute::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) {
+  results.add<RemoveUnusedGraphComputeInputsPattern>(context);
+}
+
 LogicalResult SpatScheduledCompute::fold(FoldAdaptor adaptor, ::llvm::SmallVectorImpl<::mlir::OpFoldResult>& results) {
   return foldComputeLike(*this, results);
 }
 
+template <typename ScalarComputeOpTy>
+static ScalarComputeOpTy createEmptyScalarCompute(PatternRewriter& rewriter,
+                                                  Location loc,
+                                                  TypeRange resultTypes,
+                                                  ValueRange weights,
+                                                  ValueRange inputs) {
+  auto computeOp = ScalarComputeOpTy::create(rewriter, loc, resultTypes, weights, inputs);
+  SmallVector<Type> blockArgTypes;
+  SmallVector<Location> blockArgLocs;
+  blockArgTypes.reserve(weights.size() + inputs.size());
+  blockArgLocs.reserve(weights.size() + inputs.size());
+  for (Value weight : weights) {
+    blockArgTypes.push_back(weight.getType());
+    blockArgLocs.push_back(weight.getLoc());
+  }
+  for (Value input : inputs) {
+    blockArgTypes.push_back(input.getType());
+    blockArgLocs.push_back(input.getLoc());
+  }
+  rewriter.createBlock(&computeOp.getBody(), computeOp.getBody().end(), blockArgTypes, blockArgLocs);
+  rewriter.setInsertionPointToStart(&computeOp.getBody().front());
+  return computeOp;
+}
+
+static SmallVector<OpFoldResult> remapMixedOffsets(ArrayRef<OpFoldResult> mixedOffsets, IRMapping& mapper) {
+  SmallVector<OpFoldResult> remapped;
+  remapped.reserve(mixedOffsets.size());
+  for (OpFoldResult ofr : mixedOffsets) {
+    if (auto value = dyn_cast<Value>(ofr))
+      remapped.push_back(cast<Value>(mapper.lookupOrDefault(value)));
+    else
+      remapped.push_back(cast<Attribute>(ofr));
+  }
+  return remapped;
+}
+
+static SmallVector<Value> createEmptyResults(PatternRewriter& rewriter, Location loc, TypeRange resultTypes) {
+  SmallVector<Value> resultValues;
+  resultValues.reserve(resultTypes.size());
+  for (Type resultType : resultTypes) {
+    auto tensorType = dyn_cast<RankedTensorType>(resultType);
+    if (!tensorType || !tensorType.hasStaticShape())
+      return {};
+    resultValues.push_back(tensor::EmptyOp::create(rewriter, loc, tensorType.getShape(), tensorType.getElementType()));
+  }
+  return resultValues;
+}
+
+template <typename ScalarComputeOpTy, typename ComputeBatchOpTy>
+static void copyCanonicalizedBatchAttrs(ScalarComputeOpTy compute, ComputeBatchOpTy batch, PatternRewriter& rewriter) {
+  for (NamedAttribute attr : batch->getAttrs()) {
+    if (attr.getName() == batch.getOperandSegmentSizesAttrName() || attr.getName() == batch.getLaneCountAttrName()
+        || attr.getName() == onnx_mlir::kCoreIdsAttrName)
+      continue;
+    compute->setAttr(attr.getName(), attr.getValue());
+  }
+  if constexpr (std::is_same_v<ComputeBatchOpTy, SpatScheduledComputeBatch>) {
+    if (auto coreIds = batch->template getAttrOfType<DenseI32ArrayAttr>(onnx_mlir::kCoreIdsAttrName)) {
+      assert(coreIds.size() == 1 && "single-lane scheduled compute_batch canonicalization expects exactly one core id");
+      compute->setAttr(onnx_mlir::kCoreIdAttrName, rewriter.getI32IntegerAttr(coreIds.asArrayRef().front()));
+    }
+  }
+}
+
+template <typename ComputeBatchOpTy, typename ScalarComputeOpTy>
+struct CanonicalizeSingleLaneComputeBatchPattern : OpRewritePattern<ComputeBatchOpTy> {
+  using OpRewritePattern<ComputeBatchOpTy>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ComputeBatchOpTy compute, PatternRewriter& rewriter) const override {
+    if (compute.getLaneCount() != 1)
+      return rewriter.notifyMatchFailure(compute, "lane count is not 1");
+
+    Block& oldBlock = compute.getBody().front();
+    auto oldLaneArg = compute.getLaneArgument();
+    if (!oldLaneArg)
+      return rewriter.notifyMatchFailure(compute, "missing compute_batch lane block argument");
+
+    rewriter.setInsertionPointAfter(compute);
+    auto newCompute =
+      createEmptyScalarCompute<ScalarComputeOpTy>(rewriter, compute.getLoc(), compute.getResultTypes(), compute.getWeights(), compute.getInputs());
+    copyCanonicalizedBatchAttrs(newCompute, compute, rewriter);
+    auto* newBlock = &newCompute.getBody().front();
+    rewriter.setInsertionPointToStart(newBlock);
+
+    IRMapping mapper;
+    Value zero = 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 =
@@ -204,7 +226,7 @@ static LogicalResult verifyOnlyConstantExternalValues(Operation* ownerOp, Region
 }
 
 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,11 +241,12 @@ 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");
-  for (auto& bodyOp : block) {
-    if (auto extractSlice = dyn_cast<tensor::ExtractSliceOp>(&bodyOp))
-      if (failed(verifyStaticUnitStrideExtractSliceOp(extractSlice, *laneArg, "tensor.extract_slice")))
-        return failure();
-  }
+  if (verifyLaneSliceOffsets)
+    for (auto& bodyOp : block) {
+      if (auto extractSlice = dyn_cast<tensor::ExtractSliceOp>(&bodyOp))
+        if (failed(verifyStaticUnitStrideExtractSliceOp(extractSlice, *laneArg, "tensor.extract_slice")))
+          return failure();
+    }
   return success();
 }
 
@@ -436,6 +459,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 +547,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 +601,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 +629,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 +697,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,57 +710,65 @@ 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)
-    return compute.emitOpError("compute body must have weight and input block arguments");
+  bool isScheduled = isa<SpatScheduledCompute>(compute.getOperation());
+  if (compute.getBody().empty())
+    return compute.emitOpError("compute body must have at least one block");
+  if (isScheduled && !compute.getBody().hasOneBlock())
+    return compute.emitOpError("scheduled compute must have exactly one block");
 
-  for (auto [weightIndex, weight] : llvm::enumerate(compute.getWeights())) {
-    auto blockArg = compute.getWeightArgument(weightIndex);
-    if (!blockArg || blockArg->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())
-      return compute.emitOpError("compute input block argument types must match input operand types exactly");
-  }
+  SmallVector<Type> yieldedTypes;
+  for (Block &block : compute.getBody()) {
+    if (block.getNumArguments() != expectedArgCount)
+      return compute.emitOpError("compute body must have weight and input block arguments");
 
-  if (block.mightHaveTerminator()) {
-    auto yieldOp = dyn_cast_or_null<SpatYieldOp>(block.getTerminator());
-    if (!yieldOp)
+    for (auto [weightIndex, weight] : llvm::enumerate(compute.getWeights()))
+      if (block.getArgument(weightIndex).getType() != weight.getType())
+        return compute.emitOpError("compute weight block argument types must match weight operand types exactly");
+    for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs()))
+      if (block.getArgument(compute.getWeights().size() + inputIndex).getType() != input.getType())
+        return compute.emitOpError("compute input block argument types must match input operand types exactly");
+
+    Operation* terminator = block.getTerminator();
+    if (auto yieldOp = dyn_cast_or_null<SpatYieldOp>(terminator)) {
+      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");
+    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");
+  auto resultTypes = compute.getResultTypes();
+  if (resultTypes.size() != yieldedTypes.size())
+    return compute.emitOpError("ComputeOp must have same number of results as yielded operands");
 
-    for (auto it : llvm::reverse(llvm::zip(resultTypes, yieldTypes))) {
-      auto resultType = std::get<0>(it);
-      auto yieldType = std::get<1>(it);
+  for (auto it : llvm::reverse(llvm::zip(resultTypes, yieldedTypes))) {
+    auto resultType = std::get<0>(it);
+    auto yieldType = std::get<1>(it);
 
-      if (resultType != yieldType || failed(verifyCompatibleShape(resultType, yieldType)))
-        return compute.emitOpError("ComputeOp output must be of the same type as yieldOp operand");
+    if (resultType != yieldType || failed(verifyCompatibleShape(resultType, yieldType)))
+      return compute.emitOpError("ComputeOp output must be of the same type as yieldOp operand");
 
-      if (auto resultRankedType = dyn_cast<RankedTensorType>(resultType)) {
-        if (auto yieldRankedType = dyn_cast<RankedTensorType>(yieldType)) {
-          if (resultRankedType.getEncoding() != yieldRankedType.getEncoding())
-            return compute.emitOpError("ComputeOp output must have the same encoding as yieldOp operand");
-        }
-        else {
+    if (auto resultRankedType = dyn_cast<RankedTensorType>(resultType)) {
+      if (auto yieldRankedType = dyn_cast<RankedTensorType>(yieldType)) {
+        if (resultRankedType.getEncoding() != yieldRankedType.getEncoding())
           return compute.emitOpError("ComputeOp output has an encoding while yieldOp operand does not have one");
-        }
       }
-      else if (dyn_cast<RankedTensorType>(yieldType)) {
-        return compute.emitOpError("ComputeOp output must not have an encoding if yieldOp operand has one");
+      else {
+        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");
+    }
   }
 
-  for (unsigned inputIndex = 0; inputIndex < compute.getInputs().size(); ++inputIndex)
-    if (auto inputArg = compute.getInputArgument(inputIndex); !inputArg || inputArg->use_empty())
-      return compute.emitOpError("ComputeOp block argument is not used");
+  if (compute.getBody().hasOneBlock())
+    for (unsigned inputIndex = 0; inputIndex < compute.getInputs().size(); ++inputIndex)
+      if (auto inputArg = compute.getInputArgument(inputIndex); !inputArg || inputArg->use_empty())
+        return compute.emitOpError("ComputeOp block argument is not used");
   if (failed(verifyStaticWeights(compute, opName)))
     return failure();
   if (failed(verifyOnlyConstantExternalValues(compute.getOperation(), compute.getBody(), opName)))
@@ -705,6 +782,93 @@ 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");
+  auto specialization = (*this)->getAttrOfType<IntegerAttr>(
+    "specialization_count");
+  int64_t specializationCount = specialization ? specialization.getInt() : 1;
+  if (specializationCount <= 0)
+    return emitOpError("specialization_count must be positive");
+  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 (getBody().empty())
+    return emitOpError("spat.deferred_communication requires a body block");
+  Block &body = getBody().front();
+  unsigned expectedArguments = getSources().size()
+                             + (specializationCount > 1 ? 1 : 0);
+  if (body.getNumArguments() != expectedArguments)
+    return emitOpError("body argument count must match sources plus the grouped specialization argument");
+  for (auto [argument, source] : llvm::zip(
+         body.getArguments().take_front(getSources().size()), getSources()))
+    if (argument.getType() != source.getType())
+      return emitOpError("body source argument types must match source operand types");
+  if (specializationCount > 1
+      && !body.getArguments().back().getType().isIndex())
+    return emitOpError("grouped specialization argument must have index type");
+  auto yield = dyn_cast_or_null<SpatYieldOp>(body.getTerminator());
+  if (!yield || yield.getOutputs().size() != 1)
+    return emitOpError("body must yield exactly one fragment");
+  Type fragmentType = yield.getOutputs().front().getType();
+  Type outputType = getOutput().getType();
+  if (specializationCount == 1)
+    return fragmentType == outputType
+         ? success()
+         : emitOpError("ordinary deferred yield type must match its output type");
+  auto fragmentTensor = dyn_cast<RankedTensorType>(fragmentType);
+  auto outputTensor = dyn_cast<RankedTensorType>(outputType);
+  if (!fragmentTensor || !outputTensor || !fragmentTensor.hasStaticShape()
+      || !outputTensor.hasStaticShape())
+    return emitOpError("grouped specialization requires static ranked tensor types");
+  if (outputTensor.getRank() != fragmentTensor.getRank() + 1
+      || outputTensor.getDimSize(0) != specializationCount
+      || outputTensor.getShape().drop_front() != fragmentTensor.getShape()
+      || outputTensor.getElementType() != fragmentTensor.getElementType())
+    return emitOpError("grouped output must have shape specialization_count x fragment shape");
+  return success();
+}
+
+LogicalResult SpatDeferredSourceSelectOp::verify() {
+  if (getSources().empty())
+    return emitOpError("requires at least one source");
+  if (!getSelector().getType().isIndex())
+    return emitOpError("requires an index selector");
+  if (llvm::any_of(getSources(), [&](Value source) {
+        return source.getType() != getOutput().getType();
+      }))
+    return emitOpError("source and output types must match");
+  if (!getOperation()->getParentOfType<SpatDeferredCommunicationOp>()
+      && !getOperation()->getParentOfType<SpatScheduledCompute>()
+      && !getOperation()->getParentOfType<SpatScheduledComputeBatch>())
+    return emitOpError("must be nested in deferred or scheduled computation");
+  return success();
+}
+
 template <typename ComputeBatchOpTy>
 LogicalResult verifyComputeBatchLikeOp(ComputeBatchOpTy batch, StringRef opName) {
   int32_t count = batch.getLaneCount();
@@ -727,30 +891,33 @@ LogicalResult verifyComputeBatchLikeOp(ComputeBatchOpTy batch, StringRef opName)
         return batch.emitOpError("compute_batch coreIds values must be unique");
   }
 
-  Block& block = batch.getBody().front();
-  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 first block argument must have index type");
+  if (batch.getBody().empty())
+    return batch.emitOpError("compute_batch body must have at least one block");
 
-  for (auto [weightIndex, weight] : llvm::enumerate(batch.getWeights())) {
-    auto blockArg = batch.getWeightArgument(weightIndex);
-    if (!blockArg || blockArg->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())
-      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)
-      return batch.emitOpError("compute_batch output block argument types must match result types exactly");
+  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");
+    if (block.getNumArguments() != expectedArgCount)
+      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()))
+      if (block.getArgument(1 + weightIndex).getType() != weight.getType())
+        return batch.emitOpError("compute_batch weight block argument types must match weight operand types exactly");
+    for (auto [inputIndex, input] : llvm::enumerate(batch.getInputs()))
+      if (block.getArgument(1 + batch.getWeights().size() + inputIndex).getType() != input.getType())
+        return batch.emitOpError("compute_batch input block argument types must match input operand types exactly");
+    for (auto [resultIndex, resultType] : llvm::enumerate(batch.getResultTypes()))
+      if (block.getArgument(1 + batch.getWeights().size() + batch.getInputs().size() + resultIndex).getType()
+          != resultType)
+        return batch.emitOpError("compute_batch output block argument types must match result types exactly");
+
+    if (failed(verifyBatchBody(batch, block, verifyLaneSliceOffsets)))
+      return failure();
   }
 
   if (failed(verifyComputeResultsUses(batch.getOperation())))
@@ -759,7 +926,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,314 @@
+#include "DeferredBoundaryPlanning.hpp"
+#include "DeferredCommunicationScheduling.hpp"
+#include "DeferredTransferPlanning.hpp"
+
+namespace onnx_mlir::spatial {
+using namespace mlir;
+namespace {
+
+static BoundaryProgram &getBoundary(
+    SmallVectorImpl<BoundaryProgram> &boundaries,
+    DenseMap<BoundaryKey, unsigned> &indices, BoundaryKey key) {
+  auto [it, inserted] = indices.try_emplace(key, boundaries.size());
+  if (inserted)
+    boundaries.push_back({key, {}});
+  return boundaries[it->second];
+}
+
+static LaneSet getReceiveLanes(const ScheduledTransferSlice &slice) {
+  LaneInterval family = slice.family->targetLanes.intervals().front();
+  unsigned begin = family.begin + slice.familyOffset;
+  return LaneSet::range(begin, begin + slice.transferCount);
+}
+
+static void appendSend(BoundaryProgram &boundary,
+                       const ScheduledTransferSlice &slice) {
+  unsigned lane = slice.family->requirement->producer->scheduledLane;
+  LaneSet lanes = LaneSet::range(lane, lane + 1);
+  if (!boundary.instructions.empty())
+    if (auto *run = std::get_if<EmitSendRun>(&boundary.instructions.back());
+        run && haveSameTransferEmissionSignature(
+                 *run->slices.back().family, *slice.family)) {
+      run->slices.push_back(slice);
+      run->lanes = run->lanes.unite(lanes);
+      return;
+    }
+  boundary.instructions.push_back(EmitSendRun {{slice}, lanes});
+}
+
+static LogicalResult addCoverage(
+    RequirementFamily &requirement, const LaneSet &lanes,
+    DenseMap<RequirementFamily *, LaneSet> &coverage) {
+  LaneSet &covered = coverage[&requirement];
+  if (!covered.intersect(lanes).empty())
+    return requirement.exchange->deferred.emitOpError(
+      "deferred availability covers a target lane more than once");
+  covered = covered.unite(lanes);
+  return success();
+}
+
+static bool canGroupLocalAvailability(RequirementFamily &lhs,
+                                      RequirementFamily &rhs) {
+  if (!(lhs.coordinate == rhs.coordinate)
+      || lhs.publicationFragmentType != rhs.publicationFragmentType
+      || lhs.producer->payload != rhs.producer->payload
+      || lhs.producerProjection.has_value()
+           != rhs.producerProjection.has_value())
+    return false;
+  if (!lhs.producerProjection)
+    return true;
+  auto ranks = [](const DeferredStaticSliceGeometry &geometry) {
+    return std::tuple(geometry.offsets.size(), geometry.sizes.size(),
+                      geometry.strides.size());
+  };
+  return ranks(*lhs.producerProjection) == ranks(*rhs.producerProjection);
+}
+
+struct CollectionTarget {
+  const FragmentCollectionPlan *collection = nullptr;
+  unsigned position = 0;
+};
+
+static bool sameCollectionEmissionContract(
+    const CollectionTarget &lhs, const CollectionTarget &rhs) {
+  if (lhs.collection != rhs.collection)
+    return false;
+  if (lhs.collection->key.kind != FragmentCollectionKind::InsertAssembly)
+    return true;
+  const auto &entries =
+    lhs.collection->key.exchange->program.insertAssembly->entries;
+  const auto &left = entries[lhs.position];
+  const auto &right = entries[rhs.position];
+  return left.sourceTransform == right.sourceTransform
+      && left.sourceType == right.sourceType;
+}
+
+static std::optional<EmitLocalCollectionRun> buildLocalConcat(
+    const FragmentCollectionPlan &collection,
+    const DenseMap<RequirementFamily *, LocalAvailabilityFamily *> &locals,
+    unsigned targetLaneCount) {
+  RankedTensorType type = collection.collectionType;
+  if (collection.key.kind != FragmentCollectionKind::Leaf
+      || targetLaneCount != 1 || type.getRank() == 0
+      || collection.positionCount == 0
+      || type.getDimSize(0) != collection.positionCount)
+    return std::nullopt;
+  SmallVector<RequirementFamily *> requirements(collection.positionCount);
+  for (const auto &entry : collection.requirements) {
+    if (entry.position >= requirements.size() || requirements[entry.position])
+      return std::nullopt;
+    requirements[entry.position] = entry.family;
+  }
+  LaneSet all = LaneSet::all(targetLaneCount);
+  EmitLocalCollectionRun run {&collection, 0, {}, all, true};
+  Value payload;
+  int64_t payloadBegin = 0;
+  int64_t payloadEnd = 0;
+  for (auto [position, requirement] : llvm::enumerate(requirements)) {
+    if (!requirement)
+      return std::nullopt;
+    LocalAvailabilityFamily *local = locals.lookup(requirement);
+    if (!local || !(requirement->targetLanes == all)
+        || !(local->targetLanes == all) || !requirement->graphLanes
+        || requirement->graphLanes->size() != 1
+        || requirement->graphLanes->valueAt(0) != static_cast<int64_t>(position)
+        || !requirement->producerLocalOffsets
+        || requirement->producerLocalOffsets->size() != 1)
+      return std::nullopt;
+    ProducedValue *producer = requirement->producer;
+    int64_t payloadOffset =
+      requirement->producerLocalOffsets->valueAt(0);
+    if (static_cast<int64_t>(position) == payloadEnd) {
+      if (!producer)
+        return std::nullopt;
+      auto payloadType = dyn_cast<RankedTensorType>(producer->payload.getType());
+      if (!payloadType || payloadOffset != 0
+          || payloadType.getRank() != type.getRank()
+          || payloadType.getElementType() != type.getElementType()
+          || payloadType.getShape().drop_front() != type.getShape().drop_front())
+        return std::nullopt;
+      payloadBegin = position;
+      payloadEnd = payloadBegin + payloadType.getDimSize(0);
+      if (payloadEnd > collection.positionCount)
+        return std::nullopt;
+      payload = producer->payload;
+      run.families.push_back(local);
+    }
+    if (producer->payload != payload
+        || payloadOffset != static_cast<int64_t>(position) - payloadBegin)
+      return std::nullopt;
+  }
+  return payloadEnd == collection.positionCount
+       ? std::optional<EmitLocalCollectionRun>(std::move(run)) : std::nullopt;
+}
+
+static void appendReceive(BoundaryProgram &boundary,
+                          const ScheduledTransferSlice &slice,
+                          CollectionTarget target) {
+  RequirementFamily *requirement = slice.family->requirement;
+  LaneSet lanes = getReceiveLanes(slice);
+  if (!boundary.instructions.empty())
+    if (auto *run = std::get_if<EmitReceiveAssemblyRun>(
+          &boundary.instructions.back())) {
+      RequirementFamily *previous = run->slices[
+        run->entryOffsets[run->entryOffsets.size() - 2]].family->requirement;
+      CollectionTarget previousTarget {run->collection, run->positions.back()};
+      bool sameEntry = previous == requirement;
+      if (sameEntry
+          || (sameCollectionEmissionContract(previousTarget, target)
+              && previous->publicationFragmentType
+                   == requirement->publicationFragmentType)) {
+        run->slices.push_back(slice);
+        if (sameEntry) {
+          run->entryOffsets.back() = run->slices.size();
+          run->entryLanes.back() = run->entryLanes.back().unite(lanes);
+        } else {
+          run->entryOffsets.push_back(run->slices.size());
+          run->positions.push_back(target.position);
+          run->entryLanes.push_back(lanes);
+        }
+        run->lanes = run->lanes.unite(lanes);
+        return;
+      }
+    }
+  boundary.instructions.push_back(EmitReceiveAssemblyRun {
+    target.collection, {slice}, {0, 1}, {target.position}, {lanes}, lanes});
+}
+
+} // namespace
+
+FailureOr<DeferredBoundaryPlan> buildDeferredBoundaryPlan(
+    DeferredTransferPlan &transfers,
+    const ScheduledCommunicationPlan &schedule) {
+  DeferredBoundaryPlan result;
+  SmallVector<BoundaryProgram> boundaries;
+  DenseMap<BoundaryKey, unsigned> indices;
+  DenseMap<DeferredExchangePlan *, unsigned> resultSteps;
+  DenseMap<RequirementFamily *, LaneSet> coverage;
+
+  for (const std::unique_ptr<DeferredExchangePlan> &exchange :
+       transfers.exchanges) {
+    auto plan = buildDeferredResultPlan(*exchange);
+    if (failed(plan))
+      return exchange->deferred.emitOpError(
+        "cannot evaluate deferred result lane functions"), failure();
+    result.results.push_back(std::move(*plan));
+  }
+  DenseMap<RequirementFamily *, CollectionTarget> collections;
+  for (const DeferredResultPlan &plan : result.results)
+    for (const FragmentCollectionPlan &collection : plan.collections)
+      for (const FragmentCollectionPlan::Requirement &requirement :
+           collection.requirements)
+        if (!collections.try_emplace(
+              requirement.family,
+              CollectionTarget{&collection, requirement.position}).second)
+          return requirement.family->exchange->deferred.emitOpError(
+            "deferred requirement is owned by multiple fragment collections"),
+            failure();
+
+  for (const ScheduledTransferSlice &slice : schedule.slices) {
+    ExternalTransferFamily &family = *slice.family;
+    BoundaryProgram &source = getBoundary(
+      boundaries, indices,
+      {family.sourceScheduled, slice.sourceInsertionStep});
+    appendSend(source, slice);
+    BoundaryProgram &target = getBoundary(
+      boundaries, indices,
+      {family.targetScheduled, slice.targetInsertionStep});
+    CollectionTarget collection = collections.lookup(family.requirement);
+    if (!collection.collection)
+      return family.requirement->exchange->deferred.emitOpError(
+        "deferred requirement has no complete result-owned collection"),
+        failure();
+    appendReceive(target, slice, collection);
+    resultSteps[family.requirement->exchange] = std::max(
+      resultSteps.lookup(family.requirement->exchange),
+      slice.targetInsertionStep);
+    if (failed(addCoverage(*family.requirement, getReceiveLanes(slice),
+                           coverage)))
+      return failure();
+  }
+
+  for (const std::unique_ptr<DeferredExchangePlan> &exchange :
+       transfers.exchanges) {
+    unsigned resultStep = resultSteps.lookup(exchange.get());
+    for (LocalAvailabilityFamily &local : exchange->local)
+      resultStep = std::max(
+        resultStep, local.requirement->producer->step + 1);
+    if (exchange->requirements.empty())
+      resultStep = exchange->consumerStep;
+    if (resultStep > exchange->consumerStep)
+      return exchange->deferred.emitOpError(
+        "deferred result boundary is later than its consumer"), failure();
+
+    BoundaryProgram &boundary = getBoundary(
+      boundaries, indices, {exchange->target, resultStep});
+    DenseMap<RequirementFamily *, LocalAvailabilityFamily *> localByRequirement;
+    for (LocalAvailabilityFamily &local : exchange->local) {
+      auto [it, inserted] = localByRequirement.try_emplace(local.requirement, &local);
+      if (!inserted)
+        it->second = nullptr;
+    }
+    DenseMap<const FragmentCollectionPlan *, std::optional<EmitLocalCollectionRun>> concatRuns;
+    llvm::SmallPtrSet<const FragmentCollectionPlan *, 4> emittedConcats;
+    SmallVector<EmitLocalCollectionRun> localUpdates;
+    for (LocalAvailabilityFamily &local : exchange->local) {
+      CollectionTarget target = collections.lookup(local.requirement);
+      if (!target.collection)
+        return exchange->deferred.emitOpError(
+          "local availability has no complete result-owned collection"),
+          failure();
+      auto [concat, inserted] = concatRuns.try_emplace(target.collection);
+      if (inserted)
+        concat->second = buildLocalConcat(*target.collection,
+                                          localByRequirement,
+                                          exchange->targetLaneCount);
+      if (concat->second) {
+        if (emittedConcats.insert(target.collection).second)
+          localUpdates.push_back(std::move(*concat->second));
+        if (failed(addCoverage(*local.requirement, local.targetLanes, coverage)))
+          return failure();
+        continue;
+      }
+      auto grouped = llvm::find_if(localUpdates, [&](EmitLocalCollectionRun &update) {
+        return update.lanes.intersect(local.targetLanes).empty()
+            && update.collection == target.collection
+            && update.collectionPosition == target.position
+            && canGroupLocalAvailability(*update.families.front()->requirement,
+                                         *local.requirement);
+      });
+      if (grouped == localUpdates.end()) {
+        localUpdates.push_back(EmitLocalCollectionRun {
+          target.collection, target.position, {&local}, local.targetLanes,
+          false});
+      } else {
+        grouped->families.push_back(&local);
+        grouped->lanes = grouped->lanes.unite(local.targetLanes);
+      }
+      if (failed(addCoverage(*local.requirement, local.targetLanes, coverage)))
+        return failure();
+    }
+    for (EmitLocalCollectionRun &update : localUpdates)
+      boundary.instructions.push_back(std::move(update));
+    for (RequirementFamily &requirement : exchange->requirements)
+      if (!(coverage.lookup(&requirement) == requirement.targetLanes))
+        return exchange->deferred.emitOpError(
+          "deferred availability does not cover every target lane exactly once"),
+          failure();
+    boundary.instructions.push_back(ProduceDeferredResult {exchange.get()});
+
+  }
+
+  DenseMap<ScheduledInfo *, unsigned> scheduledOrder;
+  for (auto [index, scheduled] : llvm::enumerate(transfers.scheduled))
+    scheduledOrder[&scheduled] = index;
+  llvm::stable_sort(boundaries, [&](const BoundaryProgram &lhs,
+                                    const BoundaryProgram &rhs) {
+    return std::tie(scheduledOrder[lhs.key.first], lhs.key.second)
+         < std::tie(scheduledOrder[rhs.key.first], rhs.key.second);
+  });
+  result.boundaries = std::move(boundaries);
+  return result;
+}
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,51 @@
+#pragma once
+
+#include "DeferredCommunicationScheduling.hpp"
+#include "DeferredResultRealization.hpp"
+
+namespace onnx_mlir::spatial {
+
+struct DeferredTransferPlan;
+
+using BoundaryKey = std::pair<ScheduledInfo *, unsigned>;
+
+struct EmitSendRun {
+  llvm::SmallVector<ScheduledTransferSlice> slices;
+  LaneSet lanes;
+};
+struct EmitLocalCollectionRun {
+  const FragmentCollectionPlan* collection = nullptr;
+  unsigned collectionPosition = 0;
+  llvm::SmallVector<LocalAvailabilityFamily*> families;
+  LaneSet lanes;
+  bool concatenatePayloads = false;
+};
+struct EmitReceiveAssemblyRun {
+  const FragmentCollectionPlan* collection = nullptr;
+  llvm::SmallVector<ScheduledTransferSlice> slices;
+  llvm::SmallVector<size_t> entryOffsets;
+  llvm::SmallVector<unsigned> positions;
+  llvm::SmallVector<LaneSet> entryLanes;
+  LaneSet lanes;
+};
+struct ProduceDeferredResult {
+  DeferredExchangePlan* exchange = nullptr;
+};
+
+using BoundaryInstruction =
+  std::variant<EmitSendRun, EmitLocalCollectionRun, EmitReceiveAssemblyRun,
+               ProduceDeferredResult>;
+struct BoundaryProgram {
+  BoundaryKey key;
+  llvm::SmallVector<BoundaryInstruction, 0> instructions;
+};
+
+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

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

@@ -0,0 +1,766 @@
+#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/LoopUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/StaticIntGrid.hpp"
+#include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
+#include <array>
+namespace onnx_mlir::spatial {
+using namespace mlir;
+namespace {
+struct LogicalTransferMetadataView {
+  StaticIntSequenceChain channels;
+  StaticIntSequenceChain parents;
+  StaticIntSequenceChain parentCounts;
+  StaticIntSequenceChain sourceCores;
+  StaticIntSequenceChain targetCores;
+  StaticIntSequenceChain targetLanes;
+  StaticIntSequenceChain localOffsets;
+  SmallVector<StaticIntSequenceChain> projectionOffsets;
+  SmallVector<StaticIntSequenceChain> projectionSizes;
+  SmallVector<StaticIntSequenceChain> projectionStrides;
+
+  size_t size() const { return channels.size(); }
+};
+using MetadataMember = StaticIntSequenceChain LogicalTransferMetadataView::*;
+static constexpr std::array<MetadataMember, 3> transferMetadataMembers{
+    &LogicalTransferMetadataView::channels, &LogicalTransferMetadataView::sourceCores, &LogicalTransferMetadataView::targetCores};
+struct TransferGrids {
+  std::array<StaticIntGrid, 3> values;
+  StaticIntGrid &channels() { return values[0]; }
+  StaticIntGrid &sourceCores() { return values[1]; }
+  StaticIntGrid &targetCores() { return values[2]; }
+};
+template <typename Build> static FailureOr<TransferGrids> buildTransferGrids(Build build) {
+  auto channels = build(transferMetadataMembers[0]);
+  auto sourceCores = build(transferMetadataMembers[1]);
+  auto targetCores = build(transferMetadataMembers[2]);
+  if (failed(channels) || failed(sourceCores) || failed(targetCores))
+    return failure();
+  return TransferGrids{{std::move(*channels), std::move(*sourceCores), std::move(*targetCores)}};
+}
+using GridGeometry = DeferredGridSliceGeometry;
+using StaticGeometryMember = SmallVector<StaticIntSequence> DeferredStaticSliceGeometry::*;
+using MetadataGeometryMember = SmallVector<StaticIntSequenceChain> LogicalTransferMetadataView::*;
+static constexpr std::array<StaticGeometryMember, 3> staticGeometryMembers{&DeferredStaticSliceGeometry::offsets, &DeferredStaticSliceGeometry::sizes,
+                                                                           &DeferredStaticSliceGeometry::strides};
+static constexpr std::array<MetadataGeometryMember, 3> metadataGeometryMembers{
+    &LogicalTransferMetadataView::projectionOffsets, &LogicalTransferMetadataView::projectionSizes, &LogicalTransferMetadataView::projectionStrides};
+static MixedSliceGeometry lookupGeometry(const GridGeometry &geometry, Value row, Value lane, Operation *anchor, DeferredEmissionContext &context,
+                                         Location loc) {
+  MixedSliceGeometry result;
+  std::array<SmallVectorImpl<OpFoldResult> *, 3> targets{&result.offsets, &result.sizes, &result.strides};
+  for (auto [source, target] : llvm::zip_equal(geometry, targets))
+    for (const StaticIntGrid &grid : source)
+      target->push_back(grid.emitFoldedLookup(row, lane, anchor, context.constants, context.rewriter, loc));
+  return result;
+}
+static FailureOr<Value> emitLaneCondition(const LaneSet &lanes, Value lane, unsigned laneCount, Operation *anchor, DeferredEmissionContext &context,
+                                          Location loc) {
+  SmallVector<std::pair<size_t, size_t>> intervals;
+  for (LaneInterval interval : lanes.intervals())
+    intervals.push_back({interval.begin, interval.end});
+  auto active = StaticIntGrid::laneIntervals(laneCount, intervals, 1, 0);
+  if (failed(active))
+    return failure();
+  Value selected = active->emitLookup(context.constants.getIndex(0), lane, anchor, context.constants, context.rewriter, loc);
+  return arith::CmpIOp::create(context.rewriter, loc, arith::CmpIPredicate::ne, selected, context.constants.getIndex(0)).getResult();
+}
+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));
+  if (family.requirement->producerProjection) {
+    const DeferredStaticSliceGeometry &geometry = *family.requirement->producerProjection;
+    if (metadata.projectionOffsets.empty()) {
+      metadata.projectionOffsets.resize(geometry.offsets.size());
+      metadata.projectionSizes.resize(geometry.sizes.size());
+      metadata.projectionStrides.resize(geometry.strides.size());
+    }
+    size_t geometryIndex = targetLane - requirementLanes.begin;
+    for (auto [target, source] : llvm::zip_equal(metadata.projectionOffsets, geometry.offsets))
+      target.append(source, geometryIndex, count);
+    for (auto [target, source] : llvm::zip_equal(metadata.projectionSizes, geometry.sizes))
+      target.append(source, geometryIndex, count);
+    for (auto [target, source] : llvm::zip_equal(metadata.projectionStrides, geometry.strides))
+      target.append(source, geometryIndex, count);
+  }
+}
+static LogicalTransferMetadataView buildMetadataView(ArrayRef<ScheduledTransferSlice> slices) {
+  LogicalTransferMetadataView metadata;
+  for (const ScheduledTransferSlice &slice : slices)
+    appendMetadata(slice, metadata);
+  return metadata;
+}
+static StaticIntSequence canonicalizePadded(const StaticIntSequenceChain &chain, size_t count, int64_t defaultValue) {
+  if (chain.size() == count)
+    return chain.canonicalize();
+  StaticIntSequenceChain padded;
+  chain.forEachSegment([&](const StaticIntSequence &sequence, size_t begin, size_t length) { padded.append(sequence, begin, length); });
+  padded.append(StaticIntSequence::uniform(defaultValue, count - chain.size()));
+  return padded.canonicalize();
+}
+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 FailureOr<Value> materializeSendPayload(const RequirementFamily &requirement, Value localOffset, const MixedSliceGeometry *producerProjection,
+                                               DeferredEmissionContext &context, Location loc) {
+  Value payload = requirement.producer->payload;
+  if (producerProjection) {
+    auto fragmentType = dyn_cast<RankedTensorType>(requirement.publicationFragmentType);
+    if (!fragmentType)
+      return failure();
+    return extractMixedSliceOrIdentity(context.rewriter, loc, payload, fragmentType, *producerProjection);
+  }
+  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));
+  return extractMixedSliceOrIdentity(context.rewriter, loc, payload, fragmentType, geometry);
+}
+
+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());
+  auto buildGrid = [&](auto member, int64_t defaultValue) {
+    SmallVector<StaticIntSequence> columns;
+    for (const LogicalTransferMetadataView &metadata : metadataByLane)
+      columns.push_back(canonicalizePadded(metadata.*member, actionCount, defaultValue));
+    return StaticIntGrid::fromColumns(actionCount, columns, defaultValue);
+  };
+  auto transferGrids = buildTransferGrids([&](MetadataMember member) { return buildGrid(member, (logical.*member).valueAt(0)); });
+  FailureOr<StaticIntGrid> localOffsets = buildGrid(&LogicalTransferMetadataView::localOffsets, logical.localOffsets.valueAt(0));
+  if (failed(transferGrids) || failed(localOffsets))
+    return failure();
+  GridGeometry projectionGrids;
+  for (auto [geometryIndex, sourceMember] : llvm::enumerate(metadataGeometryMembers)) {
+    const auto &logicalValues = logical.*sourceMember;
+    for (size_t dimension = 0; dimension < logicalValues.size(); ++dimension) {
+      int64_t defaultValue = logicalValues[dimension].valueAt(0);
+      SmallVector<StaticIntSequence> columns;
+      for (const LogicalTransferMetadataView &metadata : metadataByLane) {
+        const auto &values = metadata.*sourceMember;
+        columns.push_back(dimension < values.size() ? canonicalizePadded(values[dimension], actionCount, defaultValue)
+                                                    : StaticIntSequence::uniform(defaultValue, actionCount));
+      }
+      auto grid = StaticIntGrid::fromColumns(actionCount, columns, defaultValue);
+      if (failed(grid))
+        return failure();
+      projectionGrids[geometryIndex].push_back(std::move(*grid));
+    }
+  }
+  SmallVector<int64_t> counts(laneCount);
+  for (unsigned sourceLane = 0; sourceLane < laneCount; ++sourceLane) {
+    const LogicalTransferMetadataView &source = metadataByLane[sourceLane];
+    counts[sourceLane] = source.size();
+  }
+  ExternalTransferFamily &firstFamily = *run.slices.front().family;
+  RequirementFamily &requirement = *firstFamily.requirement;
+  Operation *anchor = requirement.exchange->deferred;
+  Location loc = requirement.exchange->deferred.getLoc();
+  auto emitOne = [&](Value action, Value runtimeLane) -> LogicalResult {
+    Value localOffset = localOffsets->emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc);
+    MixedSliceGeometry projection = lookupGeometry(projectionGrids, action, runtimeLane, anchor, context, loc);
+    auto payload = materializeSendPayload(requirement, localOffset, projectionGrids[0].empty() ? nullptr : &projection, context, loc);
+    if (failed(payload))
+      return failure();
+    auto send = SpatChannelSendOp::create(
+        context.rewriter, loc, transferGrids->channels().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc),
+        transferGrids->sourceCores().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc),
+        transferGrids->targetCores().emitLookup(action, runtimeLane, anchor, context.constants, context.rewriter, loc), *payload);
+    setLogicalTransferMetadata(send, logical);
+    return success();
+  };
+  Value runtimeLane = lane ? lane : context.constants.getIndex(0);
+  if (actionCount == 1)
+    return emitOne(context.constants.getIndex(0), 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> &) { return emitOne(index, runtimeLane); });
+  return success(succeeded(loop));
+}
+
+static FailureOr<Value> emitReceiveValue(ArrayRef<ScheduledTransferSlice> slices, Value lane, unsigned laneCount, DeferredEmissionContext &context) {
+  LogicalTransferMetadataView metadata = buildMetadataView(slices);
+  RequirementFamily &requirement = *slices.front().family->requirement;
+  Operation *anchor = requirement.exchange->deferred;
+  auto buildGrid = [&](const StaticIntSequenceChain &values) {
+    int64_t defaultValue = values.valueAt(0);
+    SmallVector<int64_t> lanes(laneCount, defaultValue);
+    for (size_t index = 0; index < metadata.size(); ++index) lanes[metadata.targetLanes.valueAt(index)] = values.valueAt(index);
+    StaticIntSequence row = StaticIntSequence::fromValues(lanes);
+    return StaticIntGrid::fromRows(ArrayRef<StaticIntSequence>(row));
+  };
+  auto grids = buildTransferGrids([&](MetadataMember member) { return buildGrid(metadata.*member); });
+  if (failed(grids)) return failure();
+  Value position = lane ? lane : context.constants.getIndex(0);
+  Value row = context.constants.getIndex(0);
+  auto receive = SpatChannelReceiveOp::create(context.rewriter, anchor->getLoc(), requirement.publicationFragmentType,
+    grids->channels().emitLookup(row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
+    grids->sourceCores().emitLookup(row, position, anchor, context.constants, context.rewriter, anchor->getLoc()),
+    grids->targetCores().emitLookup(row, position, anchor, context.constants, context.rewriter, anchor->getLoc()));
+  setLogicalTransferMetadata(receive, metadata);
+  return receive.getOutput();
+}
+
+template <typename Insert>
+static FailureOr<Value> emitReceiveAssembly(const EmitReceiveAssemblyRun &run, Value lane, unsigned laneCount, Value initial,
+                                            DeferredEmissionContext &context, Insert insert) {
+  if (run.entryOffsets.size() != run.positions.size() + 1 || run.entryLanes.size() != run.positions.size() || run.positions.empty())
+    return failure();
+  LogicalTransferMetadataView logical = buildMetadataView(run.slices);
+  size_t actionCount = run.positions.size();
+  SmallVector<int64_t> counts(laneCount);
+  std::optional<TransferGrids> transferGrids;
+  std::optional<StaticIntGrid> positions;
+  SmallVector<LogicalTransferMetadataView, 0> metadataByEntry;
+  bool rectangular = run.lanes.size() == laneCount && llvm::all_of(run.entryLanes, [&](const LaneSet &lanes) { return lanes.size() == laneCount; });
+  for (size_t entry = 0; rectangular && entry < run.positions.size(); ++entry) {
+    ArrayRef<ScheduledTransferSlice> slices =
+        ArrayRef(run.slices).slice(run.entryOffsets[entry], run.entryOffsets[entry + 1] - run.entryOffsets[entry]);
+    SmallVector<const ScheduledTransferSlice *> ordered;
+    for (const ScheduledTransferSlice &slice : slices)
+      ordered.push_back(&slice);
+    llvm::sort(ordered, [](const ScheduledTransferSlice *left, const ScheduledTransferSlice *right) {
+      LaneInterval leftFamily = left->family->targetLanes.intervals().front();
+      LaneInterval rightFamily = right->family->targetLanes.intervals().front();
+      return leftFamily.begin + left->familyOffset < rightFamily.begin + right->familyOffset;
+    });
+    unsigned covered = 0;
+    LogicalTransferMetadataView metadata;
+    for (const ScheduledTransferSlice *slice : ordered) {
+      LaneInterval family = slice->family->targetLanes.intervals().front();
+      unsigned begin = family.begin + slice->familyOffset;
+      if (begin != covered) {
+        rectangular = false;
+        break;
+      }
+      appendMetadata(*slice, metadata);
+      covered += slice->transferCount;
+    }
+    rectangular &= covered == laneCount;
+    if (rectangular)
+      metadataByEntry.push_back(std::move(metadata));
+  }
+  if (rectangular) {
+    auto buildRows = [&](auto member) {
+      SmallVector<StaticIntSequence> rows;
+      for (const LogicalTransferMetadataView &metadata : metadataByEntry)
+        rows.push_back((metadata.*member).canonicalize());
+      return StaticIntGrid::fromRows(rows);
+    };
+    auto grids = buildTransferGrids(buildRows);
+    SmallVector<StaticIntSequence> positionRows;
+    for (unsigned position : run.positions)
+      positionRows.push_back(StaticIntSequence::uniform(position, laneCount));
+    auto positionGrid = StaticIntGrid::fromRows(positionRows);
+    if (failed(grids) || failed(positionGrid))
+      return failure();
+    transferGrids = std::move(*grids);
+    positions = std::move(*positionGrid);
+    llvm::fill(counts, actionCount);
+  } else {
+    SmallVector<LogicalTransferMetadataView, 0> metadataByLane(laneCount);
+    SmallVector<StaticIntSequenceChain, 0> positionsByLane(laneCount);
+    for (size_t entry = 0; entry < run.positions.size(); ++entry) {
+      ArrayRef<ScheduledTransferSlice> slices =
+          ArrayRef(run.slices).slice(run.entryOffsets[entry], run.entryOffsets[entry + 1] - run.entryOffsets[entry]);
+      for (const ScheduledTransferSlice &slice : slices) {
+        LaneInterval family = slice.family->targetLanes.intervals().front();
+        unsigned begin = family.begin + slice.familyOffset;
+        for (unsigned targetLane = begin; targetLane < begin + slice.transferCount; ++targetLane) {
+          ScheduledTransferSlice selected = slice;
+          selected.familyOffset += targetLane - begin;
+          selected.transferCount = 1;
+          appendMetadata(selected, metadataByLane[targetLane]);
+          positionsByLane[targetLane].append(
+            StaticIntSequence::uniform(run.positions[entry], 1));
+        }
+      }
+    }
+    actionCount = 0;
+    for (unsigned targetLane = 0; targetLane < laneCount; ++targetLane) {
+      counts[targetLane] = metadataByLane[targetLane].size();
+      actionCount = std::max(actionCount, metadataByLane[targetLane].size());
+    }
+    if (actionCount == 0)
+      return failure();
+    auto buildGrid = [&](auto member) {
+      const StaticIntSequenceChain &first = logical.*member;
+      int64_t defaultValue = first.valueAt(0);
+      SmallVector<StaticIntSequence> columns;
+      for (const LogicalTransferMetadataView &metadata : metadataByLane)
+        columns.push_back(canonicalizePadded(metadata.*member, actionCount, defaultValue));
+      return StaticIntGrid::fromColumns(actionCount, columns, defaultValue);
+    };
+    auto grids = buildTransferGrids(buildGrid);
+    SmallVector<StaticIntSequence> positionColumns;
+    for (const StaticIntSequenceChain &values : positionsByLane)
+      positionColumns.push_back(
+        canonicalizePadded(values, actionCount, run.positions.front()));
+    auto positionGrid = StaticIntGrid::fromColumns(
+      actionCount, positionColumns, run.positions.front());
+    if (failed(grids) || failed(positionGrid))
+      return failure();
+    transferGrids = std::move(*grids);
+    positions = std::move(*positionGrid);
+  }
+  if (!run.collection)
+    return failure();
+  Operation *anchor = run.collection->key.exchange->deferred.getOperation();
+  Location loc = anchor->getLoc();
+  Value runtimeLane = lane ? lane : context.constants.getIndex(0);
+  auto emitEntry = [&](Value entry, Value current) -> FailureOr<Value> {
+    Type fragmentType = run.slices.front().family->requirement->publicationFragmentType;
+    auto receive =
+        SpatChannelReceiveOp::create(context.rewriter, loc, fragmentType,
+                                     transferGrids->channels().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
+                                     transferGrids->sourceCores().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc),
+                                     transferGrids->targetCores().emitLookup(entry, runtimeLane, anchor, context.constants, context.rewriter, loc));
+    setLogicalTransferMetadata(receive, logical);
+    Value position = positions->emitLookup(
+      entry, runtimeLane, anchor, context.constants, context.rewriter, loc);
+    return insert(receive.getOutput(), position, entry, runtimeLane, current);
+  };
+  if (actionCount == 1 && llvm::all_of(counts, [](int64_t count) { return count == 1; }))
+    return emitEntry(context.constants.getIndex(0), initial);
+  Value count = emitStaticIntLookup(StaticIntSequence::fromValues(counts), runtimeLane, anchor, context.constants, context.rewriter, loc);
+  auto loop = buildNormalizedScfFor(context.rewriter, loc, context.constants.getIndex(0), count, 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();
+  return loop->results.front();
+}
+
+template <typename Emit>
+static LogicalResult emitCollectionUpdate(const LaneSet &lanes, Value lane, unsigned laneCount, FragmentCollectionKey key, Value current,
+                                          Operation *anchor, DeferredEmissionContext &context, Emit emit, bool local = false) {
+  FailureOr<Value> value;
+  if (local && lanes.size() != laneCount) {
+    if (!lane) return failure();
+    auto condition = emitLaneCondition(lanes, lane, laneCount, anchor, context, anchor->getLoc());
+    if (failed(condition)) return failure();
+    auto conditional = scf::IfOp::create(context.rewriter, anchor->getLoc(), TypeRange{current.getType()}, *condition, true);
+    OpBuilder::InsertionGuard guard(context.rewriter);
+    context.rewriter.setInsertionPointToStart(&conditional.getThenRegion().front());
+    value = emit(current);
+    if (failed(value)) return failure();
+    scf::YieldOp::create(context.rewriter, anchor->getLoc(), *value);
+    context.rewriter.setInsertionPointToStart(&conditional.getElseRegion().front());
+    scf::YieldOp::create(context.rewriter, anchor->getLoc(), current);
+    value = conditional.getResult(0);
+  } else value = emit(current);
+  if (failed(value))
+    return failure();
+  context.fragmentCollections[key] = *value;
+  return success();
+}
+
+static FailureOr<Value> insertProjectionFragment(Value fragment, Value specialization, Value position, Value geometryRow, Value runtimeLane,
+                                                 Value assembled, const DeferredProjectionLeafTemplate &leaf, const GridGeometry &geometry,
+                                                 DeferredExchangePlan &exchange, bool grouped, DeferredEmissionContext &context) {
+  Value shaped = fragment;
+  if (leaf.form == DeferredLeafForm::GraphBatchProjection) {
+    SmallVector<int64_t> shape(leaf.leadingRankReduced ? leaf.reconstructedType.getShape() : leaf.reconstructedType.getShape().drop_front());
+    shaped = extractMixedSliceOrIdentity(context.rewriter, exchange.deferred.getLoc(), shaped,
+                                         RankedTensorType::get(shape, leaf.reconstructedType.getElementType()),
+                                         lookupGeometry(geometry, geometryRow, runtimeLane, exchange.deferred, context, exchange.deferred.getLoc()));
+    if (!shaped) return failure();
+  }
+  auto sourceType = dyn_cast<RankedTensorType>(shaped.getType());
+  auto assembledType = dyn_cast<RankedTensorType>(assembled.getType());
+  int64_t rankDifference = sourceType && assembledType ? assembledType.getRank() - sourceType.getRank() : 0;
+  if (rankDifference < 0 || rankDifference > 2 || (grouped && rankDifference == 0)) return failure();
+  if (rankDifference == 0 && sourceType != assembledType) return failure();
+  MixedSliceGeometry slice;
+  slice.offsets.assign(assembledType.getRank(), context.rewriter.getIndexAttr(0));
+  if (rankDifference) slice.offsets.front() = grouped ? specialization : position;
+  if (rankDifference == 2) slice.offsets[1] = position;
+  slice.sizes.assign(rankDifference, context.rewriter.getIndexAttr(1));
+  for (int64_t dimension : sourceType.getShape()) slice.sizes.push_back(context.rewriter.getIndexAttr(dimension));
+  slice.strides.assign(assembledType.getRank(), context.rewriter.getIndexAttr(1));
+  return insertMixedSlice(context.rewriter, exchange.deferred.getLoc(), shaped, assembled, slice);
+}
+
+static Value createCollectionInitial(const FragmentCollectionPlan &collection, DeferredEmissionContext &context) {
+  DeferredExchangePlan &exchange = *collection.key.exchange;
+  if (collection.key.kind == FragmentCollectionKind::InsertAssembly)
+    return context.rewriter.clone(*exchange.program.insertAssembly->initialValue)->getResult(0);
+  return tensor::EmptyOp::create(context.rewriter, exchange.deferred.getLoc(), collection.collectionType.getShape(),
+                                 collection.collectionType.getElementType());
+}
+
+static LogicalResult emitLeafCollectionUpdate(const EmitReceiveAssemblyRun &run, Value lane, unsigned laneCount,
+                                              const DeferredResultPlan &resultPlan, DeferredEmissionContext &context) {
+  const FragmentCollectionPlan &collection = *run.collection;
+  DeferredExchangePlan &exchange = *collection.key.exchange;
+  unsigned leafIndex = collection.key.leafIndex;
+  const DeferredProjectionLeafTemplate &leaf = exchange.program.leaves[leafIndex];
+  bool grouped = collection.key.kind == FragmentCollectionKind::GroupedLeaf;
+  FragmentCollectionKey key = collection.key;
+  if (collection.key.kind == FragmentCollectionKind::Leaf && collection.positionCount == 1
+      && run.lanes.size() == laneCount
+      && llvm::all_of(run.positions, [](unsigned position) { return position == 0; })
+      && run.slices.front().family->requirement->publicationFragmentType == collection.collectionType) {
+    auto value = emitReceiveValue(run.slices, lane, laneCount, context);
+    if (failed(value)) return failure();
+    context.fragmentCollections[key] = *value;
+    return success();
+  }
+  Value current = context.fragmentCollections.lookup(collection.key);
+  if (!current) current = createCollectionInitial(collection, context);
+  const GridGeometry &geometry = resultPlan.innerGeometry[leafIndex];
+  auto emit = [&](Value initial) {
+    return emitReceiveAssembly(run, lane, laneCount, initial, context,
+        [&](Value fragment, Value position, Value, Value runtimeLane, Value assembled) -> FailureOr<Value> {
+          Value specialization = context.constants.getIndex(0);
+          Value leafPosition = position;
+          if (grouped) {
+            Value divisor = context.constants.getIndex(collection.positionCount);
+            specialization = arith::DivUIOp::create(context.rewriter, exchange.deferred.getLoc(), position, divisor);
+            leafPosition = arith::RemUIOp::create(context.rewriter, exchange.deferred.getLoc(), position, divisor);
+          }
+          return insertProjectionFragment(fragment, specialization, leafPosition, grouped ? specialization : context.constants.getIndex(0), runtimeLane,
+                                          assembled, leaf, geometry, exchange, grouped, context);
+        });
+  };
+  return emitCollectionUpdate(run.lanes, lane, laneCount, key, current, exchange.deferred, context, emit);
+}
+
+static FailureOr<Value> transformAssemblySource(Value fragment, const DeferredInsertAssemblyEntryTemplate &entry,
+                                                DeferredExchangePlan &exchange, DeferredEmissionContext &context) {
+  switch (entry.sourceTransform) {
+  case DeferredAssemblySourceTransform::Identity:
+    return fragment.getType() == entry.sourceType ? FailureOr<Value>(fragment) : FailureOr<Value>(failure());
+  case DeferredAssemblySourceTransform::AddLeadingUnitDimension:
+    return addLeadingUnitTensorDimension(context.rewriter, exchange.deferred.getLoc(), fragment);
+  case DeferredAssemblySourceTransform::RemoveLeadingUnitDimension:
+    return removeLeadingUnitTensorDimension(context.rewriter, exchange.deferred.getLoc(), fragment, entry.sourceType);
+  }
+  llvm_unreachable("unknown deferred assembly source transform");
+}
+
+static LogicalResult emitInsertAssemblyUpdate(const EmitReceiveAssemblyRun &run, Value lane, unsigned laneCount,
+                                              const DeferredResultPlan &resultPlan, DeferredEmissionContext &context) {
+  const FragmentCollectionPlan &collection = *run.collection;
+  DeferredExchangePlan &exchange = *collection.key.exchange;
+  const DeferredInsertAssemblyTemplate &assembly = *exchange.program.insertAssembly;
+  if (run.positions.empty()) return failure();
+  const DeferredInsertAssemblyEntryTemplate &sourceEntry = assembly.entries[run.positions.front()];
+  Value current = context.fragmentCollections.lookup(collection.key);
+  if (!current) current = createCollectionInitial(collection, context);
+  auto emit = [&](Value initial) {
+    return emitReceiveAssembly(run, lane, laneCount, initial, context,
+      [&](Value fragment, Value position, Value, Value runtimeLane, Value assembled) -> FailureOr<Value> {
+        auto shaped = transformAssemblySource(fragment, sourceEntry, exchange, context);
+        if (failed(shaped) || shaped->getType() != sourceEntry.sourceType) return failure();
+        return insertMixedSlice(context.rewriter, exchange.deferred.getLoc(), *shaped, assembled,
+          lookupGeometry(resultPlan.assemblyGeometry, position, runtimeLane, exchange.deferred, context, exchange.deferred.getLoc()));
+      });
+  };
+  return emitCollectionUpdate(run.lanes, lane, laneCount, collection.key, current, exchange.deferred, context, emit);
+}
+
+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();
+  Operation *anchor = run.slices.front().family->sourceScheduled->op;
+  Location loc = run.slices.front().family->requirement->exchange->deferred.getLoc();
+  auto condition = emitLaneCondition(run.lanes, lane, laneCount, anchor, context, loc);
+  if (failed(condition))
+    return failure();
+  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 EmitLocalCollectionRun &local, Value lane, unsigned laneCount, DeferredEmissionContext &context) {
+  RequirementFamily &reference = *local.families.front()->requirement;
+  Value fragment = reference.producer->payload;
+  if (reference.producerProjection || fragment.getType() != reference.publicationFragmentType) {
+    auto buildLaneGrid = [&](auto getSequence, int64_t defaultValue) {
+      SmallVector<int64_t> values(laneCount, defaultValue);
+      for (LocalAvailabilityFamily *family : local.families) {
+        RequirementFamily &requirement = *family->requirement;
+        const StaticIntSequence *sequence = getSequence(requirement);
+        if (!sequence)
+          continue;
+        LaneInterval lanes = requirement.targetLanes.intervals().front();
+        size_t laneSize = lanes.end - lanes.begin;
+        if (lanes.end > laneCount || sequence->size() < laneSize)
+          return FailureOr<StaticIntGrid>(failure());
+        for (size_t offset = 0; offset < laneSize; ++offset)
+          values[lanes.begin + offset] = sequence->valueAt(offset);
+      }
+      StaticIntSequence row = StaticIntSequence::fromValues(values);
+      return StaticIntGrid::fromRows(ArrayRef<StaticIntSequence>(row));
+    };
+    auto offsets = buildLaneGrid(
+        [](RequirementFamily &requirement) { return requirement.producerLocalOffsets ? &*requirement.producerLocalOffsets : nullptr; }, 0);
+    GridGeometry projectionGrids;
+    for (auto [geometryIndex, sourceMember] : llvm::enumerate(staticGeometryMembers)) {
+      if (!reference.producerProjection)
+        break;
+      StaticGeometryMember member = sourceMember;
+      const auto &referenceValues = reference.producerProjection.value().*member;
+      for (size_t dimension = 0; dimension < referenceValues.size(); ++dimension) {
+        auto grid = buildLaneGrid(
+            [&](RequirementFamily &requirement) -> const StaticIntSequence * {
+              if (!requirement.producerProjection)
+                return nullptr;
+              const auto &values = requirement.producerProjection.value().*member;
+              return dimension < values.size() ? &values[dimension] : nullptr;
+            },
+            referenceValues[dimension].valueAt(0));
+        if (failed(grid))
+          return failure();
+        projectionGrids[geometryIndex].push_back(std::move(*grid));
+      }
+    }
+    if (failed(offsets))
+      return failure();
+    Location loc = reference.exchange->deferred.getLoc();
+    Value position = lane ? lane : context.constants.getIndex(0);
+    Value row = context.constants.getIndex(0);
+    MixedSliceGeometry projection = lookupGeometry(projectionGrids, row, position, reference.exchange->deferred, context, loc);
+    auto materialized =
+        materializeSendPayload(reference, offsets->emitLookup(row, position, reference.exchange->deferred, context.constants, context.rewriter, loc),
+                               projectionGrids[0].empty() ? nullptr : &projection, 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 LogicalResult emitLocalCollectionUpdate(const EmitLocalCollectionRun &update, Value lane, unsigned laneCount, const DeferredResultPlan &resultPlan,
+                                               DeferredEmissionContext &context) {
+  const FragmentCollectionPlan &collection = *update.collection;
+  RequirementFamily &requirement = *update.families.front()->requirement;
+  DeferredExchangePlan &exchange = *requirement.exchange;
+  if (update.concatenatePayloads) {
+    SmallVector<Value> payloads;
+    for (LocalAvailabilityFamily *family : update.families)
+      payloads.push_back(family->requirement->producer->payload);
+    context.fragmentCollections[collection.key] = payloads.size() == 1
+      ? payloads.front()
+      : SpatConcatOp::create(
+          context.rewriter, exchange.deferred.getLoc(),
+          collection.collectionType, context.rewriter.getI64IntegerAttr(0),
+          payloads).getOutput();
+    return success();
+  }
+  auto materialize = [&]() { return materializeLocalValue(update, lane, laneCount, context); };
+  if (collection.key.kind == FragmentCollectionKind::Leaf
+      && collection.positionCount == 1
+      && update.lanes.size() == laneCount
+      && requirement.publicationFragmentType == collection.collectionType) {
+    auto fragment = materialize();
+    if (failed(fragment)) return failure();
+    context.fragmentCollections[collection.key] = *fragment;
+    return success();
+  }
+  Value current = context.fragmentCollections.lookup(collection.key);
+  if (!current) current = createCollectionInitial(collection, context);
+  if (collection.key.kind == FragmentCollectionKind::InsertAssembly) {
+    const auto &entry = exchange.program.insertAssembly->entries[update.collectionPosition];
+    auto emit = [&](Value assembled) -> FailureOr<Value> {
+      auto fragment = materialize();
+      if (failed(fragment)) return failure();
+      auto shaped = transformAssemblySource(*fragment, entry, exchange, context);
+      if (failed(shaped) || shaped->getType() != entry.sourceType) return failure();
+      return insertMixedSlice(context.rewriter, exchange.deferred.getLoc(), *shaped, assembled,
+        lookupGeometry(resultPlan.assemblyGeometry, context.constants.getIndex(update.collectionPosition), lane ? lane : context.constants.getIndex(0),
+                       exchange.deferred, context, exchange.deferred.getLoc()));
+    };
+    return emitCollectionUpdate(update.lanes, lane, laneCount, collection.key, current, exchange.deferred, context, emit, true);
+  }
+  bool grouped = collection.key.kind == FragmentCollectionKind::GroupedLeaf;
+  unsigned specialization = grouped ? update.collectionPosition / collection.positionCount : 0;
+  unsigned leafPosition = grouped ? update.collectionPosition % collection.positionCount : update.collectionPosition;
+  unsigned leafIndex = collection.key.leafIndex;
+  const DeferredProjectionLeafTemplate &leaf = exchange.program.leaves[leafIndex];
+  auto emit = [&](Value assembled) -> FailureOr<Value> {
+    auto fragment = materialize();
+    if (failed(fragment)) return failure();
+    return insertProjectionFragment(*fragment, context.constants.getIndex(specialization),
+                                    context.constants.getIndex(leafPosition), context.constants.getIndex(specialization),
+                                    lane ? lane : context.constants.getIndex(0), assembled, leaf, resultPlan.innerGeometry[leafIndex], exchange, grouped, context);
+  };
+  return emitCollectionUpdate(update.lanes, lane, laneCount, collection.key, current, exchange.deferred, context, emit, true);
+}
+
+static FailureOr<SmallVector<Value>> emitInstructions(ArrayRef<BoundaryInstruction> instructions, Value lane, unsigned laneCount,
+                                                      ArrayRef<DeferredResultPlan> results, DeferredEmissionContext &context) {
+  SmallVector<Value> produced;
+  for (const BoundaryInstruction &instruction : instructions) {
+    if (auto send = std::get_if<EmitSendRun>(&instruction)) {
+      if (failed(emitConditionalSendRun(*send, lane, laneCount, context)))
+        return failure();
+    } else if (auto update = std::get_if<EmitLocalCollectionRun>(&instruction)) {
+      DeferredExchangePlan *exchange = update->collection->key.exchange;
+      const DeferredResultPlan *resultPlan = findResultPlan(results, exchange);
+      LogicalResult emitted = resultPlan
+        ? emitLocalCollectionUpdate(*update, lane, laneCount, *resultPlan, context)
+        : failure();
+      if (failed(emitted))
+        return exchange->deferred.emitOpError(
+          "failed to update fragment collection from local availability"), failure();
+    } else if (auto assembly = std::get_if<EmitReceiveAssemblyRun>(&instruction)) {
+      DeferredExchangePlan *exchange = assembly->collection->key.exchange;
+      const DeferredResultPlan *resultPlan = findResultPlan(results, exchange);
+      LogicalResult emitted = resultPlan
+        ? (assembly->collection->key.kind == FragmentCollectionKind::InsertAssembly
+          ? emitInsertAssemblyUpdate(*assembly, lane, laneCount, *resultPlan, context)
+          : emitLeafCollectionUpdate(*assembly, lane, laneCount, *resultPlan, context))
+        : failure();
+      if (failed(emitted))
+        return exchange->deferred.emitOpError(
+          "failed to update fragment collection from received availability"), failure();
+    } else if (auto result = std::get_if<ProduceDeferredResult>(&instruction)) {
+      const DeferredResultPlan *plan = findResultPlan(results, result->exchange);
+      if (!plan)
+        return failure();
+      auto value = realizeDeferredResult(*plan, lane, context);
+      if (failed(value))
+        return result->exchange->deferred.emitOpError(
+          "failed to realize deferred result from fragment collections"), failure();
+      produced.push_back(*value);
+    }
+  }
+  return produced;
+}
+
+static SmallVector<DeferredExchangePlan *> getProducedExchanges(ArrayRef<BoundaryInstruction> instructions) {
+  SmallVector<DeferredExchangePlan *> exchanges;
+  for (const BoundaryInstruction &instruction : instructions)
+    if (auto result = std::get_if<ProduceDeferredResult>(&instruction))
+      exchanges.push_back(result->exchange);
+  return exchanges;
+}
+
+static void setInsertionAtBoundary(IRRewriter &rewriter, const BoundaryKey &key) {
+  ScheduledInfo &scheduled = *key.first;
+  if (key.second < scheduled.stepAnchors.size() && scheduled.stepAnchors[key.second]->getBlock()) {
+    rewriter.setInsertionPoint(scheduled.stepAnchors[key.second]);
+    return;
+  }
+  rewriter.setInsertionPoint(scheduled.blocks.front()->getTerminator());
+}
+
+static LogicalResult replaceResults(ArrayRef<DeferredExchangePlan *> exchanges, ValueRange replacements,
+                                    DeferredReplacementMap &deferredReplacements) {
+  if (exchanges.size() != replacements.size())
+    return failure();
+  for (auto [exchange, replacement] : llvm::zip_equal(exchanges, replacements)) {
+    deferredReplacements.insert({exchange->deferred, replacement});
+  }
+  return success();
+}
+
+static LogicalResult emitBoundary(const BoundaryProgram &boundary, ArrayRef<DeferredResultPlan> results, DeferredEmissionContext &context,
+                                  DeferredReplacementMap &replacements) {
+  setInsertionAtBoundary(context.rewriter, boundary.key);
+  unsigned laneCount = boundary.key.first->cores.size();
+  Value lane;
+  if (auto batch = dyn_cast<SpatScheduledComputeBatch>(boundary.key.first->op))
+    lane = *batch.getLaneArgument();
+  SmallVector<DeferredExchangePlan *> exchanges = getProducedExchanges(boundary.instructions);
+  auto values = emitInstructions(boundary.instructions, lane, laneCount, results, context);
+  return failed(values) ? failure() : replaceResults(exchanges, *values, replacements);
+}
+
+} // namespace
+
+LogicalResult realizeDeferredBoundaries(ArrayRef<BoundaryProgram> boundaries, ArrayRef<DeferredResultPlan> results, DeferredEmissionContext &context,
+                                        DeferredReplacementMap &replacements) {
+  ScheduledInfo *scheduled = nullptr;
+  for (const BoundaryProgram &boundary : boundaries) {
+    if (scheduled != boundary.key.first) {
+      context.fragmentCollections.clear();
+      scheduled = boundary.key.first;
+    }
+    if (failed(emitBoundary(boundary, results, context, replacements)))
+      return boundary.key.first->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,48 @@
+#pragma once
+
+#include "llvm/ADT/MapVector.h"
+
+#include "mlir/IR/PatternMatch.h"
+
+#include "DeferredBoundaryPlanning.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+
+namespace onnx_mlir::spatial {
+
+struct FragmentCollectionKeyInfo {
+  static FragmentCollectionKey getEmptyKey() {
+    return {llvm::DenseMapInfo<DeferredExchangePlan*>::getEmptyKey()};
+  }
+  static FragmentCollectionKey getTombstoneKey() {
+    return {llvm::DenseMapInfo<DeferredExchangePlan*>::getTombstoneKey()};
+  }
+  static unsigned getHashValue(const FragmentCollectionKey& key) {
+    return llvm::hash_combine(
+      key.exchange, key.kind, key.leafIndex);
+  }
+  static bool isEqual(const FragmentCollectionKey& lhs,
+                      const FragmentCollectionKey& rhs) {
+    return lhs == rhs;
+  }
+};
+
+struct DeferredEmissionContext {
+  DeferredEmissionContext(mlir::IRRewriter& rewriter,
+                          ConstantPool& constants)
+    : rewriter(rewriter), constants(constants) {}
+
+  mlir::IRRewriter& rewriter;
+  ConstantPool& constants;
+  llvm::DenseMap<FragmentCollectionKey, mlir::Value,
+                 FragmentCollectionKeyInfo> fragmentCollections;
+};
+
+using DeferredReplacementMap =
+  llvm::MapVector<mlir::Operation*, mlir::Value>;
+
+mlir::LogicalResult realizeDeferredBoundaries(mlir::ArrayRef<BoundaryProgram> boundaries,
+                                              mlir::ArrayRef<DeferredResultPlan> results,
+                                              DeferredEmissionContext& context,
+                                              DeferredReplacementMap& replacements);
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,372 @@
+#include "DeferredCommunicationDeadlock.hpp"
+
+#include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+#include "llvm/ADT/DenseMap.h"
+
+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 {
+  bool send = false;
+  StaticIntSequence channels;
+  StaticIntSequence parents;
+  StaticIntSequence counts;
+  StaticIntSequence sources;
+  StaticIntSequence targets;
+};
+
+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();
+  }
+  size_t size = 1;
+  if (batch) {
+    auto count = op->getAttrOfType<IntegerAttr>(
+      "raptor.batch_transfer_count");
+    if (!count || count.getInt() <= 0)
+      return op->emitOpError("has invalid compact transfer count"), failure();
+    size = count.getInt();
+  }
+  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 {isa<SpatChannelSendOp>(op),
+    *channels, *parents, *counts, *sources, *targets};
+}
+
+static void appendEventsByCore(
+    DenseMap<int64_t, StaticIntSequenceChain> &result,
+    const StaticIntSequence &channels, const StaticIntSequence &cores,
+    size_t begin, size_t count, bool send) {
+  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)
+        return;
+      SmallVector<int64_t> events;
+      events.reserve(selectedEnd - selectedBegin);
+      for (size_t index = selectedBegin; index < selectedEnd; ++index)
+        events.push_back(2 * channels.valueAt(index) + (send ? 0 : 1));
+      result[core].append(StaticIntSequence::fromValues(events));
+    });
+}
+
+static LogicalResult compareEventSequences(
+    func::FuncOp funcOp,
+    const DenseMap<int64_t, StaticIntSequenceChain> &expected,
+    const DenseMap<int64_t, StaticIntSequenceChain> &actual) {
+  if (expected.size() != actual.size())
+    return funcOp.emitOpError(
+      "realized communication stream set differs from plan");
+  for (const auto &[core, sequence] : expected) {
+    auto found = actual.find(core);
+    if (found == actual.end())
+      return funcOp.emitOpError()
+             << "realized communication 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;
+    auto describe = [](StaticIntSequenceChainCursor &cursor) {
+      if (cursor.done())
+        return std::pair<StringRef, int64_t>("none", -1);
+      int64_t event = cursor.value();
+      return std::pair<StringRef, int64_t>(
+        event % 2 == 0 ? "send" : "receive", event / 2);
+    };
+    auto [expectedKind, expectedChannel] = describe(expectedCursor);
+    auto [actualKind, actualChannel] = describe(actualCursor);
+    return funcOp.emitOpError()
+           << "realized communication order differs on core " << core
+           << " at ordinal " << ordinal << ": expected " << expectedKind
+           << " channel " << expectedChannel << ", actual " << actualKind
+           << " channel " << actualChannel;
+  }
+  return success();
+}
+
+} // namespace
+
+LogicalResult verifyPlannedCommunicationDeadlockFree(
+    Operation *anchor, ArrayRef<unsigned> stepCounts,
+    const ScheduledCommunicationPlan &plan) {
+  SmallVector<std::pair<int64_t, int64_t>> familyChannels;
+  DenseMap<ExternalTransferFamily *, unsigned> familyIndex;
+  for (const ScheduledTransferSlice &slice : plan.slices) {
+    ExternalTransferFamily *family = slice.family;
+    if (!familyIndex.try_emplace(family, familyIndex.size()).second)
+      continue;
+    size_t count = family->channelIds.size();
+    if (count == 0)
+      return anchor->emitError(
+        "planned communication family has no channels");
+    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 anchor->emitError(
+          "planned communication family has non-consecutive channels");
+    familyChannels.emplace_back(
+      first, first + static_cast<int64_t>(count));
+  }
+  llvm::sort(familyChannels);
+  int64_t nextChannel = 0;
+  for (auto [firstChannel, endChannel] : familyChannels) {
+    if (firstChannel != nextChannel)
+      return anchor->emitError(
+        "planned communication channels are not exactly contiguous");
+    nextChannel = endChannel;
+  }
+  if (static_cast<uint64_t>(nextChannel) != plan.logicalTransferCount)
+    return anchor->emitError(
+      "planned communication channel count is inconsistent");
+
+  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<ExternalTransferFamily *> familyByChannel(
+    plan.logicalTransferCount);
+  DenseMap<ExternalTransferFamily *, unsigned> familyIndex;
+  for (const ScheduledTransferSlice &slice : plan.slices) {
+    ExternalTransferFamily *family = slice.family;
+    if (!familyIndex.try_emplace(family, familyIndex.size()).second)
+      continue;
+    for (size_t index = 0; index < family->channelIds.size(); ++index)
+      familyByChannel[family->channelIds.valueAt(index)] = family;
+  }
+
+  DenseMap<int64_t, StaticIntSequenceChain> expected;
+  for (const ScheduledTransferSlice &slice : plan.slices) {
+    ExternalTransferFamily &family = *slice.family;
+    appendEventsByCore(expected, family.channelIds, family.sourceCores,
+                       slice.familyOffset, slice.transferCount, true);
+    appendEventsByCore(expected, family.channelIds, family.targetCores,
+                       slice.familyOffset, slice.transferCount, false);
+  }
+
+  DenseMap<int64_t, StaticIntSequenceChain> 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);
+      if (channel < 0
+          || static_cast<uint64_t>(channel) >= familyByChannel.size()) {
+        op->emitOpError("references an unknown logical channel");
+        invalid = true;
+        return;
+      }
+      ExternalTransferFamily *family = familyByChannel[channel];
+      if (!family) {
+        op->emitOpError("references an unknown logical channel");
+        invalid = true;
+        return;
+      }
+      size_t familyOffset = channel - family->channelIds.valueAt(0);
+      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)));
+    appendEventsByCore(actual, *actualChannels.back(),
+                       realized->send ? realized->sources : realized->targets,
+                       0, actualChannels.back()->size(), realized->send);
+  });
+  if (invalid)
+    return failure();
+  return compareEventSequences(funcOp, expected, actual);
+}
+
+} // 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,248 @@
+#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 specializationIndex = 0;
+  unsigned leafIndex = 0;
+  unsigned selectedPosition = 0;
+
+  bool operator==(const RequirementCoordinate& other) const {
+    return specializationIndex == other.specializationIndex
+        && leafIndex == other.leafIndex
+        && selectedPosition == other.selectedPosition;
+  }
+};
+
+enum class DeferredLeafForm {
+  DirectSource,
+  ScalarProjection,
+  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 DeferredStaticSliceGeometry {
+  llvm::SmallVector<StaticIntSequence> offsets;
+  llvm::SmallVector<StaticIntSequence> sizes;
+  llvm::SmallVector<StaticIntSequence> strides;
+};
+
+struct DeferredProjectionLeafTemplate {
+  DeferredLeafForm form = DeferredLeafForm::DirectSource;
+  mlir::Value sourceRoot;
+  mlir::Value replacementRoot;
+  DeferredSliceTemplate leadingGeometry;
+  DeferredSliceTemplate innerGeometry;
+  mlir::RankedTensorType reconstructedType;
+  bool leadingRankReduced = false;
+};
+
+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;
+  unsigned specializationCount = 1;
+  mlir::Value specializationArgument;
+  mlir::RankedTensorType specializationFragmentType;
+  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;
+  unsigned stepCount = 0;
+  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;
+  std::optional<DeferredStaticSliceGeometry> producerProjection;
+};
+
+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,485 @@
+#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/Common/IR/StaticIntSequence.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.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));
+    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));
+    }
+    SmallVector<int64_t> selectedFragmentShape;
+    selectedFragmentShape.reserve(rank);
+    for (int64_t dim = 0; dim < rank; ++dim)
+      selectedFragmentShape.push_back(sizes[fragmentIndex * rank + dim]);
+    auto fragmentType = RankedTensorType::get(
+      selectedFragmentShape, resultType.getElementType());
+    Value fragment = extractMixedSliceOrIdentity(
+      builder, loc, sourceBlockArgs[operandIndex], fragmentType,
+      {sliceOffsets, sliceSizes, sliceStrides});
+    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> buildSelectedDeferredSource(OpBuilder &builder, Location loc,
+                                                    Operation *diagnosticOwner,
+                                                    Value scheduledLane,
+                                                    ValueRange sourceBlockArgs,
+                                                    ArrayRef<int64_t> sourceOperandForScheduledLane) {
+  if (sourceBlockArgs.size() == 1)
+    return sourceBlockArgs.front();
+  if (!scheduledLane || sourceOperandForScheduledLane.empty())
+    return diagnosticOwner->emitOpError("multiple deferred sources require the enclosing scheduled lane"), failure();
+  for (int64_t sourceIndex : sourceOperandForScheduledLane) {
+    if (sourceIndex < 0 || sourceIndex >= static_cast<int64_t>(sourceBlockArgs.size()))
+      return diagnosticOwner->emitOpError("deferred source mapping operand is out of range"), failure();
+  }
+  StaticIntSequence sourceIndices = StaticIntSequence::fromValues(
+    sourceOperandForScheduledLane);
+  if (sourceIndices.getKind() == StaticIntSequenceKind::Uniform)
+    return sourceBlockArgs[sourceIndices.valueAt(0)];
+  Value selector;
+  if (sourceIndices.getKind() == StaticIntSequenceKind::Affine) {
+    int64_t step = sourceIndices.valueAt(1) - sourceIndices.valueAt(0);
+    selector = affineMulConst(
+      builder, loc, scheduledLane, step, diagnosticOwner);
+    selector = affineAddConst(
+      builder, loc, selector, sourceIndices.valueAt(0),
+      diagnosticOwner);
+  } else {
+    struct Run { int64_t value; size_t end; };
+    SmallVector<Run> runs;
+    sourceIndices.forEachEqualRun(
+      [&](int64_t value, size_t begin, size_t count) {
+        runs.push_back({value, begin + count});
+      });
+    selector = arith::ConstantIndexOp::create(
+      builder, loc, runs.back().value);
+    for (const Run &run : llvm::reverse(ArrayRef(runs).drop_back())) {
+      Value end = arith::ConstantIndexOp::create(builder, loc, run.end);
+      Value before = arith::CmpIOp::create(
+        builder, loc, arith::CmpIPredicate::ult, scheduledLane, end);
+      Value value = arith::ConstantIndexOp::create(builder, loc, run.value);
+      selector = arith::SelectOp::create(
+        builder, loc, before, value, selector);
+    }
+  }
+  return SpatDeferredSourceSelectOp::create(
+    builder, loc, sourceBlockArgs.front().getType(), selector,
+    sourceBlockArgs).getOutput();
+}
+
+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 (auto argument = dyn_cast<BlockArgument>(value);
+        argument && argument.getOwner() == &transfer.getBody().front()
+        && argument.getArgNumber() >= transfer.getSources().size())
+      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,
+                                    const DenseMap<ProducerValueKey, MaterializedProducerRef> &availableValues,
+                                    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 batch = dyn_cast<SpatComputeBatch>(producer->instance.op);
+  bool hasCompleteValue = !batch
+                       || (producer->instance.laneStart == 0
+                           && producer->instance.laneCount
+                                == static_cast<uint32_t>(batch.getLaneCount()));
+  if (auto available = availableValues.find(key);
+      hasCompleteValue && available != availableValues.end()) {
+    plan.availableValue = available->second.payload;
+    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());
+    Value sharedSource;
+    if (!plan.blueprint) {
+      OpBuilder::InsertPoint restore = builder.saveInsertionPoint();
+      if (plan.scalarizedGraphLaneBase && plan.originalSources.size() > 1) {
+        Operation *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 source is missing a hoist point";
+      }
+      auto selected = buildSelectedDeferredSource(
+        builder, loc, plan.graphInput.getOwner()->getParentOp(),
+        plan.scheduledLane, plan.originalSources,
+        plan.sourceOperandForScheduledLane);
+      if (failed(selected))
+        return failure();
+      sharedSource = *selected;
+      builder.restoreInsertionPoint(restore);
+    }
+    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";
+      }
+      unsigned count = scalarize ? plan.scalarizedLaneCount : 1;
+      bool grouped = count > 1;
+      auto fragmentType = dyn_cast<RankedTensorType>(root.getType());
+      if (!fragmentType || !fragmentType.hasStaticShape())
+        return emitError(loc) << "phase 1 deferred payload requires a static ranked tensor";
+      Type outputType = fragmentType;
+      if (grouped) {
+        SmallVector<int64_t> groupedShape {static_cast<int64_t>(count)};
+        llvm::append_range(groupedShape, fragmentType.getShape());
+        outputType = RankedTensorType::get(groupedShape,
+                                           fragmentType.getElementType());
+      }
+      auto specialization = scalarize
+                          ? builder.getI64IntegerAttr(count)
+                          : IntegerAttr();
+      SmallVector<Value> transferSources;
+      if (plan.blueprint)
+        llvm::append_range(transferSources, plan.originalSources);
+      else
+        transferSources.push_back(sharedSource);
+      auto transfer = SpatDeferredCommunicationOp::create(
+        builder, loc, outputType, transferSources, specialization);
+      SmallVector<Type> bodyTypes(transfer.getSources().getTypes());
+      if (grouped)
+        bodyTypes.push_back(builder.getIndexType());
+      Block *deferred = builder.createBlock(&transfer.getBody(), transfer.getBody().end(),
+        bodyTypes, SmallVector<Location>(bodyTypes.size(), loc));
+      builder.setInsertionPointToStart(deferred);
+      ValueRange sourceArguments = deferred->getArguments().take_front(
+        transfer.getSources().size());
+      auto selected = plan.blueprint
+                        ? buildBlueprintReconstruction(builder, loc, plan.blueprint,
+                                                       sourceArguments)
+                        : FailureOr<Value>(sourceArguments.front());
+      if (failed(selected)) return failure();
+      Value boundGraphLane;
+      if (grouped)
+        boundGraphLane = arith::AddIOp::create(
+          builder, loc, plan.scalarizedGraphLaneBase,
+          deferred->getArguments().back());
+      else if (scalarize)
+        boundGraphLane = plan.scalarizedGraphLaneBase;
+      auto payload = clonePayloadRoot(root, body, plan, builder, transfer, *selected, boundGraphLane);
+      if (failed(payload)) return failure();
+      SpatYieldOp::create(builder, loc, *payload);
+      builder.setInsertionPointAfter(transfer);
+      if (grouped) {
+        builder.restoreInsertionPoint(restore);
+        SmallVector<OpFoldResult> offsets {
+          plan.scalarizedLocalLane};
+        SmallVector<OpFoldResult> sizes {builder.getIndexAttr(1)};
+        SmallVector<OpFoldResult> strides {builder.getIndexAttr(1)};
+        for (int64_t dimension : fragmentType.getShape()) {
+          offsets.push_back(builder.getIndexAttr(0));
+          sizes.push_back(builder.getIndexAttr(dimension));
+          strides.push_back(builder.getIndexAttr(1));
+        }
+        mapper.map(root, extractMixedSliceOrIdentity(
+          builder, loc, transfer.getOutput(), fragmentType,
+          {offsets, sizes, strides}));
+      } else {
+        mapper.map(root, transfer.getOutput());
+      }
+      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,
+                                    const DenseMap<ProducerValueKey, MaterializedProducerRef> &availableValues,
+                                    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,171 @@
+#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 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()) {
+        Operation *user = use.getOwner();
+        if (isa<SpatGraphCompute, SpatGraphComputeBatch,
+                SpatDeferredCommunicationOp>(user))
+          continue;
+        if (auto blueprint = dyn_cast<SpatBlueprintOp>(user)) {
+          bool blueprintEscapes = llvm::any_of(
+            blueprint.getOutput().getUses(), [](OpOperand &blueprintUse) {
+              return !isa<SpatGraphCompute, SpatGraphComputeBatch,
+                          SpatDeferredCommunicationOp>(
+                blueprintUse.getOwner());
+            });
+          if (!blueprintEscapes)
+            continue;
+        }
+        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
+            && 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 *> oldGraph;
+  for (Operation &op : funcOp.getOps())
+    if (isa<SpatGraphCompute, SpatGraphComputeBatch, SpatBlueprintOp>(op))
+      oldGraph.push_back(&op);
+  for (Operation *op : llvm::reverse(oldGraph)) {
+    if (auto blueprint = dyn_cast<SpatBlueprintOp>(op)) {
+      if (blueprint.getOutput().use_empty())
+        rewriter.eraseOp(blueprint);
+      continue;
+    }
+    if (!op->use_empty())
+      return op->emitOpError(
+        "phase 2 cannot erase an old graph compute with live results");
+    rewriter.eraseOp(op);
+  }
+  return success();
+}
+
+static LogicalResult eraseDeferredSourceSelectors(
+    func::FuncOp funcOp, IRRewriter &rewriter) {
+  SmallVector<SpatDeferredSourceSelectOp> selectors;
+  funcOp.walk([&](SpatDeferredSourceSelectOp selector) {
+    selectors.push_back(selector);
+  });
+  for (SpatDeferredSourceSelectOp selector : llvm::reverse(selectors)) {
+    if (!selector.getOutput().use_empty())
+      return selector.emitOpError(
+        "phase 2 left a live deferred source selection");
+    rewriter.eraseOp(selector);
+  }
+  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,
+    const ScheduledComputeMaterializationResult &materialization) {
+  auto transfers = buildDeferredTransferPlan(funcOp, materialization);
+  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(retargetDeferredPublications(funcOp, *transfers))
+      || failed(replaceFinalGraphPublications(funcOp, *transfers)))
+    return failure();
+  ConstantPool constants(funcOp, rewriter);
+  DeferredEmissionContext context(rewriter, constants);
+  DeferredReplacementMap replacements;
+  if (failed(realizeDeferredBoundaries(
+        boundaries->boundaries, boundaries->results, context, replacements)))
+    return failure();
+  for (auto [op, replacement] : replacements) {
+    if (op->getResult(0) == replacement)
+      return op->emitOpError(
+        "phase 2 cannot replace deferred communication with itself");
+    op->getResult(0).replaceAllUsesWith(replacement);
+    if (!op->use_empty())
+      return op->emitOpError(
+        "phase 2 cannot erase deferred communication with live uses");
+    rewriter.eraseOp(op);
+  }
+  if (failed(eraseDeferredSourceSelectors(funcOp, rewriter))
+      || 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,13 @@
+#pragma once
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+
+namespace onnx_mlir::spatial {
+
+struct ScheduledComputeMaterializationResult;
+
+mlir::LogicalResult realizeDeferredCommunication(
+    mlir::func::FuncOp funcOp,
+    const ScheduledComputeMaterializationResult &materialization);
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,397 @@
+#include "llvm/ADT/MapVector.h"
+
+#include <limits>
+#include <queue>
+
+#include "DeferredCommunicationScheduling.hpp"
+#include "DeferredTransferPlanning.hpp"
+
+namespace onnx_mlir::spatial {
+using namespace mlir;
+namespace {
+
+using TransferEmissionSignature =
+  std::tuple<ScheduledInfo*, Value, Type, bool, bool, bool>;
+
+static TransferEmissionSignature getTransferEmissionSignature(
+    const ExternalTransferFamily& family) {
+  ProducedValue* producer = family.requirement->producer;
+  return {producer->scheduled,
+          producer->payload,
+          family.requirement->publicationFragmentType,
+          family.requirement->graphLanes.has_value(),
+          family.requirement->producerProjection.has_value(),
+          producer->scheduled->isBatch()};
+}
+
+struct StreamThreshold {
+  unsigned stream = 0;
+  unsigned completedStep = 0;
+};
+
+struct TransferGroup {
+  DeferredExchangePlan* exchange = nullptr;
+  SmallVector<ScheduledTransferSlice> ordered;
+  SmallVector<StreamThreshold> dependencies;
+  unsigned unsatisfied = 0;
+  bool scheduled = false;
+  std::tuple<unsigned, unsigned, unsigned, uint64_t> originalPriority;
+  std::tuple<unsigned, unsigned, unsigned,
+             std::tuple<unsigned, unsigned, unsigned, uint64_t>> staticPriority;
+};
+
+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(std::get<0>(signature),
+                            std::get<1>(signature).getAsOpaquePointer(),
+                            std::get<2>(signature).getAsOpaquePointer(),
+                            std::get<3>(signature),
+                            std::get<4>(signature),
+                            std::get<5>(signature));
+}
+
+static bool betterStaticPriority(const TransferGroup& lhs, const TransferGroup& rhs) {
+  return lhs.staticPriority < rhs.staticPriority;
+}
+
+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]);
+  }
+  unsigned instructionCount = 0;
+  unsigned absorbedTransfers = 0;
+  unsigned transferCount = 0;
+  std::optional<TransferEmissionSignature> previousSignature;
+  RequirementFamily* previousRequirement = nullptr;
+  for (const ScheduledTransferSlice& slice : group.ordered) {
+    TransferEmissionSignature signature =
+      getTransferEmissionSignature(*slice.family);
+    if (!previousSignature || *previousSignature != signature)
+      ++instructionCount;
+    else
+      absorbedTransfers += slice.transferCount;
+    previousSignature = signature;
+    if (previousRequirement != slice.family->requirement)
+      ++instructionCount;
+    previousRequirement = slice.family->requirement;
+    transferCount += slice.transferCount;
+  }
+  group.staticPriority =
+    {instructionCount,
+     std::numeric_limits<unsigned>::max() - absorbedTransfers,
+     transferCount,
+     group.originalPriority};
+}
+
+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
+
+bool haveSameTransferEmissionSignature(
+    const ExternalTransferFamily& lhs, const ExternalTransferFamily& rhs) {
+  return getTransferEmissionSignature(lhs) == getTransferEmissionSignature(rhs);
+}
+
+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];
+    ready.push(index);
+    auto& bucket = bySignature[hashSignature(
+      getTransferEmissionSignature(*group.ordered.front().family))];
+    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.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].scheduled)
+        chosen = candidate;
+    }
+    if (chosen) {
+      TransferGroup& group = groups[*chosen];
+      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,30 @@
+#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;
+};
+
+bool haveSameTransferEmissionSignature(
+  const ExternalTransferFamily& lhs, const ExternalTransferFamily& rhs);
+
+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,621 @@
+#include "DeferredProjectionAnalysis.hpp"
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.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 compactSelection = result
+    ? dyn_cast<SpatDeferredSourceSelectOp>(result.getOwner())
+    : SpatDeferredSourceSelectOp();
+  if (compactSelection && result.getResultNumber() == 0) {
+    auto selector = evaluateDeferredIndex(
+      compactSelection.getSelector(), environment);
+    if (failed(selector) || *selector < 0
+        || *selector >= static_cast<int64_t>(
+             compactSelection.getSources().size()))
+      return failure();
+    return sourceArgument(
+      compactSelection.getSources()[*selector], deferred, environment);
+  }
+  return std::optional<unsigned>();
+}
+
+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 compactSelection = result
+    ? dyn_cast<SpatDeferredSourceSelectOp>(result.getOwner())
+    : SpatDeferredSourceSelectOp();
+  if (compactSelection && result.getResultNumber() == 0) {
+    for (Value source : compactSelection.getSources())
+      collectSourceArguments(source, 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 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 && 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 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 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());
+}
+
+static LogicalResult validateDeferredProgram(
+    SpatDeferredCommunicationOp deferred, Block &body, Value scheduledLane,
+    Value specialization, int64_t laneCount, int64_t specializationCount) {
+  auto specializes = [&](Value value, auto &&accept) {
+    for (int64_t specializationIndex = 0;
+         specializationIndex < specializationCount; ++specializationIndex)
+      for (int64_t lane = 0; lane < laneCount; ++lane) {
+        StaticIndexEnvironment environment;
+        if (scheduledLane)
+          environment.bindings[scheduledLane] = lane;
+        if (specialization)
+          environment.bindings[specialization] = specializationIndex;
+        auto result = evaluateDeferredIndex(value, environment);
+        if (failed(result) || !accept(*result))
+          return false;
+      }
+    return true;
+  };
+  WalkResult result = body.walk([&](Operation *op) -> WalkResult {
+    auto reject = [&](StringRef message) {
+      op->emitOpError(message);
+      return WalkResult::interrupt();
+    };
+    if (isa<SpatYieldOp, scf::YieldOp>(op)
+        || (isa<linalg::YieldOp>(op)
+            && isa<linalg::TransposeOp>(op->getParentOp())))
+      return WalkResult::advance();
+    if (auto selection = dyn_cast<SpatDeferredSourceSelectOp>(op)) {
+      if (selection->getBlock() != &body)
+        return reject("must be a top-level deferred source selector");
+      if (selection.getSources().empty()
+          || !selection.getSelector().getType().isIndex()
+          || llvm::any_of(selection.getSources(), [&](Value source) {
+               return source.getType() != selection.getOutput().getType()
+                   || failed(getPossibleDeferredSourceOperandIndices(
+                        source, deferred));
+             }))
+        return reject("has invalid deferred source operands or types");
+      if (!specializes(selection.getSelector(), [&](int64_t index) {
+            return index >= 0
+                && index < static_cast<int64_t>(selection.getSources().size());
+          }))
+        return reject("source selector does not specialize to a valid source");
+      return WalkResult::advance();
+    }
+    if (auto loop = dyn_cast<scf::ForOp>(op)) {
+      for (Value bound : {loop.getLowerBound(), loop.getUpperBound()})
+        if (!specializes(bound, [](int64_t) { return true; }))
+          return reject("has a loop bound that does not specialize");
+      if (!specializes(loop.getStep(), [](int64_t step) { return step > 0; }))
+        return reject("has a non-positive or non-static deferred loop step");
+      return WalkResult::advance();
+    }
+    if (op->getNumRegions() != 0 && !isa<linalg::TransposeOp>(op))
+      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 WalkResult::advance();
+  });
+  return failure(result.wasInterrupted());
+}
+
+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 = {0, 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
+
+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();
+}
+
+DeferredLaneValueEvaluator::DeferredLaneValueEvaluator(
+    const DeferredProgramTemplate &program, unsigned laneCount,
+    unsigned specializationIndex)
+  : program(program), laneCount(laneCount),
+    specializationIndex(specializationIndex) {}
+
+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;
+    if (program.specializationArgument)
+      environment.bindings[program.specializationArgument] =
+        specializationIndex;
+    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;
+    if (program.specializationArgument)
+      environment.bindings[program.specializationArgument] =
+        specializationIndex;
+    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) {
+  if (!deferred.getBody().hasOneBlock())
+    return deferred.emitOpError(
+      "deferred program must have exactly one body block"), failure();
+  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();
+  auto specialization = deferred->getAttrOfType<IntegerAttr>(
+    "specialization_count");
+  program.specializationCount = specialization
+                              ? specialization.getInt()
+                              : 1;
+  if (program.specializationCount > 1) {
+    program.specializationArgument = body.getArguments().back();
+    program.specializationFragmentType = dyn_cast<RankedTensorType>(
+      program.yieldedValue.getType());
+    if (!program.specializationFragmentType)
+      return deferred.emitOpError(
+        "grouped deferred fragment must be a ranked tensor"), failure();
+  }
+  if (auto scheduled = deferred->getParentOfType<SpatScheduledComputeBatch>())
+    program.scheduledLane = getEnclosingScheduledLane(deferred, scheduled);
+  int64_t laneCount = 1;
+  if (auto scheduled = deferred->getParentOfType<SpatScheduledComputeBatch>())
+    laneCount = scheduled.getLaneCount();
+  if (failed(validateDeferredProgram(
+        deferred, body, program.scheduledLane,
+        program.specializationArgument, laneCount,
+        program.specializationCount)))
+    return failure();
+
+  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());
+           });
+      bool scalarProjection = succeeded(sources)
+        && llvm::all_of(*sources, [&](unsigned index) {
+             auto result = dyn_cast<OpResult>(deferred.getSources()[index]);
+             return result && isa<SpatGraphCompute>(result.getOwner());
+           });
+      if (graphProjection || scalarProjection) {
+        DeferredProjectionLeafTemplate leaf;
+        leaf.form = graphProjection ? DeferredLeafForm::GraphBatchProjection
+                                    : DeferredLeafForm::ScalarProjection;
+        leaf.sourceRoot = slice.getSource();
+        leaf.replacementRoot = value;
+        leaf.leadingGeometry = {
+          SmallVector<OpFoldResult>(slice.getMixedOffsets()),
+          SmallVector<OpFoldResult>(slice.getMixedSizes()),
+          SmallVector<OpFoldResult>(slice.getMixedStrides())};
+        if (graphProjection)
+          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());
+        if (graphProjection
+            && slice.getSourceType().getRank()
+                 == leaf.reconstructedType.getRank() + 1
+            && slice.getMixedSizes().size()
+                 == static_cast<size_t>(slice.getSourceType().getRank())) {
+          ArrayRef<OpFoldResult> innerSizes =
+            ArrayRef(slice.getMixedSizes()).drop_front();
+          leaf.leadingRankReduced = llvm::equal(
+            innerSizes, leaf.reconstructedType.getShape(),
+            [](OpFoldResult size, int64_t dimension) {
+              return getConstantIntValue(size) == dimension;
+            });
+        }
+        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) {
+  std::pair<Operation *, unsigned> 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,55 @@
+#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::FailureOr<DeferredProgramTemplate> analyzeDeferredProgramTemplate(
+  SpatDeferredCommunicationOp deferred);
+
+class DeferredLaneValueEvaluator {
+public:
+  DeferredLaneValueEvaluator(const DeferredProgramTemplate &program,
+                             unsigned laneCount,
+                             unsigned specializationIndex = 0);
+
+  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;
+  unsigned specializationIndex;
+  llvm::DenseMap<mlir::Value, StaticIntSequence> values;
+  llvm::DenseMap<mlir::Value, StaticIntSequence> sourceOperands;
+};
+
+struct GraphBatchPublicationMap {
+  mlir::RankedTensorType physicalResultType;
+  mlir::RankedTensorType publicationFragmentType;
+  llvm::SmallVector<int64_t> graphLaneToPhysicalSlot;
+  llvm::SmallVector<int64_t> physicalSlotToGraphLane;
+};
+
+using GraphBatchPublicationCache =
+  llvm::DenseMap<std::pair<mlir::Operation *, unsigned>,
+                 GraphBatchPublicationMap>;
+
+mlir::FailureOr<const GraphBatchPublicationMap *> getGraphBatchPublicationMap(
+  SpatGraphComputeBatch graphBatch, unsigned resultIndex,
+  GraphBatchPublicationCache &cache);
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,338 @@
+#include "DeferredResultRealization.hpp"
+#include "DeferredBoundaryRealization.hpp"
+#include "DeferredProjectionAnalysis.hpp"
+#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/DenseSet.h"
+#include <array>
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
+namespace onnx_mlir::spatial {
+using namespace mlir; namespace {
+static MixedSliceGeometry leadingSlice(OpBuilder &builder, RankedTensorType type, OpFoldResult position) {
+  MixedSliceGeometry geometry;
+  geometry.offsets.assign(type.getRank() + 1, builder.getIndexAttr(0));
+  geometry.offsets.front() = position;
+  geometry.sizes.push_back(builder.getIndexAttr(1));
+  for (int64_t dimension : type.getShape()) geometry.sizes.push_back(builder.getIndexAttr(dimension));
+  geometry.strides.assign(type.getRank() + 1, builder.getIndexAttr(1));
+  return geometry;
+}
+static LogicalResult verifyCollectionCoverage(
+    DeferredExchangePlan &exchange, FragmentCollectionPlan &collection,
+    unsigned slotCount) {
+  LaneSet all = LaneSet::all(exchange.targetLaneCount);
+  SmallVector<LaneSet, 0> coverage(slotCount);
+  for (const FragmentCollectionPlan::Requirement &entry :
+       collection.requirements) {
+    if (entry.position >= coverage.size()
+        || !coverage[entry.position]
+              .intersect(entry.family->targetLanes).empty())
+      return exchange.deferred.emitOpError(
+               "fragment collection has overlapping coverage at slot ")
+             << entry.position;
+    coverage[entry.position] =
+      coverage[entry.position].unite(entry.family->targetLanes);
+  }
+  for (auto [position, lanes] : llvm::enumerate(coverage))
+    if (!(lanes == all))
+      return exchange.deferred.emitOpError(
+               "fragment collection has incomplete coverage at slot ")
+             << position;
+  return success();
+}
+
+static LogicalResult buildLeafCollections(DeferredExchangePlan &exchange,
+                                          DeferredResultPlan &plan) {
+  unsigned specializationCount = exchange.program.specializationCount;
+  for (auto [leafIndex, leaf] : llvm::enumerate(exchange.program.leaves)) {
+    SmallVector<RequirementFamily *> requirements;
+    unsigned positionCount = 0;
+    for (RequirementFamily &requirement : exchange.requirements)
+      if (requirement.coordinate.leafIndex == leafIndex) {
+        requirements.push_back(&requirement);
+        positionCount = std::max(
+          positionCount, requirement.coordinate.selectedPosition + 1);
+      }
+    if (requirements.empty())
+      return exchange.deferred.emitOpError(
+               "cannot form fragment collection for leaf ")
+             << leafIndex << ": no requirements";
+    auto fragmentType = dyn_cast<RankedTensorType>(
+      requirements.front()->publicationFragmentType);
+    if (!fragmentType || llvm::any_of(requirements,
+          [&](RequirementFamily *item) {
+            return item->publicationFragmentType != fragmentType;
+          }))
+      return exchange.deferred.emitOpError(
+               "cannot form fragment collection for leaf ")
+             << leafIndex
+             << ": publication fragment types differ or are unranked";
+    RankedTensorType normalized = fragmentType;
+    if (leaf.form == DeferredLeafForm::GraphBatchProjection)
+      normalized = leaf.leadingRankReduced
+        ? leaf.reconstructedType
+        : RankedTensorType::get(
+            leaf.reconstructedType.getShape().drop_front(),
+            leaf.reconstructedType.getElementType());
+    bool direct = positionCount == 1 && normalized == leaf.reconstructedType;
+    bool leading = leaf.reconstructedType.getRank() == normalized.getRank() + 1
+                && leaf.reconstructedType.getDimSize(0) == positionCount
+                && leaf.reconstructedType.getShape().drop_front()
+                     == normalized.getShape();
+    if (!direct && !leading)
+      return exchange.deferred.emitOpError(
+               "cannot form fragment collection for leaf ")
+             << leafIndex << ": " << positionCount << " fragments of "
+             << normalized << " do not reconstruct " << leaf.reconstructedType;
+    FragmentCollectionKind kind = specializationCount == 1
+                                ? FragmentCollectionKind::Leaf
+                                : FragmentCollectionKind::GroupedLeaf;
+    SmallVector<int64_t> shape;
+    if (specializationCount > 1)
+      shape.push_back(specializationCount);
+    llvm::append_range(shape, leaf.reconstructedType.getShape());
+    FragmentCollectionPlan collection;
+    collection.key = {&exchange, kind, static_cast<unsigned>(leafIndex)};
+    collection.collectionType = RankedTensorType::get(
+      shape, leaf.reconstructedType.getElementType());
+    collection.positionCount = positionCount;
+    for (RequirementFamily *requirement : requirements)
+      collection.requirements.push_back({
+        requirement,
+        requirement->coordinate.specializationIndex * positionCount
+          + requirement->coordinate.selectedPosition});
+    if (failed(verifyCollectionCoverage(
+          exchange, collection, specializationCount * positionCount)))
+      return failure();
+    plan.collections.push_back(std::move(collection));
+  }
+  return success();
+}
+
+static bool supportsAssemblyTransform(
+    Type publicationType, const DeferredInsertAssemblyEntryTemplate &entry) {
+  auto publication = dyn_cast<RankedTensorType>(publicationType);
+  RankedTensorType source = entry.sourceType;
+  if (!publication || !source)
+    return false;
+  switch (entry.sourceTransform) {
+  case DeferredAssemblySourceTransform::Identity:
+    return publication == source;
+  case DeferredAssemblySourceTransform::AddLeadingUnitDimension:
+    return source.getRank() == publication.getRank() + 1
+        && source.getDimSize(0) == 1
+        && source.getShape().drop_front() == publication.getShape();
+  case DeferredAssemblySourceTransform::RemoveLeadingUnitDimension:
+    return publication.getRank() == source.getRank() + 1
+        && publication.getDimSize(0) == 1
+        && publication.getShape().drop_front() == source.getShape();
+  }
+  llvm_unreachable("unknown deferred assembly source transform");
+}
+
+static LogicalResult buildInsertAssemblyCollection(
+    DeferredExchangePlan &exchange, DeferredResultPlan &plan) {
+  if (exchange.program.specializationCount != 1)
+    return exchange.deferred.emitOpError(
+      "grouped specialization insert assembly is unsupported");
+  const DeferredInsertAssemblyTemplate &assembly =
+    *exchange.program.insertAssembly;
+  FragmentCollectionPlan collection;
+  collection.key = {&exchange, FragmentCollectionKind::InsertAssembly, 0};
+  collection.collectionType = assembly.resultType;
+  collection.positionCount = assembly.entries.size();
+  llvm::DenseSet<RequirementFamily *> collected;
+  for (auto [position, entry] : llvm::enumerate(assembly.entries)) {
+    bool found = false;
+    for (RequirementFamily &requirement : exchange.requirements) {
+      if (!(requirement.coordinate == entry.coordinate))
+        continue;
+      if (!supportsAssemblyTransform(
+            requirement.publicationFragmentType, entry))
+        return exchange.deferred.emitOpError(
+                 "insert assembly source transform does not match publication type at entry ")
+               << position;
+      if (!collected.insert(&requirement).second)
+        return exchange.deferred.emitOpError(
+                 "insert assembly requirement is owned by multiple entries at entry ")
+               << position;
+      collection.requirements.push_back(
+        {&requirement, static_cast<unsigned>(position)});
+      found = true;
+    }
+    if (!found)
+      return exchange.deferred.emitOpError(
+               "insert assembly entry has no requirement at entry ")
+             << position;
+  }
+  if (collected.size() != exchange.requirements.size())
+    return exchange.deferred.emitOpError(
+      "insert assembly does not own every deferred requirement");
+  if (failed(verifyCollectionCoverage(
+        exchange, collection, assembly.entries.size())))
+    return failure();
+  plan.collections.push_back(std::move(collection));
+  return success();
+}
+
+using TemplateGeometryMember = SmallVector<OpFoldResult> DeferredSliceTemplate::*;
+static constexpr std::array<TemplateGeometryMember, 3> templateGeometryMembers{
+  &DeferredSliceTemplate::offsets, &DeferredSliceTemplate::sizes,
+  &DeferredSliceTemplate::strides};
+
+template <typename GetValue>
+static FailureOr<StaticIntGrid> buildGrid(
+    const DeferredProgramTemplate &program, unsigned laneCount,
+    unsigned rowCount, bool specializeRows, GetValue getValue) {
+  SmallVector<StaticIntSequence> rows;
+  for (unsigned row = 0; row < rowCount; ++row) {
+    DeferredLaneValueEvaluator evaluator(
+      program, laneCount, specializeRows ? row : 0);
+    auto sequence = evaluator.evaluate(getValue(row));
+    if (failed(sequence)) return failure();
+    rows.push_back(std::move(*sequence));
+  }
+  return StaticIntGrid::fromRows(rows);
+}
+
+template <typename GetGeometry>
+static FailureOr<DeferredGridSliceGeometry> buildGeometryGrids(
+    const DeferredProgramTemplate &program, unsigned laneCount,
+    unsigned rowCount, bool specializeRows, GetGeometry getGeometry) {
+  DeferredGridSliceGeometry result;
+  const DeferredSliceTemplate &first = getGeometry(0);
+  for (auto [group, sourceMember] : llvm::enumerate(templateGeometryMembers)) {
+    TemplateGeometryMember member = sourceMember;
+    for (unsigned dimension = 0;
+         dimension < (first.*member).size(); ++dimension) {
+      auto grid = buildGrid(program, laneCount, rowCount, specializeRows,
+        [&](unsigned row) { return (getGeometry(row).*member)[dimension]; });
+      if (failed(grid)) return failure();
+      result[group].push_back(std::move(*grid));
+    }
+  }
+  return result;
+}
+static Value cloneResidual(DeferredExchangePlan &exchange, IRMapping &mapping, DeferredEmissionContext &context) {
+  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);
+  }
+  return mapping.lookupOrDefault(exchange.program.yieldedValue);
+}
+static FailureOr<Value> realizeOne(const DeferredResultPlan &plan, Value lane, DeferredEmissionContext &context) {
+  DeferredExchangePlan &exchange = *plan.exchange;
+  if (exchange.program.insertAssembly) {
+    Value collection = context.fragmentCollections.lookup(
+      {&exchange, FragmentCollectionKind::InsertAssembly, 0});
+    return collection && collection.getType() == exchange.program.insertAssembly->resultType
+      ? FailureOr<Value>(collection) : FailureOr<Value>(failure());
+  }
+  IRMapping mapping;
+  if (exchange.program.scheduledLane) mapping.map(exchange.program.scheduledLane, lane);
+  for (auto &[value, grid] : plan.residualValues) {
+    if (mapping.contains(value)) continue;
+    Value selected = grid.emitLookup(
+      context.constants.getIndex(0),
+      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 index = 0; index < exchange.program.leaves.size(); ++index) {
+    Value value = context.fragmentCollections.lookup(
+      {&exchange, FragmentCollectionKind::Leaf, index});
+    if (!value || value.getType() != exchange.program.leaves[index].reconstructedType) {
+      exchange.deferred.emitOpError("failed to reconstruct deferred result leaf ") << index;
+      return failure();
+    }
+    mapping.map(exchange.program.leaves[index].replacementRoot, value);
+  }
+  Value result = cloneResidual(exchange, mapping, context);
+  Type expected = exchange.program.specializationCount > 1 ? Type(exchange.program.specializationFragmentType) : exchange.deferred.getOutput().getType();
+  return result && result.getType() == expected ? FailureOr<Value>(result) : FailureOr<Value>(failure());
+}
+} // namespace
+FailureOr<DeferredResultPlan> buildDeferredResultPlan(DeferredExchangePlan &exchange) {
+  DeferredResultPlan plan;
+  plan.exchange = &exchange;
+  if (failed(exchange.program.insertAssembly
+      ? buildInsertAssemblyCollection(exchange, plan)
+      : buildLeafCollections(exchange, plan))) return failure();
+  unsigned specializations = exchange.program.specializationCount;
+  for (const auto &leaf : exchange.program.leaves) {
+    auto geometry = buildGeometryGrids(exchange.program,
+      exchange.targetLaneCount, specializations, true,
+      [&](unsigned) -> const DeferredSliceTemplate & {
+        return leaf.innerGeometry;
+      });
+    if (failed(geometry)) return failure();
+    plan.innerGeometry.push_back(std::move(*geometry));
+  }
+  if (exchange.program.insertAssembly) {
+    const auto &entries = exchange.program.insertAssembly->entries;
+    auto geometry = buildGeometryGrids(exchange.program,
+      exchange.targetLaneCount, entries.size(), false,
+      [&](unsigned row) -> const DeferredSliceTemplate & {
+        return entries[row].targetGeometry;
+      });
+    if (failed(geometry)) return failure();
+    plan.assemblyGeometry = std::move(*geometry);
+  }
+  llvm::DenseSet<Value> residualValues;
+  for (Operation *op : exchange.program.residualOps)
+    op->walk([&](Operation *nested) {
+      for (Value operand : nested->getOperands())
+        if (operand.getType().isIndex() || isa<IntegerType>(operand.getType()))
+          residualValues.insert(operand);
+    });
+  for (Value value : residualValues) {
+    auto grid = buildGrid(exchange.program, exchange.targetLaneCount,
+      specializations, true,
+      [&](unsigned) { return OpFoldResult(value); });
+    if (succeeded(grid))
+      plan.residualValues.try_emplace(value, std::move(*grid));
+  }
+  return plan;
+}
+FailureOr<Value> realizeDeferredResult(const DeferredResultPlan &plan, Value lane, DeferredEmissionContext &context) {
+  DeferredExchangePlan &exchange = *plan.exchange;
+  if (exchange.program.specializationCount == 1) return realizeOne(plan, lane, context);
+  auto outputType = dyn_cast<RankedTensorType>(exchange.deferred.getOutput().getType());
+  RankedTensorType fragmentType = exchange.program.specializationFragmentType;
+  if (!outputType || !fragmentType || outputType.getRank() != fragmentType.getRank() + 1 || outputType.getDimSize(0) != exchange.program.specializationCount || outputType.getShape().drop_front() != fragmentType.getShape()) return failure();
+  if (exchange.program.insertAssembly) return failure();
+  SmallVector<Value> leafStacks;
+  for (auto [index, leaf] : llvm::enumerate(exchange.program.leaves)) {
+      FragmentCollectionKey key{&exchange, FragmentCollectionKind::GroupedLeaf, static_cast<unsigned>(index)};
+      Value collection = context.fragmentCollections.lookup(key);
+      SmallVector<int64_t> shape{exchange.program.specializationCount};
+      llvm::append_range(shape, leaf.reconstructedType.getShape());
+      if (!collection || collection.getType() != RankedTensorType::get(shape, leaf.reconstructedType.getElementType())) return failure();
+      leafStacks.push_back(collection);
+  }
+  Location loc = exchange.deferred.getLoc();
+  Value initial = tensor::EmptyOp::create(context.rewriter, loc, outputType.getShape(), outputType.getElementType());
+  auto loop = buildNormalizedScfFor(context.rewriter, loc, context.constants.getIndex(0), context.constants.getIndex(exchange.program.specializationCount), context.constants.getIndex(1), ValueRange{initial},
+                                    [&](OpBuilder &, Location, Value specialization, ValueRange iterArgs, SmallVectorImpl<Value> &yielded) -> LogicalResult {
+                                      IRMapping mapping;
+                                      if (exchange.program.scheduledLane) mapping.map(exchange.program.scheduledLane, lane);
+                                      if (exchange.program.specializationArgument) mapping.map(exchange.program.specializationArgument, specialization);
+                                      Value runtimeLane = lane ? lane : context.constants.getIndex(0);
+                                      for (auto &[value, grid] : plan.residualValues) {
+                                          Value selected = grid.emitLookup(specialization, runtimeLane, exchange.deferred, context.constants, context.rewriter, loc);
+                                          if (!value.getType().isIndex()) selected = arith::IndexCastOp::create(context.rewriter, loc, value.getType(), selected);
+                                          mapping.map(value, selected);
+                                      }
+                                      for (auto [index, leaf] : llvm::enumerate(exchange.program.leaves)) { Value selected = extractMixedSliceOrIdentity(context.rewriter, loc, leafStacks[index], leaf.reconstructedType, leadingSlice(context.rewriter, leaf.reconstructedType, specialization)); if (!selected) return failure(); mapping.map(leaf.replacementRoot, selected); }
+                                      Value fragment = cloneResidual(exchange, mapping, context);
+                                      if (!fragment || fragment.getType() != fragmentType) return failure();
+                                      yielded.push_back(insertMixedSlice(context.rewriter, loc, fragment, iterArgs.front(), leadingSlice(context.rewriter, fragmentType, specialization)));
+                                      return success();
+                                    });
+  return failed(loop) || loop->results.size() != 1 ? FailureOr<Value>(failure()) : FailureOr<Value>(loop->results.front());
+}
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,57 @@
+#pragma once
+
+#include "DeferredCommunicationModel.hpp"
+#include "src/Accelerators/PIM/Common/IR/StaticIntGrid.hpp"
+#include <array>
+
+namespace onnx_mlir::spatial {
+
+struct DeferredEmissionContext;
+
+enum class FragmentCollectionKind {
+  Leaf,
+  GroupedLeaf,
+  InsertAssembly
+};
+
+struct FragmentCollectionKey {
+  DeferredExchangePlan *exchange = nullptr;
+  FragmentCollectionKind kind = FragmentCollectionKind::Leaf;
+  unsigned leafIndex = 0;
+
+  bool operator==(const FragmentCollectionKey &other) const {
+    return exchange == other.exchange && kind == other.kind
+        && leafIndex == other.leafIndex;
+  }
+};
+
+struct FragmentCollectionPlan {
+  FragmentCollectionKey key;
+  mlir::RankedTensorType collectionType;
+  unsigned positionCount = 0;
+  struct Requirement {
+    RequirementFamily *family = nullptr;
+    unsigned position = 0;
+  };
+  llvm::SmallVector<Requirement> requirements;
+};
+
+using DeferredGridSliceGeometry =
+  std::array<llvm::SmallVector<StaticIntGrid, 0>, 3>;
+
+struct DeferredResultPlan {
+  DeferredExchangePlan* exchange = nullptr;
+  llvm::SmallVector<FragmentCollectionPlan, 0> collections;
+  llvm::SmallVector<DeferredGridSliceGeometry, 0> innerGeometry;
+  DeferredGridSliceGeometry assemblyGeometry;
+  llvm::DenseMap<mlir::Value, StaticIntGrid> residualValues;
+};
+
+mlir::FailureOr<DeferredResultPlan>
+buildDeferredResultPlan(DeferredExchangePlan& exchange);
+
+mlir::FailureOr<mlir::Value> realizeDeferredResult(const DeferredResultPlan& plan,
+                                                   mlir::Value lane,
+                                                   DeferredEmissionContext& context);
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,475 @@
+#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<unsigned> getStepIndex(
+    ScheduledInfo& info, SpatDeferredCommunicationOp deferred) {
+  Operation *position = deferred;
+  Block *block = info.blocks.front();
+  while (position && position->getBlock() != block)
+    position = position->getParentOp();
+  if (!position)
+    return failure();
+  for (unsigned step : llvm::reverse(
+         llvm::seq<unsigned>(0, info.stepAnchors.size())))
+    if (info.stepAnchors[step] == position
+        || info.stepAnchors[step]->isBeforeInBlock(position))
+      return step;
+  return failure();
+}
+
+static LogicalResult collectScheduledOperations(
+    const ScheduledComputeMaterializationResult &materialization,
+    DeferredTransferPlan &plan) {
+  unsigned nextStream = 0;
+  for (const ScheduledMaterializationRecord &record :
+       materialization.materializedSchedules) {
+    Operation &op = *record.scheduledOp;
+    ScheduledInfo info;
+    info.op = &op;
+    Region& body = isa<SpatScheduledCompute>(op) ? cast<SpatScheduledCompute>(op).getBody()
+                                                 : cast<SpatScheduledComputeBatch>(op).getBody();
+    if (!body.hasOneBlock())
+      return op.emitOpError("phase 2 requires canonical one-block scheduled IR");
+    info.blocks.push_back(&body.front());
+    info.stepCount = record.runs.empty() ? record.stepPlans.size()
+                                        : record.runs.size();
+    info.stepAnchors = record.stepAnchors;
+    if (llvm::any_of(info.stepAnchors,
+                     [](Operation *anchor) { return !anchor; }))
+      return op.emitOpError("phase 2 scheduled step anchor is missing");
+    for (size_t core : record.cpus)
+      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.stepCount;
+  return success();
+}
+
+static LogicalResult collectProducedValues(
+    const ScheduledComputeMaterializationResult &materialization,
+    DeferredTransferPlan &plan) {
+  if (plan.scheduled.size() != materialization.materializedSchedules.size())
+    return failure();
+  for (auto [recordIndex, record] :
+       llvm::enumerate(materialization.materializedSchedules)) {
+    ScheduledInfo &info = plan.scheduled[recordIndex];
+    DenseMap<ComputeInstance, std::pair<unsigned, unsigned>> coordinates;
+    if (record.runs.empty()) {
+      for (auto [step, stepPlan] : llvm::enumerate(record.stepPlans))
+        for (auto [lane, instance] :
+             llvm::enumerate(stepPlan.stepTuple.instances))
+          coordinates[instance] = {step, lane};
+    } else {
+      for (auto [step, run] : llvm::enumerate(record.runs))
+        for (const ComputeInstance &instance : run.instances)
+          coordinates[instance] = {step, 0};
+    }
+
+    for (const MaterializedStepValue &value : record.stepValues) {
+      auto coordinate = coordinates.find(value.instance);
+      if (coordinate == coordinates.end()
+          || coordinate->second.second >= info.cores.size()
+          || value.graphId < 0 || value.laneStart < 0
+          || value.laneCount <= 0 || !value.payload)
+        return info.op->emitOpError(
+          "phase 2 received an invalid scheduled materialization record");
+      unsigned step = coordinate->second.first;
+      unsigned lane = coordinate->second.second;
+      auto produced = std::make_unique<ProducedValue>(ProducedValue {
+        &info, step, value.resultIndex, value.graphId, info.cores[lane],
+        value.laneStart, value.laneCount, lane, value.payloadLaneStart,
+        value.payloadLaneCount, value.payload, value.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;
+}
+
+static LogicalResult buildRequirementFamilies(DeferredTransferPlan& plan,
+                                              DeferredExchangePlan& exchange,
+                                              GraphBatchPublicationCache& publicationCache) {
+  for (unsigned specialization = 0;
+       specialization < exchange.program.specializationCount;
+       ++specialization) {
+    DeferredLaneValueEvaluator evaluator(
+      exchange.program, exchange.targetLaneCount, specialization);
+    for (auto leafItem : llvm::enumerate(exchange.program.leaves)) {
+      unsigned leafIndex = leafItem.index();
+      const DeferredProjectionLeafTemplate &leaf = leafItem.value();
+      auto sourceIndices = evaluator.resolveSourceOperandIndices(leaf.sourceRoot);
+      if (failed(sourceIndices))
+        return failure();
+      DeferredStaticSliceGeometry geometry;
+      auto evaluate = [&](ArrayRef<OpFoldResult> values,
+                          SmallVectorImpl<StaticIntSequence> &target) {
+        for (OpFoldResult value : values) {
+          auto sequence = evaluator.evaluate(value);
+          if (failed(sequence))
+            return failure();
+          target.push_back(std::move(*sequence));
+        }
+        return success();
+      };
+      if (failed(evaluate(leaf.leadingGeometry.offsets, geometry.offsets))
+          || failed(evaluate(leaf.leadingGeometry.sizes, geometry.sizes))
+          || failed(evaluate(leaf.leadingGeometry.strides, geometry.strides)))
+        return failure();
+      struct Source {
+        OpResult value;
+        int64_t graphId;
+        const GraphBatchPublicationMap *publication = nullptr;
+        ProducedValue *scalarProducer = nullptr;
+      };
+      SmallVector<Source> sources;
+      unsigned positionCount = 1;
+      for (unsigned lane = 0; lane < exchange.targetLaneCount; ++lane) {
+        Value value = exchange.deferred.getSources()[sourceIndices->valueAt(lane)];
+        while (auto selection = value.getDefiningOp<SpatDeferredSourceSelectOp>()) {
+          auto selectors = evaluator.evaluate(selection.getSelector());
+          if (failed(selectors))
+            return exchange.deferred.emitOpError(
+              "phase 2 cannot evaluate deferred source selection"), failure();
+          int64_t selected = selectors->valueAt(lane);
+          if (selected < 0 || selected >= static_cast<int64_t>(selection.getSources().size()))
+            return exchange.deferred.emitOpError(
+              "phase 2 deferred source selection is out of range"), failure();
+          value = selection.getSources()[selected];
+        }
+        auto source = dyn_cast<OpResult>(value);
+        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();
+        Source resolved {source, graphId.getInt()};
+        if (auto batch = dyn_cast<SpatGraphComputeBatch>(source.getOwner())) {
+          auto publication = getGraphBatchPublicationMap(
+            batch, source.getResultNumber(), publicationCache);
+          if (failed(publication))
+            return failure();
+          if (leaf.form != DeferredLeafForm::GraphBatchProjection)
+            positionCount = std::max<unsigned>(
+              positionCount, (*publication)->physicalSlotToGraphLane.size());
+        }
+        else if (isa<SpatGraphCompute>(source.getOwner())) {
+          auto producer = findProducer(plan, exchange.deferred, graphId.getInt(),
+                                       source.getResultNumber(), std::nullopt);
+          if (failed(producer))
+            return failure();
+          resolved.scalarProducer = *producer;
+        }
+        sources.push_back(resolved);
+      }
+      for (Source &source : sources)
+        if (auto batch = dyn_cast<SpatGraphComputeBatch>(source.value.getOwner()))
+          source.publication = &publicationCache.find(
+            {batch, source.value.getResultNumber()})->second;
+      if (leaf.form == DeferredLeafForm::GraphBatchProjection)
+        for (unsigned lane = 0; lane < exchange.targetLaneCount; ++lane)
+          positionCount = std::max<unsigned>(
+            positionCount, geometry.sizes.front().valueAt(lane));
+
+      for (unsigned position = 0; position < positionCount; ++position) {
+        SmallVector<ProducedValue *> producers(exchange.targetLaneCount);
+        SmallVector<Type> fragmentTypes(exchange.targetLaneCount);
+        SmallVector<int64_t> graphLanes(exchange.targetLaneCount, -1);
+        SmallVector<int64_t> localOffsets(exchange.targetLaneCount, -1);
+        for (unsigned lane = 0; lane < exchange.targetLaneCount; ++lane) {
+          Source source = sources[lane];
+          if (source.publication) {
+            int64_t slot = position;
+            if (leaf.form == DeferredLeafForm::GraphBatchProjection) {
+              if (position >= geometry.sizes.front().valueAt(lane))
+                continue;
+              slot = geometry.offsets.front().valueAt(lane)
+                   + static_cast<int64_t>(position)
+                       * geometry.strides.front().valueAt(lane);
+            }
+            else if (position >= source.publication->physicalSlotToGraphLane.size())
+              continue;
+            if (slot < 0 || slot >= static_cast<int64_t>(
+                  source.publication->physicalSlotToGraphLane.size()))
+              return exchange.deferred.emitOpError(
+                "projection physical slot is outside publication map"),
+                failure();
+            graphLanes[lane] = source.publication->physicalSlotToGraphLane[slot];
+            fragmentTypes[lane] = source.publication->publicationFragmentType;
+            auto producer = findProducer(
+              plan, exchange.deferred, source.graphId,
+              source.value.getResultNumber(),
+              graphLanes[lane]);
+            if (failed(producer))
+              return failure();
+            producers[lane] = *producer;
+            localOffsets[lane] = graphLanes[lane] - (*producer)->laneStart;
+            if (!(*producer)->scheduled->isBatch())
+              localOffsets[lane] += (*producer)->publishedSlotStart;
+          }
+          else if (position == 0 && source.scalarProducer) {
+            fragmentTypes[lane] = leaf.form == DeferredLeafForm::ScalarProjection
+                                ? Type(leaf.reconstructedType)
+                                : source.value.getType();
+            producers[lane] = source.scalarProducer;
+          }
+        }
+
+        StaticIntSequence graphLaneSequence =
+          StaticIntSequence::fromValues(graphLanes);
+        StaticIntSequence localOffsetSequence =
+          StaticIntSequence::fromValues(localOffsets);
+        for (unsigned begin = 0; begin < exchange.targetLaneCount;) {
+          if (!producers[begin]) {
+            ++begin;
+            continue;
+          }
+          unsigned end = begin + 1;
+          while (end < exchange.targetLaneCount
+                 && producers[end] == producers[begin]
+                 && fragmentTypes[end] == fragmentTypes[begin])
+            ++end;
+          RequirementFamily family;
+          family.exchange = &exchange;
+          family.coordinate = {specialization, leafIndex, position};
+          family.targetLanes = LaneSet::range(begin, end);
+          family.producer = producers[begin];
+          family.publicationFragmentType = fragmentTypes[begin];
+          if (graphLanes[begin] >= 0) {
+            family.graphLanes = graphLaneSequence.slice(begin, end - begin);
+            family.producerLocalOffsets =
+              localOffsetSequence.slice(begin, end - begin);
+          }
+          else if (leaf.form == DeferredLeafForm::ScalarProjection) {
+            family.producerProjection.emplace();
+            auto slice = [&](ArrayRef<StaticIntSequence> source,
+                             SmallVectorImpl<StaticIntSequence> &target) {
+              for (const StaticIntSequence &sequence : source)
+                target.push_back(sequence.slice(begin, end - begin));
+            };
+            slice(geometry.offsets, family.producerProjection->offsets);
+            slice(geometry.sizes, family.producerProjection->sizes);
+            slice(geometry.strides, family.producerProjection->strides);
+          }
+          exchange.requirements.push_back(std::move(family));
+          begin = end;
+        }
+      }
+    }
+  }
+  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, deferred)
+                       : 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();
+  bool escapesScheduledGraph = llvm::any_of(
+    blueprint.getOutput().getUses(), [](OpOperand &use) {
+      return !isa<SpatGraphCompute, SpatGraphComputeBatch,
+                  SpatDeferredCommunicationOp>(use.getOwner());
+    });
+  if (!escapesScheduledGraph)
+    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();
+      if (!(*producer)->published)
+        return blueprint.emitOpError(
+          "phase 2 Blueprint source has no scheduled publication"), 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,
+    const ScheduledComputeMaterializationResult &materialization) {
+  DeferredTransferPlan plan;
+  if (failed(collectScheduledOperations(materialization, plan))
+      || failed(collectProducedValues(materialization, 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,24 @@
+#pragma once
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+
+#include "DeferredCommunicationModel.hpp"
+#include "ScheduledComputeMaterialization.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,
+                          const ScheduledComputeMaterializationResult &materialization);
+
+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,398 +1,115 @@
-#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 "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 "ScheduledComputeMaterialization.hpp"
+#include "ScheduledComputeReport.hpp"
+#include "ScheduledComputeVerification.hpp"
 #include "SpatialDataflowCsvExporter.hpp"
-#include "Scheduling/ComputeGraph.hpp"
-#include "Scheduling/ComputeInstanceUtils.hpp"
+#include "DeferredCommunicationRealization.hpp"
+
+#include "mlir/Pass/Pass.h"
+
 #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();
-    spatial::SpatialDataflowExportStage exportMode = spatial::getSpatialDataflowExportStage();
-    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;
+    }
+    // Phase 1 is intentionally dumped before its verifier: malformed deferred
+    // payloads must be diagnosed from the producer-owned body.
+    dumpModule(moduleOp, "spatial3_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(verifyDeferredTransferPhase1Invariants(funcOp))) {
+      moduleOp.emitError("scheduled Spatial deferred communication verification failed");
+      signalPassFailure();
+      return;
+    }
+    if (failed(verifyScheduledMaterializationRecords(materialization->materializedSchedules))) {
+      moduleOp.emitError("scheduled Spatial materialization record verification failed");
+      signalPassFailure();
+      return;
+    }
+    if (failed(verifyScheduledSpatialInvariants(funcOp))) {
+      moduleOp.emitError("scheduled Spatial phase 1 verification failed");
       signalPassFailure();
       return;
     }
 
-    if (!sortTopologically(&func.getBody().front())) {
-      func.emitOpError("failed to topologically order merged Spatial IR");
+    SpatialDataflowExportStage exportMode = getSpatialDataflowExportStage();
+    if (shouldExportSpatialDataflowStage(exportMode, SpatialDataflowExportStage::Spatial3)
+        && failed(exportSpatialDataflowCsvScheduled(
+          funcOp, materialization->materializedSchedules,
+          "spatial3_scheduled_no_comm", "spatial3"))) {
       signalPassFailure();
       return;
     }
-    if (failed(verifyScheduledSpatialInvariants(func))) {
-      func.emitOpError("scheduled Spatial verification failed after merge materialization");
+
+    dumpScheduledComputeReport(moduleOp,
+                               funcOp,
+                               schedule,
+                               materialization->peftClassPlans,
+                               materialization->materializedSchedules);
+    if (failed(realizeDeferredCommunication(funcOp, *materialization))) {
+      moduleOp.emitError("MergeComputeNodes phase 2 communication realization failed");
       signalPassFailure();
       return;
     }
-    if (spatial::shouldExportSpatialDataflowStage(exportMode, spatial::SpatialDataflowExportStage::Post)
-        && failed(spatial::exportSpatialDataflowCsvPost(func))) {
+    dumpModule(moduleOp, "spatial4_scheduled", /*assumeVerified=*/true);
+    if (failed(verifyScheduledResultsLive(materialization->materializedSchedules))
+        || failed(verifyScheduledSpatialInvariants(funcOp))) {
+      moduleOp.emitError("scheduled Spatial phase 2 verification failed");
       signalPassFailure();
       return;
     }
-    dumpModule(cast<ModuleOp>(func->getParentOp()), "spatial2_merged");
-    generateReport(func, "spatial_merge_report", analysisResult->cpuToLastComputeMap.size());
+    if (shouldExportSpatialDataflowStage(exportMode, SpatialDataflowExportStage::Spatial4)
+        && failed(exportSpatialDataflowCsvScheduled(
+          funcOp, materialization->materializedSchedules,
+          "spatial4_scheduled", "spatial4"))) {
+      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,812 @@
+#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/IR/TensorSliceUtils.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+
+#include <map>
+
+#include "llvm/ADT/SmallPtrSet.h"
+
+namespace onnx_mlir {
+namespace spatial {
+
+using namespace mlir;
+namespace {
+
+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 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 void appendComputeBlockArguments(SmallVectorImpl<Type> &argTypes,
+                                        SmallVectorImpl<Location> &argLocs,
+                                        ValueRange weights,
+                                        ValueRange inputs,
+                                        Location loc) {
+  for (Value weight : weights)
+    argTypes.push_back(weight.getType());
+  for (Value input : inputs)
+    argTypes.push_back(input.getType());
+  argLocs.append(argTypes.size(), loc);
+}
+
+static Block *createScheduledComputeBlock(PatternRewriter &rewriter,
+                                          SpatScheduledCompute scheduled,
+                                          Location loc) {
+  SmallVector<Type> argTypes;
+  SmallVector<Location> argLocs;
+  appendComputeBlockArguments(argTypes, argLocs, scheduled.getWeights(),
+                              scheduled.getInputs(), loc);
+  return rewriter.createBlock(&scheduled.getBody(), scheduled.getBody().end(), TypeRange(argTypes), argLocs);
+}
+
+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 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 materializeResultfulBatchRun(
+    PatternRewriter &rewriter, SpatComputeBatch batch,
+    const ScheduledInstanceRun &run, ValueRange scheduledWeights,
+    ValueRange scheduledInputs, Block &block,
+    const MergeScheduleResult &schedule,
+    const DenseMap<ProducerValueKey, MaterializedProducerRef> &availableValues,
+    SmallVectorImpl<Value> &yieldedValues) {
+  const ComputeInstance &first = run.instances.front();
+  const ComputeInstance &last = run.instances.back();
+  uint32_t runLaneCount = last.laneStart + last.laneCount - first.laneStart;
+  ComputeInstance runInstance {batch.getOperation(), first.laneStart,
+                               runLaneCount};
+  SmallVector<Value> initResults;
+  SmallVector<BatchFragmentSpec> fragmentSpecs;
+  for (unsigned resultIndex = 0; resultIndex < batch.getNumResults(); ++resultIndex) {
+    auto spec = getBatchFragmentSpec(batch, resultIndex, runLaneCount);
+    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(), first.laneStart);
+  Value upper = getOrCreateIndexConstant(rewriter, batch.getOperation(), first.laneStart + runLaneCount);
+  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 = runLaneCount == 1
+        ? getOrCreateIndexConstant(rewriter, batch.getOperation(), 0)
+        : arith::SubIOp::create(
+            builder, bodyLoc, originalLane,
+            getOrCreateIndexConstant(
+              rewriter, batch.getOperation(), first.laneStart));
+      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),
+                                     runInstance,
+                                     schedule,
+                                     scheduledInputs,
+                                     block,
+                                     scheduledWeights.size(),
+                                     availableValues,
+                                     *batch.getLaneArgument(),
+                                     originalLane,
+                                     plan)))
+          return failure();
+        plan.scalarizedLocalLane = localLane;
+        plan.scalarizedGraphLaneBase = lower;
+        plan.scalarizedLaneCount = runLaneCount;
+        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;
+  Block *block = createScheduledComputeBlock(
+    rewriter, scheduled, instances.front().op->getLoc());
+  DenseMap<ProducerValueKey, MaterializedProducerRef> availableValues;
+  DenseMap<ProducerValueKey, unsigned> publicationIndices;
+  for (const ScheduledPublication &publication : record.publications)
+    publicationIndices[publication.producer] = publication.scheduledResultIndex;
+  SmallVector<Value> publishedOutputs(scheduled.getNumResults());
+  record.runs = buildScheduledInstanceRuns(instances);
+  for (const ScheduledInstanceRun &run : record.runs) {
+    const ComputeInstance &instance = run.instances.front();
+    for (const ComputeInstance &member : run.instances) {
+      GraphComputeBlockKey key = getGraphComputeBlockKey(member);
+      graphComputeToBlockMap[key] = block;
+      record.computeKeys.push_back(key);
+      record.blocks.push_back(block);
+    }
+
+    Operation *previous = block->empty() ? nullptr : &block->back();
+    rewriter.setInsertionPointToEnd(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(),
+                                     availableValues,
+                                     {},
+                                     {},
+                                     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(materializeResultfulBatchRun(rewriter, batch, run,
+                                              scheduled.getWeights(),
+                                              scheduled.getInputs(), *block,
+                                              schedule, availableValues,
+                                              yieldedValues)))
+        return failure();
+    }
+
+    if (yieldedValues.size()
+        != getComputeInstanceResultValueCount(instance))
+      return scheduled.emitOpError(
+        "scheduled scalar step produced an unexpected result count");
+    for (auto [resultIndex, value] : llvm::enumerate(yieldedValues)) {
+      int64_t runLaneStart = instance.laneStart;
+      const ComputeInstance &last = run.instances.back();
+      int64_t runLaneCount = last.laneStart + last.laneCount - runLaneStart;
+      availableValues[{{instance.op, static_cast<uint32_t>(runLaneStart),
+                        static_cast<uint32_t>(runLaneCount)}, resultIndex}] =
+        {value, 0, runLaneCount};
+      for (const ComputeInstance &member : run.instances) {
+        ProducerValueKey key {member, resultIndex};
+        int64_t payloadLaneStart = member.laneStart - runLaneStart;
+        availableValues[key] = {value, payloadLaneStart, member.laneCount};
+        Value published;
+        if (auto it = publicationIndices.find(key);
+            it != publicationIndices.end()) {
+          Value publication = value;
+          if (run.instances.size() != 1) {
+            auto spec = getBatchFragmentSpec(
+              cast<SpatComputeBatch>(member.op), resultIndex,
+              member.laneCount);
+            if (failed(spec))
+              return failure();
+            MixedSliceGeometry geometry;
+            auto payloadType = cast<RankedTensorType>(value.getType());
+            geometry.offsets.assign(payloadType.getRank(),
+                                    rewriter.getIndexAttr(0));
+            geometry.offsets.front() =
+              rewriter.getIndexAttr(payloadLaneStart);
+            geometry.sizes.reserve(payloadType.getRank());
+            for (int64_t dimension : spec->resultType.getShape())
+              geometry.sizes.push_back(rewriter.getIndexAttr(dimension));
+            geometry.strides.assign(payloadType.getRank(),
+                                    rewriter.getIndexAttr(1));
+            publication = extractMixedSliceOrIdentity(
+              rewriter, member.op->getLoc(), value, spec->resultType,
+              geometry);
+          }
+          publishedOutputs[it->second] = publication;
+          published = scheduled.getResult(it->second);
+        }
+        auto graphId = member.op->getAttrOfType<IntegerAttr>(
+          "scheduled.graph_id");
+        if (!graphId)
+          return member.op->emitOpError(
+            "scheduled materialization requires graph identity metadata");
+        record.stepValues.push_back(
+          {member, static_cast<unsigned>(resultIndex), value,
+           graphId.getInt(), member.laneStart, member.laneCount,
+           payloadLaneStart, member.laneCount, published});
+      }
+    }
+    Operation *anchor = previous ? previous->getNextNode() : &block->front();
+    record.stepAnchors.push_back(anchor);
+  }
+  if (llvm::any_of(publishedOutputs, [](Value value) { return !value; }))
+    return scheduled.emitOpError(
+      "scheduled scalar materialization did not produce every declared publication");
+  rewriter.setInsertionPointToEnd(block);
+  SpatYieldOp::create(rewriter, scheduled.getLoc(), publishedOutputs);
+  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,
+                                                                   int64_t destinationRank,
+                                                                   Operation *constantAnchor) {
+  SmallVector<OpFoldResult> offsets;
+  Value scheduledOutputLane = scheduledLane;
+  if (lanesPerScheduledLane != 1) {
+    scheduledOutputLane = affineMulConst(
+      builder, loc, scheduledLane, lanesPerScheduledLane, constantAnchor);
+  }
+  offsets.push_back(scheduledOutputLane);
+  offsets.append(destinationRank - 1, OpFoldResult(builder.getIndexAttr(0)));
+  return offsets;
+}
+
+struct BatchPublication {
+  Value fragment;
+  BlockArgument destination;
+  SmallVector<OpFoldResult> offsets;
+  SmallVector<OpFoldResult> sizes;
+  SmallVector<OpFoldResult> strides;
+};
+
+static LogicalResult collectMultiCpuPublications(
+  PatternRewriter &rewriter,
+  SpatScheduledComputeBatch scheduled,
+  const ComputeStepTuple &stepTuple,
+  const ComputeInstance &representative,
+  ArrayRef<Value> localFragments,
+  const DenseMap<ProducerValueKey, unsigned> &publicationIndices,
+  Block &block,
+  ScheduledMaterializationRecord &record,
+  SmallVectorImpl<BatchPublication> &publications) {
+  Value scheduledLane = block.getArgument(0);
+  for (auto [resultIndex, localFragment] : llvm::enumerate(localFragments)) {
+    auto publicationIt = publicationIndices.find({representative, resultIndex});
+    Value published = publicationIt == publicationIndices.end()
+      ? Value()
+      : scheduled.getResult(publicationIt->second);
+    for (auto [lane, instance] : llvm::enumerate(stepTuple.instances)) {
+      auto instancePublication = publicationIndices.find({instance, resultIndex});
+      if ((instancePublication == publicationIndices.end()) != !published
+          || (instancePublication != publicationIndices.end()
+              && instancePublication->second != publicationIt->second))
+        return scheduled.emitOpError("scheduled batch tuple has inconsistent publication ownership");
+      auto graphId = instance.op->getAttrOfType<IntegerAttr>("scheduled.graph_id");
+      if (!graphId)
+        return instance.op->emitOpError("scheduled materialization requires graph identity metadata");
+      record.stepValues.push_back({instance, static_cast<unsigned>(resultIndex),
+                                   localFragment, graphId.getInt(),
+                                   instance.laneStart, instance.laneCount,
+                                   static_cast<int64_t>(lane) * instance.laneCount,
+                                   instance.laneCount, published});
+    }
+    if (!published)
+      continue;
+
+    auto localType = cast<RankedTensorType>(localFragment.getType());
+    auto destination = block.getArgument(
+      getScheduledBatchResultArgBase(scheduled) + publicationIt->second);
+    auto destinationType = cast<RankedTensorType>(destination.getType());
+    if (destinationType.getRank() != localType.getRank()
+        && destinationType.getRank() != localType.getRank() + 1)
+      return scheduled.emitOpError(
+        "scheduled publication source must match or rank-reduce into its destination");
+
+    int64_t lanesPerScheduledLane = isa<SpatCompute>(representative.op)
+      ? 1 : representative.laneCount;
+    SmallVector<OpFoldResult> offsets = buildScheduledOutputInsertOffsets(
+      rewriter, scheduled.getLoc(), scheduledLane, lanesPerScheduledLane,
+      destinationType.getRank(), scheduled.getOperation());
+    SmallVector<OpFoldResult> sizes;
+    SmallVector<OpFoldResult> strides;
+    if (destinationType.getRank() == localType.getRank() + 1) {
+      sizes.push_back(rewriter.getIndexAttr(1));
+      strides.push_back(rewriter.getIndexAttr(1));
+    }
+    for (int64_t dim : localType.getShape()) {
+      sizes.push_back(rewriter.getIndexAttr(dim));
+      strides.push_back(rewriter.getIndexAttr(1));
+    }
+    publications.push_back({localFragment, destination, std::move(offsets),
+                            std::move(sizes), std::move(strides)});
+  }
+  return success();
+}
+
+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;
+  DenseMap<ProducerValueKey, unsigned> publicationIndices;
+  for (const ScheduledPublication &publication : record.publications)
+    publicationIndices[publication.producer] = publication.scheduledResultIndex;
+  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);
+  SmallVector<BatchPublication> publications;
+  for (const ScheduledStepPlan &stepPlan : stepPlans) {
+    const ComputeStepTuple &stepTuple = stepPlan.stepTuple;
+    SourceLaneSelector sourceLaneSelector =
+      buildSourceLaneSelector(rewriter, stepTuple, scheduled.getOperation(), laneStartTableCache);
+    for (const ComputeInstance &instance : stepTuple.instances) {
+      GraphComputeBlockKey key = getGraphComputeBlockKey(instance);
+      graphComputeToBlockMap[key] = block;
+      record.computeKeys.push_back(key);
+      record.blocks.push_back(block);
+    }
+
+    Operation *previous = block->empty() ? nullptr : &block->back();
+    rewriter.setInsertionPointToEnd(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");
+        finalLocalFragments.push_back(yielded);
+      }
+    } 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());
+    }
+
+    if (failed(collectMultiCpuPublications(
+          rewriter, scheduled, stepTuple, representative,
+          finalLocalFragments, publicationIndices, *block, record,
+          publications)))
+      return failure();
+    Operation *anchor = previous ? previous->getNextNode()
+                                 : (block->empty() ? nullptr : &block->front());
+    if (!anchor)
+      return scheduled.emitOpError(
+        "scheduled batch step did not materialize a body operation");
+    record.stepAnchors.push_back(anchor);
+  }
+  rewriter.setInsertionPointToEnd(block);
+  if (publications.empty()) {
+    SpatYieldOp::create(rewriter, scheduled.getLoc(), ValueRange {});
+    return success();
+  }
+  auto inParallel = SpatInParallelOp::create(rewriter, scheduled.getLoc());
+  rewriter.setInsertionPointToEnd(&inParallel.getRegion().front());
+  for (const BatchPublication &publication : publications)
+    tensor::ParallelInsertSliceOp::create(
+      rewriter, scheduled.getLoc(), publication.fragment,
+      publication.destination, 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);
+    record.publications = peftClassPlan.publications;
+
+    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("scheduled.realized", rewriter.getBoolAttr(true));
+      scheduled->setAttr(kCoreIdAttrName, rewriter.getI32IntegerAttr(static_cast<int32_t>(peftClassPlan.cpus.front())));
+      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("scheduled.realized", rewriter.getBoolAttr(true));
+      scheduled->setAttr(kCoreIdsAttrName, rewriter.getDenseI32ArrayAttr(toI32Array(peftClassPlan.cpus)));
+      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,37 @@
+#pragma once
+
+#include "ScheduledComputePlan.hpp"
+
+namespace onnx_mlir {
+namespace spatial {
+
+struct BatchFragmentSpec {
+  RankedTensorType resultType;
+  RankedTensorType sourceSliceType;
+};
+
+struct ScheduledComputeMaterializationResult {
+  llvm::MapVector<size_t, PeftClassPlan> peftClassPlans;
+  std::vector<ScheduledMaterializationRecord> materializedSchedules;
+  DenseMap<GraphComputeBlockKey, Block *> graphComputeToBlockMap;
+};
+
+FailureOr<BatchFragmentSpec>
+getBatchFragmentSpec(SpatComputeBatch batch,
+                     unsigned resultIndex,
+                     uint32_t fragmentLaneCount);
+
+LogicalResult verifyPeftClassPlan(Operation *diagnosticAnchor,
+                                  const PeftClassPlan &peftClassPlan,
+                                  const MergeScheduleResult &schedule);
+
+LogicalResult collectPeftClassOperandsAndResults(
+  PeftClassPlan &peftClassPlan, const MergeScheduleResult &schedule);
+
+FailureOr<ScheduledComputeMaterializationResult>
+materializeScheduledCompute(func::FuncOp funcOp,
+                            const MergeScheduleResult &schedule,
+                            PatternRewriter &rewriter);
+
+} // namespace spatial
+} // namespace onnx_mlir