Unexpected invariant now it's clear (batched in the first tensor rank)

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`

invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md

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

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

@@ -346,6 +346,52 @@ inline void createParallelInsertSliceIntoBatchOutput(mlir::PatternRewriter& rewr
   mlir::tensor::ParallelInsertSliceOp::create(rewriter, loc, source, dest, offsets, sizes, strides);
 }
 
+inline void publishGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
+                                              mlir::Location loc,
+                                              mlir::Value fragment,
+                                              mlir::Value output,
+                                              mlir::Value physicalSlot) {
+  auto fragmentType = mlir::cast<mlir::RankedTensorType>(fragment.getType());
+  mlir::SmallVector<mlir::OpFoldResult> offsets {physicalSlot};
+  mlir::SmallVector<mlir::OpFoldResult> sizes {rewriter.getIndexAttr(1)};
+  mlir::SmallVector<mlir::OpFoldResult> strides {rewriter.getIndexAttr(1)};
+  for (int64_t dim : fragmentType.getShape()) {
+    offsets.push_back(rewriter.getIndexAttr(0));
+    sizes.push_back(rewriter.getIndexAttr(dim));
+    strides.push_back(rewriter.getIndexAttr(1));
+  }
+  createParallelInsertSliceIntoBatchOutput(rewriter, loc, fragment, output, offsets, sizes, strides);
+}
+
+inline mlir::FailureOr<mlir::Value>
+extractGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
+                                  mlir::Location loc,
+                                  mlir::Value physicalBatch,
+                                  mlir::OpFoldResult slot,
+                                  mlir::RankedTensorType fragmentType) {
+  if (fragmentType.getRank() == 0)
+    return mlir::failure();
+  auto physicalType = mlir::dyn_cast<mlir::RankedTensorType>(physicalBatch.getType());
+  if (!physicalType || physicalType.getRank() != fragmentType.getRank() + 1)
+    return mlir::failure();
+  mlir::SmallVector<int64_t> selectedShape {1};
+  llvm::append_range(selectedShape, fragmentType.getShape());
+  auto selectedType = mlir::RankedTensorType::get(selectedShape, fragmentType.getElementType(), fragmentType.getEncoding());
+  mlir::SmallVector<mlir::OpFoldResult> offsets {slot};
+  mlir::SmallVector<mlir::OpFoldResult> sizes {rewriter.getIndexAttr(1)};
+  mlir::SmallVector<mlir::OpFoldResult> strides {rewriter.getIndexAttr(1)};
+  for (int64_t dim : fragmentType.getShape()) {
+    offsets.push_back(rewriter.getIndexAttr(0));
+    sizes.push_back(rewriter.getIndexAttr(dim));
+    strides.push_back(rewriter.getIndexAttr(1));
+  }
+  mlir::Value selected = mlir::tensor::ExtractSliceOp::create(rewriter, loc, selectedType, physicalBatch, offsets, sizes, strides);
+  mlir::SmallVector<mlir::ReassociationIndices> reassociation {{0, 1}};
+  for (int64_t dim = 2; dim <= fragmentType.getRank(); ++dim)
+    reassociation.push_back({dim});
+  return mlir::tensor::CollapseShapeOp::create(rewriter, loc, fragmentType, selected, reassociation).getResult();
+}
+
 template <typename BodyFn>
 mlir::Value materializeOrComputeUnary(mlir::Value input,
                                       mlir::RankedTensorType resultType,

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

@@ -19,9 +19,7 @@ RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) {
 }
 
 RankedTensorType getRowStripStorageType(RankedTensorType logicalType) {
-  return RankedTensorType::get({logicalType.getDimSize(2), logicalType.getDimSize(1), 1, logicalType.getDimSize(3)},
-                               logicalType.getElementType(),
-                               logicalType.getEncoding());
+  return spatial::getGraphBatchPhysicalResultType(logicalType.getDimSize(2), getRowStripFragmentType(logicalType));
 }
 
 std::pair<SmallVector<int64_t>, SmallVector<int64_t>> buildRowStripMetadata(RankedTensorType type) {
@@ -39,29 +37,12 @@ std::pair<SmallVector<int64_t>, SmallVector<int64_t>> buildRowStripMetadata(Rank
   return {offsets, sizes};
 }
 
-SmallVector<OpFoldResult> buildRowStripFragmentOffsets(PatternRewriter& rewriter, OpFoldResult row) {
-  return {row, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-}
-
-SmallVector<OpFoldResult> buildRowStripFragmentSizes(PatternRewriter& rewriter, RankedTensorType logicalType) {
-  return {rewriter.getIndexAttr(1),
-          rewriter.getIndexAttr(logicalType.getDimSize(1)),
-          rewriter.getIndexAttr(1),
-          rewriter.getIndexAttr(logicalType.getDimSize(3))};
-}
-
 Value extractRowStripFragment(Value storage,
                               RankedTensorType logicalType,
                               OpFoldResult row,
                               PatternRewriter& rewriter,
                               Location loc) {
-  return tensor::ExtractSliceOp::create(rewriter,
-                                        loc,
-                                        getRowStripFragmentType(logicalType),
-                                        storage,
-                                        buildRowStripFragmentOffsets(rewriter, row),
-                                        buildRowStripFragmentSizes(rewriter, logicalType),
-                                        getUnitStrides(rewriter, 4));
+  return *extractGraphBatchPhysicalFragment(rewriter, loc, storage, row, getRowStripFragmentType(logicalType));
 }
 
 void insertRowStripFragment(Value fragment,
@@ -70,13 +51,11 @@ void insertRowStripFragment(Value fragment,
                             OpFoldResult row,
                             PatternRewriter& rewriter,
                             Location loc) {
-  createParallelInsertSliceIntoBatchOutput(rewriter,
-                                           loc,
-                                           fragment,
-                                           output,
-                                           buildRowStripFragmentOffsets(rewriter, row),
-                                           buildRowStripFragmentSizes(rewriter, logicalType),
-                                           getUnitStrides(rewriter, 4));
+  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,
@@ -145,30 +124,23 @@ FailureOr<Value> createRowStripStorageFromRows(Value rows,
 }
 
 FailureOr<Value>
-materializeRowStripStorageToDense(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
+createRowStripAssemblyBlueprint(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
   auto storageType = dyn_cast<RankedTensorType>(storage.getType());
   if (!storageType || storageType != getRowStripStorageType(logicalType))
     return failure();
 
-  auto batchOp = createSpatComputeBatch(
-    rewriter, loc, TypeRange {logicalType}, logicalType.getDimSize(2), {}, ValueRange {storage},
-    [&](detail::SpatComputeBatchBodyArgs args) {
-      Value fragment = extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
-      createParallelInsertSliceIntoBatchOutput(rewriter,
-                                               loc,
-                                               fragment,
-                                               args.outputs.front(),
-                                               SmallVector<OpFoldResult> {rewriter.getIndexAttr(0),
-                                                                          rewriter.getIndexAttr(0),
-                                                                          args.lane,
-                                                                          rewriter.getIndexAttr(0)},
-                                               buildRowStripFragmentSizes(rewriter, logicalType),
-                                               getUnitStrides(rewriter, 4));
-      return success();
-    });
-  if (failed(batchOp))
-    return failure();
-  return batchOp->getResult(0);
+  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>

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

@@ -50,7 +50,7 @@ mlir::FailureOr<mlir::Value> createRowStripStorageFromRows(mlir::Value rows,
                                                            mlir::PatternRewriter& rewriter,
                                                            mlir::Location loc);
 
-mlir::FailureOr<mlir::Value> materializeRowStripStorageToDense(mlir::Value storage,
+mlir::FailureOr<mlir::Value> createRowStripAssemblyBlueprint(mlir::Value storage,
                                                                mlir::RankedTensorType logicalType,
                                                                mlir::PatternRewriter& rewriter,
                                                                mlir::Location loc);

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -75,7 +75,7 @@ materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location
   auto [expectedOffsets, expectedSizes] = buildRowStripMetadata(rowStripValue.logicalType);
   if (!llvm::equal(rowStripValue.fragmentOffsets, expectedOffsets) || !llvm::equal(rowStripValue.fragmentSizes, expectedSizes))
     return failure();
-  return materializeRowStripStorageToDense(rowStripValue.storage, rowStripValue.logicalType, rewriter, loc);
+  return createRowStripAssemblyBlueprint(rowStripValue.storage, rowStripValue.logicalType, rewriter, loc);
 }
 
 struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
@@ -277,6 +277,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;
@@ -366,11 +368,15 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
     moduleOp.walk([&](Operation* op) {
       if (isa<ONNXEntryPointOp>(op))
         return;
-      if (isa<spatial::SpatConv2DPlanOp,
-              spatial::SpatBiasAddPlanOp,
-              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;

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

@@ -1448,12 +1448,10 @@ static FailureOr<Value> reconstructDepthwiseGemmRows(Value pieces,
             Value laneIndex = createOrFoldAffineApply(
               rewriter, tileLoc, (d0 * tiling.totalPatches) + d1, ValueRange {channelTileIndex, patchIndex}, anchorOp);
             auto rowTileType = RankedTensorType::get({1, tiling.tileOutputChannels}, piecesType.getElementType());
-            SmallVector<OpFoldResult> pieceOffsets {laneIndex, rewriter.getIndexAttr(0)};
-            SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(1),
-                                                  rewriter.getIndexAttr(tiling.tileOutputChannels)};
-            Value rowTile = tensor::ExtractSliceOp::create(
-              rewriter, tileLoc, rowTileType, piecesArg, pieceOffsets, pieceSizes, getUnitStrides(rewriter, 2));
-            Value rowNext = insertOutputTile(rowTile, rowAcc, channelTileIndex, tiling, rewriter, tileLoc);
+            FailureOr<Value> rowTile = extractGraphBatchPhysicalFragment(rewriter, tileLoc, piecesArg, laneIndex, rowTileType);
+            if (failed(rowTile))
+              return failure();
+            Value rowNext = insertOutputTile(*rowTile, rowAcc, channelTileIndex, tiling, rewriter, tileLoc);
             tileYielded.push_back(rowNext);
             return success();
           });
@@ -1556,8 +1554,9 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
 
   auto gemmOutType =
     RankedTensorType::get({tiling->totalPatches, state.outType.getDimSize(1)}, state.outType.getElementType());
-  auto piecesType = RankedTensorType::get({tiling->totalPatches * tiling->numChannelTiles, tiling->tileOutputChannels},
-                                          state.outType.getElementType());
+  auto rowTileType = RankedTensorType::get({1, tiling->tileOutputChannels}, state.outType.getElementType());
+  auto piecesType = spatial::getGraphBatchPhysicalResultType(
+    tiling->totalPatches * tiling->numChannelTiles, rowTileType);
   auto paddedInputType = cast<RankedTensorType>(paddedInput.getType());
   auto inputTileType =
     RankedTensorType::get({1, tiling->channelsPerTile, state.wType.getDimSize(2), state.wType.getDimSize(3)},
@@ -1626,7 +1625,6 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
                                             *tiling,
                                             rewriter,
                                             loc);
-      auto rowTileType = RankedTensorType::get({1, tiling->tileOutputChannels}, state.outType.getElementType());
       Value rowTile = spatial::SpatVMMOp::create(rewriter, loc, rowTileType, weightTile, inputTile).getResult();
       if (args.inputs.size() > 1) {
         Value biasArg = pickInputByRank(/*rank=*/2);
@@ -1637,11 +1635,7 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
         Value biasTile = tiling->numChannelTiles == 1 ? biasArg : createBiasTile(biasArg, channelTileIndex, *tiling, rewriter, loc);
         rowTile = spatial::SpatVAddOp::create(rewriter, loc, rowTileType, rowTile, biasTile).getResult();
       }
-      SmallVector<OpFoldResult> outputOffsets {args.lane, rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> outputSizes {rewriter.getIndexAttr(1),
-                                             rewriter.getIndexAttr(tiling->tileOutputChannels)};
-      createParallelInsertSliceIntoBatchOutput(
-        rewriter, loc, rowTile, args.outputs.front(), outputOffsets, outputSizes, getUnitStrides(rewriter, 2));
+      publishGraphBatchPhysicalFragment(rewriter, loc, rowTile, args.outputs.front(), args.lane);
       return success();
     });
   if (failed(batchOp))
@@ -1650,14 +1644,10 @@ rewriteConv(Operation* convOp, const ConvLoweringState& state, PatternRewriter&
   auto nhwcType = RankedTensorType::get(
     {state.xType.getDimSize(0), state.outType.getDimSize(2), state.outType.getDimSize(3), state.outType.getDimSize(1)},
     state.outType.getElementType());
-  Value collectedRows = batchOp->getResult(0);
-  if (tiling->numChannelTiles != 1) {
-    auto reconstructedRows =
-      reconstructDepthwiseGemmRows(batchOp->getResult(0), piecesType, gemmOutType, *tiling, rewriter, loc);
-    if (failed(reconstructedRows))
-      return failure();
-    collectedRows = *reconstructedRows;
-  }
+  auto reconstructedRows = reconstructDepthwiseGemmRows(batchOp->getResult(0), piecesType, gemmOutType, *tiling, rewriter, loc);
+  if (failed(reconstructedRows))
+    return failure();
+  Value collectedRows = *reconstructedRows;
 
   return createCollectedConvOutput(ValueRange {collectedRows},
                                    state.outType,

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

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

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

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

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

@@ -139,9 +139,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 +189,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 +204,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,14 +218,10 @@ 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();
@@ -276,10 +266,11 @@ static Value buildKeepdimsFromLanePackedBatch(Value batchValue,
   if (!pendingLeadingReducedAxes.empty())
     expandCompactToKeepdims.back().append(pendingLeadingReducedAxes.begin(), pendingLeadingReducedAxes.end());
 
+  if (batchType.getNumElements() != batchType.getDimSize(0))
+    return {};
   auto reshapeCompute =
     createSpatCompute<1>(rewriter, loc, TypeRange {keepdimsType}, {}, ValueRange {batchValue}, [&](Value input) {
-      auto flatType =
-        RankedTensorType::get({batchType.getDimSize(0)}, batchType.getElementType(), batchType.getEncoding());
+      auto flatType = RankedTensorType::get({batchType.getNumElements()}, batchType.getElementType(), batchType.getEncoding());
       Value flat = tensor::CollapseShapeOp::create(rewriter, loc, flatType, input, collapseToFlat);
       Value compact = flat;
       if (compactKeptType != flatType)

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -134,6 +134,7 @@ static spatial::SpatBlueprintOp insertRowStripBlueprint(IRRewriter& rewriter, Va
                                               nullptr,
                                               nullptr,
                                               nullptr,
+                                              nullptr,
                                               nullptr);
 }
 

src/PIM/Dialect/Spatial/CMakeLists.txt

@@ -10,6 +10,8 @@ add_pim_library(SpatialOps
   Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp
   Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp
   Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp
+  Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp
+  Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp
   Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp
   Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp
   Transforms/MergeComputeNodes/ScheduledComputeReport.cpp

src/PIM/Dialect/Spatial/Spatial.td

@@ -183,23 +183,10 @@ def SpatBlockYieldOp : SpatOp<"block_yield", [
 }
 
 def SpatDeferredCommunicationOp : SpatOp<"deferred_communication", [SingleBlock]> {
-  let summary = "Temporary scheduled communication placeholder";
+  let summary = "Temporary scheduled payload derivation placeholder";
 
   let arguments = (ins
-    Variadic<SpatTensor>:$sources,
-    StrAttr:$exchangeId,
-    StrAttr:$logicalProducer,
-    StrAttr:$logicalConsumer,
-    I64Attr:$sourceClass,
-    I64Attr:$targetClass,
-    I64Attr:$sourceCore,
-    I64Attr:$targetCore,
-    I64Attr:$sourceLane,
-    I64Attr:$targetLane,
-    StrAttr:$transferKind,
-    I64Attr:$resultIndex,
-    DenseI64ArrayAttr:$projectedTransfer,
-    I64Attr:$hostOutputOwner
+    Variadic<SpatTensor>:$sources
   );
 
   let results = (outs
@@ -312,6 +299,7 @@ def SpatBlueprintOp : SpatOp<"blueprint", []> {
     StrAttr:$indexMap,
     OptionalAttr<StrAttr>:$mode,
     OptionalAttr<DenseI64ArrayAttr>:$fragmentOperandIndices,
+    OptionalAttr<DenseI64ArrayAttr>:$fragmentSourceSlots,
     OptionalAttr<DenseI64ArrayAttr>:$fragmentSourceOffsets,
     OptionalAttr<DenseI64ArrayAttr>:$fragmentStrides,
     OptionalAttr<StrAttr>:$conflictPolicy,

src/PIM/Dialect/Spatial/SpatialOps.cpp

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

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

@@ -576,6 +576,10 @@ void SpatBlueprintOp::print(OpAsmPrinter& printer) {
     printer << " operandIndices ";
     printCompressedIntegerList(printer, *operandIndices);
   }
+  if (std::optional<ArrayRef<int64_t>> sourceSlots = getFragmentSourceSlots()) {
+    printer << " sourceSlots ";
+    printCompressedIntegerList(printer, *sourceSlots);
+  }
   if (std::optional<ArrayRef<int64_t>> sourceOffsets = getFragmentSourceOffsets()) {
     printer << " sourceOffsets ";
     printCompressedIntegerList(printer, *sourceOffsets);
@@ -597,6 +601,7 @@ void SpatBlueprintOp::print(OpAsmPrinter& printer) {
                                  getIndexMapAttrName().getValue(),
                                  getModeAttrName().getValue(),
                                  getFragmentOperandIndicesAttrName().getValue(),
+                                 getFragmentSourceSlotsAttrName().getValue(),
                                  getFragmentSourceOffsetsAttrName().getValue(),
                                  getFragmentStridesAttrName().getValue(),
                                  getConflictPolicyAttrName().getValue(),
@@ -620,6 +625,7 @@ ParseResult SpatBlueprintOp::parse(OpAsmParser& parser, OperationState& result)
   SmallVector<int64_t> fragmentOffsets;
   SmallVector<int64_t> fragmentSizes;
   SmallVector<int64_t> fragmentOperandIndices;
+  SmallVector<int64_t> fragmentSourceSlots;
   SmallVector<int64_t> fragmentSourceOffsets;
   SmallVector<int64_t> fragmentStrides;
 
@@ -637,6 +643,9 @@ ParseResult SpatBlueprintOp::parse(OpAsmParser& parser, OperationState& result)
   if (succeeded(parser.parseOptionalKeyword("operandIndices"))
       && parseCompressedIntegerList(parser, fragmentOperandIndices))
     return failure();
+  if (succeeded(parser.parseOptionalKeyword("sourceSlots"))
+      && parseCompressedIntegerList(parser, fragmentSourceSlots))
+    return failure();
   if (succeeded(parser.parseOptionalKeyword("sourceOffsets"))
       && parseCompressedIntegerList(parser, fragmentSourceOffsets))
     return failure();
@@ -667,6 +676,8 @@ ParseResult SpatBlueprintOp::parse(OpAsmParser& parser, OperationState& result)
     result.addAttribute("mode", mode);
   if (!fragmentOperandIndices.empty())
     result.addAttribute("fragmentOperandIndices", builder.getDenseI64ArrayAttr(fragmentOperandIndices));
+  if (!fragmentSourceSlots.empty())
+    result.addAttribute("fragmentSourceSlots", builder.getDenseI64ArrayAttr(fragmentSourceSlots));
   if (!fragmentSourceOffsets.empty())
     result.addAttribute("fragmentSourceOffsets", builder.getDenseI64ArrayAttr(fragmentSourceOffsets));
   if (!fragmentStrides.empty())

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -574,20 +574,26 @@ LogicalResult SpatBlueprintOp::verify() {
 
   auto stridesAttr = getFragmentStridesAttr();
   auto operandIndicesAttr = getFragmentOperandIndicesAttr();
+  auto sourceSlotsAttr = getFragmentSourceSlotsAttr();
   auto sourceOffsetsAttr = getFragmentSourceOffsetsAttr();
   if (!operandIndicesAttr)
     return emitError("fragment assembly blueprint requires fragment operand indices");
+  if (!sourceSlotsAttr)
+    return emitError("fragment assembly blueprint requires physical fragment source slots");
   if (!sourceOffsetsAttr)
     return emitError("fragment assembly blueprint requires fragment source offsets");
   if (!stridesAttr)
     return emitError("fragment assembly blueprint requires fragment strides");
   ArrayRef<int64_t> operandIndices = operandIndicesAttr.asArrayRef();
+  ArrayRef<int64_t> sourceSlots = sourceSlotsAttr.asArrayRef();
   ArrayRef<int64_t> sourceOffsets = sourceOffsetsAttr.asArrayRef();
   ArrayRef<int64_t> strides = stridesAttr.asArrayRef();
   if (strides.size() != offsets.size())
     return emitError("fragment stride and offset arrays must have the same length");
   if (sourceOffsets.size() != operandIndices.size())
     return emitError("fragment source offset count must match fragment operand index count");
+  if (sourceSlots.size() != operandIndices.size())
+    return emitError("fragment source slot count must match fragment operand index count");
   if (!getConflictPolicyAttr() || !getCoveragePolicyAttr())
     return emitError("fragment assembly blueprint requires conflict and coverage policies");
   if (getConflictPolicy() != "disjoint")
@@ -622,14 +628,19 @@ LogicalResult SpatBlueprintOp::verify() {
     int64_t operandIndex = operandIndices[fragmentIndex];
     if (operandIndex < 0 || operandIndex >= operandCount)
       return emitError("fragment assembly operand index is out of range");
+    if (sourceSlots[fragmentIndex] < 0)
+      return emitError("fragment assembly physical source slot must be nonnegative");
     if (sourceOffsets[fragmentIndex] < 0)
       return emitError("fragment assembly source offsets must be nonnegative");
 
     auto operandType = dyn_cast<RankedTensorType>(operands[operandIndex].getType());
     if (!operandType || !operandType.hasStaticShape())
       return emitError("fragment assembly blueprint requires static ranked tensor operands");
-    if (operandType.getRank() != rank)
-      return emitError("fragment assembly blueprint requires operand/result rank match");
+    if (operandType.getRank() != rank + 1)
+      return emitError("fragment assembly physical operand must have one leading source-slot dimension");
+    if (sourceSlots[fragmentIndex] >= operandType.getDimSize(0))
+      return emitError("fragment assembly physical source slot is out of range");
+    auto fragmentType = RankedTensorType::get(operandType.getShape().drop_front(), operandType.getElementType(), operandType.getEncoding());
 
     SmallVector<int64_t, 4> fragmentOffsets;
     SmallVector<int64_t, 4> fragmentSizes;
@@ -645,12 +656,12 @@ LogicalResult SpatBlueprintOp::verify() {
     int64_t fragmentElements = 1;
     for (int64_t dim = 0; dim < rank; ++dim)
       fragmentElements *= fragmentSizes[dim];
-    if (sourceOffsets[fragmentIndex] + fragmentElements > operandType.getNumElements())
-      return emitError("fragment assembly source offset exceeds the operand bounds");
+    if (sourceOffsets[fragmentIndex] + fragmentElements > fragmentType.getNumElements())
+      return emitError("fragment assembly source offset exceeds the selected physical fragment bounds");
     SmallVector<int64_t, 4> sourceSliceOffsets =
-      expandFlatElementIndex(sourceOffsets[fragmentIndex], operandType.getShape());
+      expandFlatElementIndex(sourceOffsets[fragmentIndex], fragmentType.getShape());
     for (int64_t dim = 0; dim < rank; ++dim)
-      if (sourceSliceOffsets[dim] + fragmentSizes[dim] > operandType.getDimSize(dim))
+      if (sourceSliceOffsets[dim] + fragmentSizes[dim] > fragmentType.getDimSize(dim))
         return emitError("fragment assembly source offset must describe a valid unit-stride slice");
 
     for (const auto& [existingOffsets, existingSizes] : slices) {
@@ -746,7 +757,8 @@ LogicalResult verifyComputeLikeOp(ComputeOpTy compute, StringRef opName) {
 
     Operation* terminator = block.getTerminator();
     if (auto yieldOp = dyn_cast_or_null<SpatYieldOp>(terminator)) {
-      if (isScheduled)
+      auto realized = compute->template getAttrOfType<BoolAttr>("scheduled.realized");
+      if (isScheduled && (!realized || !realized.getValue() || !compute.getBody().hasOneBlock()))
         return compute.emitOpError("scheduled compute blocks must terminate with spat.block_yield");
       llvm::append_range(yieldedTypes, yieldOp->getOperandTypes());
       continue;
@@ -820,6 +832,16 @@ LogicalResult SpatBlockYieldOp::verify() {
 LogicalResult SpatDeferredCommunicationOp::verify() {
   if (getSources().empty())
     return emitOpError("requires at least one source");
+  static constexpr StringLiteral staleAttributes[] = {
+    "exchangeId", "logicalProducer", "logicalConsumer", "sourceClass", "targetClass", "sourceCore",
+    "targetCore", "sourceLane", "targetLane", "transferKind", "resultIndex", "projectedTransfer",
+    "hostOutputOwner", "source_cpus", "source_classes", "source_lane_ranges", "target_cpus",
+    "target_classes", "target_lane_ranges", "batched", "source_operand_for_scheduled_lane",
+    "multi_source_payload"};
+  for (StringLiteral name : staleAttributes)
+    if (getOperation()->hasAttr(name))
+      return emitOpError() << "does not accept stale routing attribute '" << name
+                           << "'; source selection and shaping belong in the body and routing is derived in Phase 2";
   if (failed(verifyRegionArguments(getOperation(), getBody(), getSources(), "spat.deferred_communication")))
     return failure();
   return verifyYieldTypes(getOperation(), getBody(), getOperation()->getResultTypes(), "spat.deferred_communication");

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

@@ -0,0 +1,274 @@
+#include "DeferredCommunicationDeadlock.hpp"
+
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/DenseSet.h"
+
+namespace onnx_mlir::spatial {
+
+using namespace mlir;
+
+namespace {
+
+enum class EventKind { Compute, Send, Receive };
+
+struct Event {
+  EventKind kind = EventKind::Compute;
+  uint64_t exchangeId = 0;
+};
+
+static LogicalResult simulate(Operation *anchor,
+                              ArrayRef<SmallVector<Event>> streams,
+                              StringRef phase) {
+  SmallVector<size_t> cursor(streams.size());
+  while (true) {
+    bool allFinished = true;
+    bool progressed = false;
+    for (unsigned stream = 0; stream < streams.size(); ++stream) {
+      if (cursor[stream] == streams[stream].size())
+        continue;
+      allFinished = false;
+      if (streams[stream][cursor[stream]].kind == EventKind::Compute) {
+        ++cursor[stream];
+        progressed = true;
+      }
+    }
+    if (allFinished)
+      return success();
+
+    for (unsigned source = 0; source < streams.size(); ++source) {
+      if (cursor[source] == streams[source].size())
+        continue;
+      const Event &send = streams[source][cursor[source]];
+      if (send.kind != EventKind::Send)
+        continue;
+      for (unsigned target = 0; target < streams.size(); ++target) {
+        if (cursor[target] == streams[target].size())
+          continue;
+        const Event &receive = streams[target][cursor[target]];
+        if (receive.kind == EventKind::Receive && receive.exchangeId == send.exchangeId) {
+          ++cursor[source];
+          ++cursor[target];
+          progressed = true;
+          break;
+        }
+      }
+    }
+    if (!progressed) {
+      InFlightDiagnostic diagnostic = anchor->emitError()
+                                      << phase << " communication rendezvous simulation made no progress";
+      unsigned reported = 0;
+      for (unsigned stream = 0; stream < streams.size() && reported < 8; ++stream) {
+        if (cursor[stream] == streams[stream].size())
+          continue;
+        const Event &event = streams[stream][cursor[stream]];
+        diagnostic << (reported == 0 ? "; blocked " : ", ") << "stream " << stream
+                   << " at exchange " << event.exchangeId;
+        ++reported;
+      }
+      return failure();
+    }
+  }
+}
+
+static std::optional<int64_t> getI64Attr(Operation *op, StringRef name) {
+  if (auto attr = op->getAttrOfType<IntegerAttr>(name))
+    return attr.getInt();
+  return std::nullopt;
+}
+
+} // namespace
+
+LogicalResult verifyPlannedCommunicationDeadlockFree(
+  Operation *anchor,
+  unsigned streamCount,
+  ArrayRef<unsigned> stepCounts,
+  ArrayRef<PlannedCommunicationTransfer> transfers) {
+  if (stepCounts.size() != streamCount)
+    return anchor->emitError("communication plan stream count does not match step counts");
+
+  SmallVector<SmallVector<Event>> streams(streamCount);
+  SmallVector<SmallVector<SmallVector<Event>>> atBoundary(streamCount);
+  for (unsigned stream = 0; stream < streamCount; ++stream)
+    atBoundary[stream].resize(stepCounts[stream] + 1);
+  for (const PlannedCommunicationTransfer &transfer : transfers) {
+    if (transfer.sourceStream >= streamCount || transfer.targetStream >= streamCount
+        || transfer.producerStep >= stepCounts[transfer.sourceStream]
+        || transfer.consumerStep >= stepCounts[transfer.targetStream]
+        || transfer.sourceInsertionStep > stepCounts[transfer.sourceStream]
+        || transfer.targetInsertionStep > stepCounts[transfer.targetStream]
+        || transfer.sourceInsertionStep <= transfer.producerStep
+        || transfer.targetInsertionStep > transfer.consumerStep)
+      return anchor->emitError("communication plan references an invalid stream step");
+    atBoundary[transfer.sourceStream][transfer.sourceInsertionStep].push_back(
+      {EventKind::Send, transfer.exchangeId});
+    atBoundary[transfer.targetStream][transfer.targetInsertionStep].push_back(
+      {EventKind::Receive, transfer.exchangeId});
+  }
+  for (unsigned stream = 0; stream < streamCount; ++stream) {
+    for (unsigned step = 0; step < stepCounts[stream]; ++step) {
+      llvm::append_range(streams[stream], atBoundary[stream][step]);
+      streams[stream].push_back({EventKind::Compute, 0});
+    }
+    llvm::append_range(streams[stream], atBoundary[stream].back());
+  }
+  return simulate(anchor, streams, "planned");
+}
+
+LogicalResult verifyRealizedCommunicationDeadlockFree(func::FuncOp funcOp) {
+  DenseMap<int64_t, SmallVector<Operation *, 2>> operationsByExchange;
+  struct ParentExchange {
+    std::optional<int64_t> expectedTransfers;
+    DenseSet<int64_t> channels;
+  };
+  DenseMap<int64_t, ParentExchange> parentExchanges;
+  DenseMap<int64_t, unsigned> streamByCore;
+  SmallVector<int64_t> cores;
+  bool invalid = false;
+  funcOp.walk([&](Operation *op) {
+    if (!isa<SpatChannelSendOp, SpatChannelReceiveOp>(op))
+      return;
+    std::optional<int64_t> exchangeId = getI64Attr(op, "raptor.exchange_id");
+    if (exchangeId)
+      operationsByExchange[*exchangeId].push_back(op);
+    if (auto channels = op->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_channel_ids"))
+      for (int64_t channel : channels.asArrayRef())
+        operationsByExchange[channel].push_back(op);
+    if (std::optional<int64_t> parent = getI64Attr(op, "raptor.parent_exchange_id")) {
+      ParentExchange &group = parentExchanges[*parent];
+      std::optional<int64_t> expected =
+        getI64Attr(op, "raptor.parent_transfer_count");
+      if (!expected || *expected <= 0
+          || (group.expectedTransfers && group.expectedTransfers != expected)) {
+        op->emitOpError(
+          "realized parent exchange has missing or inconsistent transfer count metadata");
+        invalid = true;
+      } else {
+        group.expectedTransfers = expected;
+      }
+      if (exchangeId)
+        group.channels.insert(*exchangeId);
+      if (auto channels =
+            op->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_channel_ids"))
+        group.channels.insert(channels.asArrayRef().begin(),
+                              channels.asArrayRef().end());
+    }
+    for (StringRef attrName : {"raptor.source_core", "raptor.target_core"})
+      if (std::optional<int64_t> core = getI64Attr(op, attrName); core && !llvm::is_contained(cores, *core))
+        cores.push_back(*core);
+    for (StringRef attrName : {"raptor.batch_source_cores", "raptor.batch_target_cores"})
+      if (auto batchCores = op->getAttrOfType<DenseI64ArrayAttr>(attrName))
+        for (int64_t core : batchCores.asArrayRef())
+          if (!llvm::is_contained(cores, core))
+            cores.push_back(core);
+  });
+  llvm::sort(cores);
+  for (auto [index, core] : llvm::enumerate(cores))
+    streamByCore[core] = index;
+
+  SmallVector<SmallVector<Event>> streams(cores.size());
+  funcOp.walk([&](Operation *op) {
+    if (!isa<SpatChannelSendOp, SpatChannelReceiveOp>(op))
+      return;
+    auto exchangeId = getI64Attr(op, "raptor.exchange_id");
+    if (!exchangeId) {
+      auto channels = op->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_channel_ids");
+      auto sourceCores = op->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_source_cores");
+      auto targetCores = op->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_target_cores");
+      if (!channels || !sourceCores || !targetCores
+          || channels.size() != sourceCores.size() || channels.size() != targetCores.size()) {
+        op->emitOpError("realized compact batch communication is missing channel/core metadata");
+        invalid = true;
+        return;
+      }
+      if (isa<SpatChannelSendOp>(op))
+        for (auto [channel, sourceCore] : llvm::zip(channels.asArrayRef(), sourceCores.asArrayRef()))
+          streams[streamByCore.lookup(sourceCore)].push_back(
+            {EventKind::Send, static_cast<uint64_t>(channel)});
+      else
+        for (auto [channel, targetCore] : llvm::zip(channels.asArrayRef(), targetCores.asArrayRef()))
+          streams[streamByCore.lookup(targetCore)].push_back(
+            {EventKind::Receive, static_cast<uint64_t>(channel)});
+      return;
+    }
+    auto sourceCore = getI64Attr(op, "raptor.source_core");
+    auto targetCore = getI64Attr(op, "raptor.target_core");
+    if (!sourceCore || !targetCore) {
+      op->emitOpError("realized communication is missing core metadata");
+      invalid = true;
+      return;
+    }
+    if (isa<SpatChannelSendOp>(op))
+      streams[streamByCore.lookup(*sourceCore)].push_back({EventKind::Send, static_cast<uint64_t>(*exchangeId)});
+    else if (isa<SpatChannelReceiveOp>(op))
+      streams[streamByCore.lookup(*targetCore)].push_back({EventKind::Receive, static_cast<uint64_t>(*exchangeId)});
+  });
+  if (invalid)
+    return failure();
+  for (const auto &entry : parentExchanges)
+    if (!entry.second.expectedTransfers
+        || entry.second.channels.size()
+             != static_cast<size_t>(*entry.second.expectedTransfers))
+      return funcOp.emitOpError()
+             << "parent exchange " << entry.first
+             << " does not contain its declared lane transfer set";
+
+  for (const auto &entry : operationsByExchange) {
+    if (entry.second.size() != 2 || !isa<SpatChannelSendOp>(entry.second[0])
+        == !isa<SpatChannelSendOp>(entry.second[1]))
+      return funcOp.emitOpError() << "exchange " << entry.first << " does not have exactly one send and one receive";
+    auto send = dyn_cast<SpatChannelSendOp>(entry.second[0]);
+    auto receive = dyn_cast<SpatChannelReceiveOp>(entry.second[1]);
+    if (!send) {
+      send = cast<SpatChannelSendOp>(entry.second[1]);
+      receive = cast<SpatChannelReceiveOp>(entry.second[0]);
+    }
+    if (send.getInput().getType() != receive.getOutput().getType())
+      return send.emitOpError("send and receive payload types do not match");
+    int64_t sendSource = 0;
+    int64_t sendTarget = 0;
+    if (auto channels = send->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_channel_ids")) {
+      auto channelIt = llvm::find(channels.asArrayRef(), entry.first);
+      auto sources = send->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_source_cores");
+      auto targets = send->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_target_cores");
+      if (channelIt == channels.asArrayRef().end() || !sources || !targets)
+        return send.emitOpError("batch send channel metadata is incomplete");
+      size_t index = std::distance(channels.asArrayRef().begin(), channelIt);
+      sendSource = sources.asArrayRef()[index];
+      sendTarget = targets.asArrayRef()[index];
+    } else {
+      auto source = getI64Attr(send, "raptor.source_core");
+      auto target = getI64Attr(send, "raptor.target_core");
+      if (!source || !target)
+        return send.emitOpError("send core metadata is incomplete");
+      sendSource = *source;
+      sendTarget = *target;
+    }
+    int64_t receiveSource = 0;
+    int64_t receiveTarget = 0;
+    if (auto channels = receive->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_channel_ids")) {
+      auto channelIt = llvm::find(channels.asArrayRef(), entry.first);
+      auto sources = receive->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_source_cores");
+      auto targets = receive->getAttrOfType<DenseI64ArrayAttr>("raptor.batch_target_cores");
+      if (channelIt == channels.asArrayRef().end() || !sources || !targets)
+        return receive.emitOpError("batch receive channel metadata is incomplete");
+      size_t index = std::distance(channels.asArrayRef().begin(), channelIt);
+      receiveSource = sources.asArrayRef()[index];
+      receiveTarget = targets.asArrayRef()[index];
+    } else {
+      auto source = getI64Attr(receive, "raptor.source_core");
+      auto target = getI64Attr(receive, "raptor.target_core");
+      if (!source || !target)
+        return receive.emitOpError("receive core metadata is incomplete");
+      receiveSource = *source;
+      receiveTarget = *target;
+    }
+    if (receiveSource != sendSource || receiveTarget != sendTarget)
+      return receive.emitOpError("receive core metadata does not match its send");
+  }
+  return simulate(funcOp, streams, "realized");
+}
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,27 @@
+#pragma once
+
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Support/LLVM.h"
+
+namespace onnx_mlir::spatial {
+
+struct PlannedCommunicationTransfer {
+  uint64_t exchangeId = 0;
+  uint64_t parentExchangeId = 0;
+  unsigned sourceStream = 0;
+  unsigned targetStream = 0;
+  unsigned producerStep = 0;
+  unsigned consumerStep = 0;
+  unsigned sourceInsertionStep = 0;
+  unsigned targetInsertionStep = 0;
+};
+
+mlir::LogicalResult verifyPlannedCommunicationDeadlockFree(
+  mlir::Operation *anchor,
+  unsigned streamCount,
+  mlir::ArrayRef<unsigned> stepCounts,
+  mlir::ArrayRef<PlannedCommunicationTransfer> transfers);
+
+mlir::LogicalResult verifyRealizedCommunicationDeadlockFree(mlir::func::FuncOp funcOp);
+
+} // namespace onnx_mlir::spatial

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

@@ -1,21 +1,16 @@
 #include "DeferredCommunicationPlanning.hpp"
 
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/IRMapping.h"
 
-namespace onnx_mlir {
-namespace spatial {
+#include "llvm/ADT/SmallPtrSet.h"
 
+namespace onnx_mlir::spatial {
 using namespace mlir;
-
 namespace {
 
-struct LaneProducerInfo {
-  ProducerValueRef producer;
-  Value originalSource;
-  DeferredEndpoint producerEndpoint;
-  DeferredEndpoint consumerEndpoint;
-};
-
 static Value getBlockOperand(Block &block, ValueRange operands, Value value, unsigned firstArgument = 0) {
   auto it = llvm::find(operands, value);
   assert(it != operands.end() && "missing scheduled operand");
@@ -29,398 +24,226 @@ static FailureOr<Value> getOriginalProducerValue(const ProducerValueRef &produce
   return outputs[producer.resultIndex];
 }
 
-static DeferredEndpoint getDeferredEndpoint(const ComputeInstance &instance,
-                                            const MergeScheduleResult &schedule) {
-  size_t cpu = getScheduledCpuForComputeInstance(instance, schedule);
-  return DeferredEndpoint {
-    instance,
-    cpu,
-    getCanonicalPeftClassId(cpu, schedule),
-    instance.laneStart,
-    instance.laneCount
-  };
+static bool isSupportedDeferredShapingOp(Operation *op) {
+  return isa<tensor::ExtractSliceOp, tensor::InsertSliceOp, tensor::CollapseShapeOp,
+             tensor::ExpandShapeOp, tensor::CastOp, tensor::EmptyOp, tensor::ExtractOp,
+             arith::ConstantOp, arith::IndexCastOp, arith::AddIOp, arith::SubIOp,
+             arith::MulIOp, affine::AffineApplyOp>(op);
 }
 
-static void setDeferredTransferMetadata(SpatDeferredCommunicationOp transfer,
-                                        OpBuilder &builder,
-                                        const DeferredTransferMetadata &transferMetadata) {
-  SmallVector<int64_t> sourceCpus;
-  SmallVector<int64_t> sourceClasses;
-  SmallVector<int64_t> sourceLaneRanges;
-  SmallVector<int64_t> targetCpus;
-  SmallVector<int64_t> targetClasses;
-  SmallVector<int64_t> targetLaneRanges;
-  for (const DeferredEndpoint &endpoint : transferMetadata.producers) {
-    sourceCpus.push_back(static_cast<int64_t>(endpoint.cpu));
-    sourceClasses.push_back(static_cast<int64_t>(endpoint.canonicalClassId));
-    sourceLaneRanges.push_back(static_cast<int64_t>(endpoint.laneStart));
-    sourceLaneRanges.push_back(static_cast<int64_t>(endpoint.laneCount));
-  }
-  for (const DeferredEndpoint &endpoint : transferMetadata.consumers) {
-    targetCpus.push_back(static_cast<int64_t>(endpoint.cpu));
-    targetClasses.push_back(static_cast<int64_t>(endpoint.canonicalClassId));
-    targetLaneRanges.push_back(static_cast<int64_t>(endpoint.laneStart));
-    targetLaneRanges.push_back(static_cast<int64_t>(endpoint.laneCount));
-  }
-  transfer->setAttr("source_cpus", builder.getDenseI64ArrayAttr(sourceCpus));
-  transfer->setAttr("source_classes", builder.getDenseI64ArrayAttr(sourceClasses));
-  transfer->setAttr("source_lane_ranges", builder.getDenseI64ArrayAttr(sourceLaneRanges));
-  transfer->setAttr("target_cpus", builder.getDenseI64ArrayAttr(targetCpus));
-  transfer->setAttr("target_classes", builder.getDenseI64ArrayAttr(targetClasses));
-  transfer->setAttr("target_lane_ranges", builder.getDenseI64ArrayAttr(targetLaneRanges));
-  // `batched=true` means this deferred communication belongs to a
-  // spat.scheduled_compute_batch block and the array attrs are indexed by
-  // scheduled lane. The legacy scalar attrs on the op keep representative values
-  // for compatibility; Phase 2 must consume the arrays when batched=true.
-  transfer->setAttr("batched", builder.getBoolAttr(transferMetadata.isBatched));
-  if (transferMetadata.isBatched) {
-    transfer->setAttr("source_operand_for_scheduled_lane",
-                      builder.getDenseI64ArrayAttr(transferMetadata.sourceOperandForScheduledLane));
-    transfer->setAttr("multi_source_payload", builder.getBoolAttr(transferMetadata.multiSourcePayload));
-  }
-}
-
-static FailureOr<Value> projectDeferredPayload(Value selectedSource,
-                                               Value consumerInput,
-                                               OpBuilder &builder,
-                                               Location loc) {
-  if (selectedSource.getType() == consumerInput.getType())
-    return selectedSource;
-
-  auto slice = consumerInput.getDefiningOp<tensor::ExtractSliceOp>();
-  if (!slice || slice.getSource().getType() != selectedSource.getType())
-    return failure();
-
-  IRMapping mapper;
-  mapper.map(slice.getSource(), selectedSource);
-  return cast<tensor::ExtractSliceOp>(builder.clone(*slice, mapper)).getResult();
-}
-
-static FailureOr<Value> buildSelectedDeferredSource(OpBuilder &builder,
-                                                    Location loc,
+static FailureOr<Value> buildSelectedDeferredSource(OpBuilder &builder, Location loc,
                                                     SpatDeferredCommunicationOp transfer,
                                                     Value scheduledLane,
                                                     ValueRange sourceBlockArgs,
-                                                    DenseI64ArrayAttr sourceOperandForScheduledLaneAttr) {
-  SmallVector<int64_t> sourceOperandForScheduledLane(sourceOperandForScheduledLaneAttr.asArrayRef().begin(),
-                                                     sourceOperandForScheduledLaneAttr.asArrayRef().end());
-  Value sourceOperandIndexTable =
-    createI64LookupTableConstant(builder, transfer.getOperation(), sourceOperandForScheduledLane);
-  Value sourceIndexI64 =
-    tensor::ExtractOp::create(builder, loc, sourceOperandIndexTable, ValueRange {scheduledLane}).getResult();
-  Value sourceIndex = arith::IndexCastOp::create(builder, loc, builder.getIndexType(), sourceIndexI64).getResult();
-
-  Type elementType = sourceBlockArgs.front().getType();
-  auto sourceTensorType = dyn_cast<RankedTensorType>(elementType);
-  if (!sourceTensorType || !sourceTensorType.hasStaticShape())
-    return transfer.emitOpError(
-             "scheduled batch deferred communication with multiple original graph SSA sources requires compatible static ranked tensor sources"),
-           failure();
-
-  for (Value source : sourceBlockArgs) {
-    auto rankedSourceType = dyn_cast<RankedTensorType>(source.getType());
-    if (!rankedSourceType || !rankedSourceType.hasStaticShape() || rankedSourceType != sourceTensorType)
-      return transfer.emitOpError(
-               "scheduled batch deferred communication with multiple original graph SSA sources requires compatible static ranked tensor sources"),
-             failure();
+                                                    ArrayRef<int64_t> sourceOperandForScheduledLane) {
+  if (sourceBlockArgs.size() == 1)
+    return sourceBlockArgs.front();
+  if (!scheduledLane || sourceOperandForScheduledLane.empty())
+    return transfer.emitOpError("multiple deferred sources require the enclosing scheduled lane"), failure();
+  Value table = createI64LookupTableConstant(builder, transfer.getOperation(), sourceOperandForScheduledLane);
+  Value i64 = tensor::ExtractOp::create(builder, loc, table, ValueRange {scheduledLane}).getResult();
+  Value index = arith::IndexCastOp::create(builder, loc, builder.getIndexType(), i64).getResult();
+  auto type = dyn_cast<RankedTensorType>(sourceBlockArgs.front().getType());
+  if (!type || !type.hasStaticShape())
+    return transfer.emitOpError("multiple deferred sources require static ranked tensors"), failure();
+  for (Value source : sourceBlockArgs)
+    if (source.getType() != type)
+      return transfer.emitOpError("multiple deferred sources require identical tensor types"), failure();
+  SmallVector<int64_t> shape {static_cast<int64_t>(sourceBlockArgs.size())};
+  llvm::append_range(shape, type.getShape());
+  auto stacked = createEmptyTensorForType(builder, loc, RankedTensorType::get(shape, type.getElementType()));
+  if (failed(stacked))
+    return failure();
+  Value value = *stacked;
+  SmallVector<OpFoldResult> sizes, strides(shape.size(), builder.getIndexAttr(1));
+  sizes.push_back(builder.getIndexAttr(1));
+  for (int64_t dim : type.getShape()) sizes.push_back(builder.getIndexAttr(dim));
+  for (auto [i, source] : llvm::enumerate(sourceBlockArgs)) {
+    SmallVector<int64_t> expandedShape {1}; llvm::append_range(expandedShape, type.getShape());
+    SmallVector<ReassociationIndices> reassociation {{0, 1}};
+    for (int64_t dim = 1; dim < type.getRank(); ++dim) reassociation.push_back({dim + 1});
+    Value expanded = tensor::ExpandShapeOp::create(builder, loc,
+      RankedTensorType::get(expandedShape, type.getElementType()), source, reassociation).getResult();
+    SmallVector<OpFoldResult> offsets(shape.size(), builder.getIndexAttr(0));
+    offsets[0] = builder.getIndexAttr(i);
+    value = tensor::InsertSliceOp::create(builder, loc, expanded, value, offsets, sizes, strides).getResult();
   }
-  if (sourceTensorType.getRank() == 0)
-    return transfer.emitOpError(
-             "scheduled batch deferred communication with multiple original graph SSA sources cannot yet build compact tensor source selection"),
-           failure();
-
-  SmallVector<int64_t> stackedShape {static_cast<int64_t>(sourceBlockArgs.size())};
-  llvm::append_range(stackedShape, sourceTensorType.getShape());
-  RankedTensorType stackedType = RankedTensorType::get(stackedShape, sourceTensorType.getElementType());
-  auto stackedEmpty = createEmptyTensorForType(builder, loc, stackedType);
-  if (failed(stackedEmpty))
-    return transfer.emitOpError(
-             "scheduled batch deferred communication with multiple original graph SSA sources cannot yet build compact tensor source selection"),
-           failure();
-  Value stackedSources = *stackedEmpty;
-
-  SmallVector<OpFoldResult> insertSizes;
-  SmallVector<OpFoldResult> unitStrides(stackedType.getRank(), builder.getIndexAttr(1));
-  insertSizes.push_back(builder.getIndexAttr(1));
-  for (int64_t dim : sourceTensorType.getShape())
-    insertSizes.push_back(builder.getIndexAttr(dim));
-
-  for (auto [sourceIndexValue, source] : llvm::enumerate(sourceBlockArgs)) {
-    SmallVector<int64_t> expandedShape {1};
-    llvm::append_range(expandedShape, sourceTensorType.getShape());
-    SmallVector<ReassociationIndices> reassociation;
-    reassociation.push_back({0, 1});
-    for (int64_t dim = 1; dim < sourceTensorType.getRank(); ++dim)
-      reassociation.push_back({dim + 1});
-    Value expandedSource = tensor::ExpandShapeOp::create(builder,
-                                                         loc,
-                                                         RankedTensorType::get(expandedShape, sourceTensorType.getElementType()),
-                                                         source,
-                                                         reassociation)
-                             .getResult();
-    SmallVector<OpFoldResult> insertOffsets(stackedType.getRank(), builder.getIndexAttr(0));
-    insertOffsets[0] = builder.getIndexAttr(sourceIndexValue);
-    stackedSources = tensor::InsertSliceOp::create(
-                       builder, loc, expandedSource, stackedSources, insertOffsets, insertSizes, unitStrides)
-                       .getResult();
-  }
-
-  SmallVector<OpFoldResult> extractOffsets(stackedType.getRank(), builder.getIndexAttr(0));
-  SmallVector<OpFoldResult> extractSizes;
-  SmallVector<OpFoldResult> extractStrides(stackedType.getRank(), builder.getIndexAttr(1));
-  extractOffsets[0] = sourceIndex;
-  extractSizes.push_back(builder.getIndexAttr(1));
-  for (int64_t dim : sourceTensorType.getShape())
-    extractSizes.push_back(builder.getIndexAttr(dim));
-  SmallVector<int64_t> selectedSliceShape {1};
-  llvm::append_range(selectedSliceShape, sourceTensorType.getShape());
-  RankedTensorType selectedSliceType = RankedTensorType::get(selectedSliceShape, sourceTensorType.getElementType());
-  Value selectedSlice = tensor::ExtractSliceOp::create(
-                          builder,
-                          loc,
-                          selectedSliceType,
-                          stackedSources,
-                          extractOffsets,
-                          extractSizes,
-                          extractStrides)
-                          .getResult();
-
-  SmallVector<ReassociationIndices> collapseReassociation;
-  collapseReassociation.push_back({0, 1});
-  for (int64_t dim = 1; dim < sourceTensorType.getRank(); ++dim)
-    collapseReassociation.push_back({dim + 1});
-  return tensor::CollapseShapeOp::create(builder, loc, sourceTensorType, selectedSlice, collapseReassociation)
-    .getResult();
+  SmallVector<OpFoldResult> offsets(shape.size(), builder.getIndexAttr(0)); offsets[0] = index;
+  SmallVector<int64_t> sliceShape {1}; llvm::append_range(sliceShape, type.getShape());
+  auto slice = tensor::ExtractSliceOp::create(builder, loc,
+    RankedTensorType::get(sliceShape, type.getElementType()), value, offsets, sizes, strides);
+  // extract has a leading unit dimension; remove it without changing the payload.
+  SmallVector<ReassociationIndices> reassociation {{0, 1}};
+  for (int64_t dim = 1; dim < type.getRank(); ++dim) reassociation.push_back({dim + 1});
+  return tensor::CollapseShapeOp::create(builder, loc, type, slice.getResult(), reassociation).getResult();
 }
 
-static LogicalResult finalizeDeferredCommunication(OpBuilder &builder,
-                                                   Location loc,
-                                                   Value consumerInput,
-                                                   SpatDeferredCommunicationOp transfer,
-                                                   Value &result) {
-  Block *block = builder.createBlock(&transfer.getBody(),
-                                     transfer.getBody().end(),
-                                     TypeRange {transfer.getSources().getTypes()},
-                                     SmallVector<Location>(transfer.getSources().size(), loc));
-  builder.setInsertionPointToStart(block);
+static bool isTopLevelShaping(Operation *op, Block &body) {
+  return op->getBlock() == &body && isSupportedDeferredShapingOp(op);
+}
 
-  Value selectedSource = block->getArgument(0);
-  if (auto parentScheduled = dyn_cast<SpatScheduledComputeBatch>(transfer->getParentOp())) {
-    auto sourceOperandForScheduledLaneAttr =
-      transfer->getAttrOfType<DenseI64ArrayAttr>("source_operand_for_scheduled_lane");
-    auto multiSourcePayloadAttr = transfer->getAttrOfType<BoolAttr>("multi_source_payload");
-    if (sourceOperandForScheduledLaneAttr && sourceOperandForScheduledLaneAttr.size() == parentScheduled.getLaneCount()) {
-      if (multiSourcePayloadAttr && multiSourcePayloadAttr.getValue()) {
-        FailureOr<Value> multiSourceSelection = buildSelectedDeferredSource(
-          builder,
-          loc,
-          transfer,
-          transfer->getBlock()->getArgument(0),
-          block->getArguments(),
-          sourceOperandForScheduledLaneAttr);
-        if (failed(multiSourceSelection))
-          return failure();
-        selectedSource = *multiSourceSelection;
-      } else {
-        selectedSource = block->getArgument(sourceOperandForScheduledLaneAttr.asArrayRef().front());
-      }
+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 || !isTopLevelShaping(op, body) || !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) {
+  IRMapping mapping;
+  mapping.map(plan.graphInput, selectedSource);
+  std::function<FailureOr<Value>(Value)> cloneScheduledLane = [&](Value value) -> FailureOr<Value> {
+    if (mapping.contains(value)) return mapping.lookup(value);
+    if (value == plan.scheduledLane) return value;
+    if (isa<BlockArgument>(value))
+      return transfer.emitOpError("phase 1 payload shaping captures an unsupported block argument"), failure();
+    Operation *op = value.getDefiningOp();
+    if (!op || !isSupportedDeferredShapingOp(op))
+      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(plan.scheduledGraphLane);
+      if (failed(mappedLane)) return failure();
+      mapping.map(value, *mappedLane);
+      return *mappedLane;
     }
-  }
-
-  FailureOr<Value> yielded = projectDeferredPayload(selectedSource, consumerInput, builder, loc);
-  if (failed(yielded) || (*yielded).getType() != consumerInput.getType())
-    return transfer.emitOpError("cannot derive deferred communication payload from original graph producer value");
-
-  SpatYieldOp::create(builder, loc, *yielded);
-  result = transfer.getOutput();
-  builder.setInsertionPointAfter(transfer);
-  return success();
+    if (isa<BlockArgument>(value))
+      return transfer.emitOpError("phase 1 payload shaping captures an unsupported block argument"), failure();
+    Operation *op = value.getDefiningOp();
+    if (!op || (!isTopLevelShaping(op, body) && !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 LogicalResult createSingleLaneDeferredCommunication(OpBuilder &builder,
-                                                           Location loc,
-                                                           Value consumerInput,
-                                                           Value originalSource,
-                                                           const ProducerValueRef &producer,
-                                                           const ComputeInstance &consumer,
-                                                           const MergeScheduleResult &schedule,
-                                                           uint64_t exchangeId,
-                                                           Value &result) {
-  DeferredTransferMetadata transferMetadata;
-  transferMetadata.producers.push_back(getDeferredEndpoint(producer.instance, schedule));
-  transferMetadata.consumers.push_back(getDeferredEndpoint(consumer, schedule));
-  auto transfer = SpatDeferredCommunicationOp::create(
-    builder,
-    loc,
-    consumerInput.getType(),
-    ValueRange {originalSource},
-    builder.getStringAttr(("x" + std::to_string(exchangeId)).c_str()),
-    builder.getStringAttr(getInstanceName(producer.instance)),
-    builder.getStringAttr(getInstanceName(consumer)),
-    builder.getI64IntegerAttr(static_cast<int64_t>(transferMetadata.producers.front().canonicalClassId)),
-    builder.getI64IntegerAttr(static_cast<int64_t>(transferMetadata.consumers.front().canonicalClassId)),
-    builder.getI64IntegerAttr(static_cast<int64_t>(transferMetadata.producers.front().cpu)),
-    builder.getI64IntegerAttr(static_cast<int64_t>(transferMetadata.consumers.front().cpu)),
-    builder.getI64IntegerAttr(producer.instance.laneStart),
-    builder.getI64IntegerAttr(consumer.laneStart),
-    builder.getStringAttr(consumerInput.getDefiningOp<tensor::ExtractSliceOp>() ? "projected" : "ordinary"),
-    builder.getI64IntegerAttr(static_cast<int64_t>(producer.resultIndex)),
-    builder.getDenseI64ArrayAttr({}),
-    builder.getI64IntegerAttr(-1));
-  setDeferredTransferMetadata(transfer, builder, transferMetadata);
-  return finalizeDeferredCommunication(builder, loc, consumerInput, transfer, result);
-}
-
-static LogicalResult createScheduledLaneDeferredCommunication(OpBuilder &builder,
-                                                              Location loc,
-                                                              Value consumerInput,
-                                                              ValueRange originalSources,
-                                                              const ProducerValueRef &producer,
-                                                              const DeferredTransferMetadata &transferMetadata,
-                                                              uint64_t exchangeId,
-                                                              Value &result) {
-  assert(transferMetadata.isBatched && "expected scheduled-lane deferred communication metadata");
-  assert(!transferMetadata.producers.empty() && !transferMetadata.consumers.empty() && "expected non-empty endpoints");
-  auto transfer = SpatDeferredCommunicationOp::create(
-    builder,
-    loc,
-    consumerInput.getType(),
-    originalSources,
-    builder.getStringAttr(("x" + std::to_string(exchangeId)).c_str()),
-    builder.getStringAttr(getInstanceName(producer.instance)),
-    builder.getStringAttr(getInstanceName(transferMetadata.consumers.front().instance)),
-    builder.getI64IntegerAttr(static_cast<int64_t>(transferMetadata.producers.front().canonicalClassId)),
-    builder.getI64IntegerAttr(static_cast<int64_t>(transferMetadata.consumers.front().canonicalClassId)),
-    builder.getI64IntegerAttr(static_cast<int64_t>(transferMetadata.producers.front().cpu)),
-    builder.getI64IntegerAttr(static_cast<int64_t>(transferMetadata.consumers.front().cpu)),
-    builder.getI64IntegerAttr(producer.instance.laneStart),
-    builder.getI64IntegerAttr(transferMetadata.consumers.front().laneStart),
-    builder.getStringAttr(consumerInput.getDefiningOp<tensor::ExtractSliceOp>() ? "projected" : "ordinary"),
-    builder.getI64IntegerAttr(static_cast<int64_t>(producer.resultIndex)),
-    builder.getDenseI64ArrayAttr({}),
-    builder.getI64IntegerAttr(-1));
-  setDeferredTransferMetadata(transfer, builder, transferMetadata);
-  return finalizeDeferredCommunication(builder, loc, consumerInput, transfer, result);
+static void collectClosure(Value value, Block &body, const DeferredInputPlan &plan,
+                           llvm::SmallPtrSetImpl<Operation *> &ops) {
+  Operation *op = value.getDefiningOp();
+  if (!op || !isTopLevelShaping(op, body) || !ops.insert(op).second) return;
+  for (Value operand : op->getOperands())
+    if (operand != plan.graphInput && operand != plan.graphLane) collectClosure(operand, body, plan, ops);
 }
 
 } // namespace
 
-LogicalResult mapSingleCpuInput(OpBuilder &builder,
-                                       Location loc,
-                                       Value input,
-                                       const ComputeInstance &consumerInstance,
-                                       const MergeScheduleResult &schedule,
-                                       ValueRange scheduledInputs,
-                                       Block &block,
-                                       unsigned firstInputArgument,
-                                       ArrayRef<ProducerValueKey> carriedKeys,
-                                       uint64_t &exchangeId,
-                                       Value &mapped) {
-  std::optional<ProducerValueRef> producer = getProducerValueRef(input, &consumerInstance);
-  if (!producer) {
-    mapped = getBlockOperand(block, scheduledInputs, input, firstInputArgument);
+LogicalResult prepareSingleCpuInput(OpBuilder &, Location loc, Value input, BlockArgument graphInput,
+                                    const ComputeInstance &consumerInstance, const MergeScheduleResult &,
+                                    ValueRange scheduledInputs, Block &block, unsigned firstInputArgument,
+                                    ArrayRef<ProducerValueKey> carriedKeys, Value graphLane, Value scheduledGraphLane,
+                                    DeferredInputPlan &plan) {
+  plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, {}};
+  auto producer = getProducerValueRef(input, &consumerInstance);
+  if (!producer) { plan.availableValue = getBlockOperand(block, scheduledInputs, input, firstInputArgument); return success(); }
+  ProducerValueKey key {producer->instance, producer->resultIndex};
+  auto carried = llvm::find(carriedKeys, key);
+  if (carried != carriedKeys.end()) {
+    plan.availableValue = block.getArgument(firstInputArgument + scheduledInputs.size() + std::distance(carriedKeys.begin(), carried));
     return success();
   }
-
-  ProducerValueKey producerKey {producer->instance, producer->resultIndex};
-  auto carriedIt = llvm::find(carriedKeys, producerKey);
-  if (carriedIt != carriedKeys.end()) {
-    mapped = block.getArgument(firstInputArgument + scheduledInputs.size() + std::distance(carriedKeys.begin(), carriedIt));
-    return success();
-  }
-
-  FailureOr<Value> originalSource = getOriginalProducerValue(*producer);
-  if (failed(originalSource))
-    return emitError(loc) << "cannot resolve original graph producer value for " << getInstanceName(producer->instance);
-  return createSingleLaneDeferredCommunication(
-    builder, loc, input, *originalSource, *producer, consumerInstance, schedule, exchangeId++, mapped);
+  auto source = getOriginalProducerValue(*producer);
+  if (failed(source)) return emitError(loc) << "cannot resolve original graph producer value";
+  plan.originalSources.push_back(*source);
+  return success();
 }
 
-static FailureOr<unsigned> getComputeInstanceInputIndex(const ComputeInstance &instance, Value input) {
-  auto inputs = getComputeInstanceInputs(instance);
+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) {
+  plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, scheduledLane};
+  const ComputeInstance &representative = tuple.instances.front();
+  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 failure();
-  return static_cast<unsigned>(std::distance(inputs.begin(), it));
+  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 mapMultiCpuTupleInput(OpBuilder &builder,
-                                           Location loc,
-                                           Value input,
-                                           const ComputeStepTuple &stepTuple,
-                                           const PeftClassPlan &peftClassPlan,
-                                           const MergeScheduleResult &schedule,
-                                           ValueRange scheduledInputs,
-                                           Block &block,
-                                           unsigned firstInputArgument,
-                                           uint64_t &exchangeId,
-                                           Value &mapped) {
-  assert(!stepTuple.instances.empty() && "expected non-empty step tuple");
-  const ComputeInstance &representative = stepTuple.instances.front();
-  std::optional<ProducerValueRef> representativeProducer = getProducerValueRef(input, &representative);
-  if (!representativeProducer) {
-    FailureOr<unsigned> inputIndex = getComputeInstanceInputIndex(representative, input);
-    if (failed(inputIndex))
-      return emitError(loc) << "cannot resolve scheduled batch step host input index";
-    for (const ComputeInstance &instance : stepTuple.instances) {
-      auto instanceInputs = getComputeInstanceInputs(instance);
-      if (*inputIndex >= instanceInputs.size())
-        return emitError(loc) << "scheduled batch step host input index out of range";
-      if (instanceInputs[*inputIndex] != input || getProducerValueRef(instanceInputs[*inputIndex], &instance))
-        return emitError(loc) << "scheduled batch step requires identical host inputs across lanes";
+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 (!isTopLevelShaping(user, body)) { 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 isTopLevelShaping(next.getOwner(), body); });
+          bool hasOtherUse = llvm::any_of(result.getUses(), [&](OpOperand &next) { return !isTopLevelShaping(next.getOwner(), body); });
+          if (hasOtherUse) roots.push_back(result);
+          if (hasShapingUse) worklist.push_back(result);
+        }
+      }
+    }
+    if (needsIdentity) roots.push_back(plan.graphInput);
+    llvm::sort(roots, [](Value a, Value b) { return a.getAsOpaquePointer() < b.getAsOpaquePointer(); });
+    roots.erase(std::unique(roots.begin(), roots.end()), roots.end());
+    for (Value root : roots) {
+      auto transfer = SpatDeferredCommunicationOp::create(builder, loc, root.getType(), plan.originalSources);
+      Block *deferred = builder.createBlock(&transfer.getBody(), transfer.getBody().end(),
+        TypeRange {transfer.getSources().getTypes()}, SmallVector<Location>(transfer.getSources().size(), loc));
+      builder.setInsertionPointToStart(deferred);
+      auto selected = buildSelectedDeferredSource(builder, loc, transfer, plan.scheduledLane,
+                                                  deferred->getArguments(), plan.sourceOperandForScheduledLane);
+      if (failed(selected)) return failure();
+      auto payload = clonePayloadRoot(root, body, plan, builder, transfer, *selected);
+      if (failed(payload)) return failure();
+      SpatYieldOp::create(builder, loc, *payload);
+      mapper.map(root, transfer.getOutput());
+      collectClosure(root, body, plan, absorbed);
+      builder.setInsertionPointAfter(transfer);
     }
-    mapped = getBlockOperand(block, scheduledInputs, input, firstInputArgument);
-    return success();
   }
-
-  FailureOr<unsigned> inputIndex = getComputeInstanceInputIndex(representative, input);
-  if (failed(inputIndex))
-    return emitError(loc) << "cannot resolve scheduled batch step input index";
-
-  (void)peftClassPlan;
-  DeferredTransferMetadata transferMetadata;
-  transferMetadata.isBatched = true;
-  SmallVector<LaneProducerInfo> laneProducerInfos;
-  SmallVector<Value> uniqueOriginalSources;
-  SmallVector<int64_t> sourceOperandForScheduledLane;
-  for (const ComputeInstance &instance : stepTuple.instances) {
-    auto instanceInputs = getComputeInstanceInputs(instance);
-    if (*inputIndex >= instanceInputs.size())
-      return emitError(loc) << "scheduled batch step input index out of range";
-    Value laneInput = instanceInputs[*inputIndex];
-    std::optional<ProducerValueRef> laneProducer = getProducerValueRef(laneInput, &instance);
-    if (!laneProducer)
-      return emitError(loc) << "scheduled batch step mixes host and producer-derived inputs";
-    FailureOr<Value> laneSource = getOriginalProducerValue(*laneProducer);
-    if (failed(laneSource))
-      return emitError(loc) << "cannot resolve original graph producer value for " << getInstanceName(laneProducer->instance);
-
-    auto sourceIt = llvm::find(uniqueOriginalSources, *laneSource);
-    if (sourceIt == uniqueOriginalSources.end()) {
-      sourceOperandForScheduledLane.push_back(static_cast<int64_t>(uniqueOriginalSources.size()));
-      uniqueOriginalSources.push_back(*laneSource);
-    } else {
-      sourceOperandForScheduledLane.push_back(std::distance(uniqueOriginalSources.begin(), sourceIt));
-    }
-
-    laneProducerInfos.push_back(LaneProducerInfo {
-      *laneProducer,
-      *laneSource,
-      getDeferredEndpoint(laneProducer->instance, schedule),
-      getDeferredEndpoint(instance, schedule),
+  for (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) absorbed.erase(op);
   }
-
-  for (const LaneProducerInfo &laneProducerInfo : laneProducerInfos) {
-    transferMetadata.producers.push_back(laneProducerInfo.producerEndpoint);
-    transferMetadata.consumers.push_back(laneProducerInfo.consumerEndpoint);
-  }
-  transferMetadata.sourceOperandForScheduledLane = sourceOperandForScheduledLane;
-  transferMetadata.multiSourcePayload = uniqueOriginalSources.size() > 1;
-
-  return createScheduledLaneDeferredCommunication(
-    builder, loc, input, uniqueOriginalSources, *representativeProducer, transferMetadata, exchangeId++, mapped);
+  return success();
 }
 
-} // namespace spatial
-} // namespace onnx_mlir
+} // namespace onnx_mlir::spatial

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

@@ -5,29 +5,42 @@
 namespace onnx_mlir {
 namespace spatial {
 
-LogicalResult mapSingleCpuInput(OpBuilder &builder,
-                                Location loc,
-                                Value input,
-                                const ComputeInstance &consumerInstance,
-                                const MergeScheduleResult &schedule,
-                                ValueRange scheduledInputs,
-                                Block &block,
-                                unsigned firstInputArgument,
-                                ArrayRef<ProducerValueKey> carriedKeys,
-                                uint64_t &exchangeId,
-                                Value &mapped);
+// 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;
+};
 
-LogicalResult mapMultiCpuTupleInput(OpBuilder &builder,
-                                    Location loc,
-                                    Value input,
-                                    const ComputeStepTuple &stepTuple,
-                                    const PeftClassPlan &peftClassPlan,
+LogicalResult prepareSingleCpuInput(OpBuilder &builder, Location loc, Value input,
+                                    BlockArgument graphInput,
+                                    const ComputeInstance &consumerInstance,
                                     const MergeScheduleResult &schedule,
-                                    ValueRange scheduledInputs,
-                                    Block &block,
+                                    ValueRange scheduledInputs, Block &block,
                                     unsigned firstInputArgument,
-                                    uint64_t &exchangeId,
-                                    Value &mapped);
+                                    ArrayRef<ProducerValueKey> carriedKeys,
+                                    Value graphLane, Value scheduledGraphLane,
+                                    DeferredInputPlan &plan);
+
+LogicalResult prepareMultiCpuTupleInput(OpBuilder &builder, Location loc, Value input,
+                                        BlockArgument graphInput,
+                                        const ComputeStepTuple &stepTuple,
+                                        const PeftClassPlan &peftClassPlan,
+                                        const MergeScheduleResult &schedule,
+                                        ValueRange scheduledInputs, Block &block,
+                                        unsigned firstInputArgument, Value graphLane, Value scheduledGraphLane,
+                                        Value scheduledLane, DeferredInputPlan &plan);
+
+LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc,
+                                                 Block &graphBody,
+                                                 ArrayRef<DeferredInputPlan> plans,
+                                                 IRMapping &mapper,
+                                                 llvm::SmallPtrSetImpl<Operation *> &absorbed);
 
 } // namespace spatial
 } // namespace onnx_mlir

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

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

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

@@ -1,6 +1,8 @@
 #include "ScheduledComputeMaterialization.hpp"
 #include "ScheduledComputeReport.hpp"
 #include "ScheduledComputeVerification.hpp"
+#include "DeferredCommunicationRealization.hpp"
+#include "DeferredCommunicationDeadlock.hpp"
 
 #include "mlir/Pass/Pass.h"
 
@@ -80,9 +82,20 @@ struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, Operatio
                                         schedule,
                                         materialization->peftClassPlans,
                                         materialization->materializedSchedules);
-    moduleOp.emitError("MergeComputeNodes stopped after spatial1_scheduled_no_comm; "
-                       "Phase 2 communication realization is not implemented");
-    signalPassFailure();
+    if (failed(realizeDeferredCommunication(funcOp))) {
+      moduleOp.emitError("MergeComputeNodes phase 2 communication realization failed");
+      signalPassFailure();
+      return;
+    }
+    if (failed(verifyRealizedCommunicationDeadlockFree(funcOp))) {
+      moduleOp.emitError("MergeComputeNodes final communication verification failed");
+      signalPassFailure();
+      return;
+    }
+    if (failed(verifyScheduledSpatialInvariants(funcOp))) {
+      moduleOp.emitError("scheduled Spatial phase 2 verification failed");
+      signalPassFailure();
+    }
   }
 };
 

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

@@ -10,6 +10,8 @@
 
 #include <map>
 
+#include "llvm/ADT/SmallPtrSet.h"
+
 namespace onnx_mlir {
 namespace spatial {
 
@@ -19,10 +21,6 @@ namespace {
 struct BatchFragmentSpec {
   RankedTensorType resultType;
   RankedTensorType sourceSliceType;
-  unsigned laneDim = 0;
-  SmallVector<OpFoldResult> offsets;
-  SmallVector<OpFoldResult> sizes;
-  SmallVector<OpFoldResult> strides;
 };
 
 static SmallVector<OpFoldResult> remapMixedOffsets(ArrayRef<OpFoldResult> mixedOffsets, IRMapping &mapper) {
@@ -128,42 +126,38 @@ static FailureOr<BatchFragmentSpec> getBatchFragmentSpec(SpatComputeBatch batch,
     RankedTensorType sourceType = insert.getSourceType();
     if (!destType || !sourceType || !destType.hasStaticShape() || !sourceType.hasStaticShape())
       return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
-    if (destType.getRank() == 0 || sourceType.getRank() == 0 || destType.getRank() != sourceType.getRank())
-      return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+    if (destType.getRank() != sourceType.getRank() + 1 || destType.getDimSize(0) != batch.getLaneCount()
+        || destType.getElementType() != sourceType.getElementType())
+      return batch.emitOpError("graph_compute_batch result must be a leading physical-slot dimension followed by its fragment");
+    if (!llvm::equal(destType.getShape().drop_front(), sourceType.getShape()))
+      return batch.emitOpError("graph_compute_batch result trailing shape must match its published fragment");
     if (!insert.hasUnitStride())
       return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
-
     auto offsets = insert.getMixedOffsets();
     auto sizes = insert.getMixedSizes();
     auto strides = insert.getMixedStrides();
     if (offsets.size() != static_cast<size_t>(destType.getRank()) || sizes.size() != static_cast<size_t>(destType.getRank())
         || strides.size() != static_cast<size_t>(destType.getRank()))
       return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
-
-    std::optional<unsigned> laneDim;
-    for (unsigned dim = 0; dim < offsets.size(); ++dim) {
-      bool dimLooksLaneLike = sourceType.getShape()[dim] != destType.getShape()[dim];
-      if (auto value = dyn_cast<Value>(offsets[dim]))
-        dimLooksLaneLike = dimLooksLaneLike || valueTransitivelyDependsOn(value, *batch.getLaneArgument());
-      if (dimLooksLaneLike) {
-        if (laneDim && *laneDim != dim)
-          return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
-        laneDim = dim;
-      }
+    if (!isa<Value>(offsets.front()) || !valueTransitivelyDependsOn(cast<Value>(offsets.front()), *batch.getLaneArgument()))
+      return batch.emitOpError("graph_compute_batch publication must select its physical slot in dimension zero");
+    for (unsigned dim = 1; dim < offsets.size(); ++dim) {
+      auto offset = dyn_cast<Attribute>(offsets[dim]);
+      auto integer = dyn_cast_or_null<IntegerAttr>(offset);
+      if (!integer || integer.getInt() != 0)
+        return batch.emitOpError("graph_compute_batch publication must have zero trailing offsets");
     }
-    if (!laneDim)
-      return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
-
-    SmallVector<int64_t> shape(destType.getShape().begin(), destType.getShape().end());
-    shape[*laneDim] = fragmentLaneCount;
-    return BatchFragmentSpec {
-      RankedTensorType::get(shape, destType.getElementType()),
-      sourceType,
-      *laneDim,
-      SmallVector<OpFoldResult>(offsets.begin(), offsets.end()),
-      SmallVector<OpFoldResult>(sizes.begin(), sizes.end()),
-      SmallVector<OpFoldResult>(strides.begin(), strides.end())
+    auto staticIndex = [](OpFoldResult value) -> std::optional<int64_t> {
+      auto attr = dyn_cast<Attribute>(value);
+      auto integer = dyn_cast_or_null<IntegerAttr>(attr);
+      return integer ? std::optional<int64_t>(integer.getInt()) : std::nullopt;
     };
+    if (staticIndex(sizes.front()) != 1)
+      return batch.emitOpError("graph_compute_batch publication sizes must be [1] plus the fragment shape");
+    for (auto [size, dim] : llvm::zip_equal(ArrayRef<OpFoldResult>(sizes).drop_front(), sourceType.getShape()))
+      if (staticIndex(size) != dim)
+        return batch.emitOpError("graph_compute_batch publication sizes must be [1] plus the fragment shape");
+    return BatchFragmentSpec {spatial::getGraphBatchPhysicalResultType(fragmentLaneCount, sourceType), sourceType};
   }
 
   return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
@@ -347,9 +341,12 @@ static LogicalResult collectPeftClassOperandsAndResults(PeftClassPlan &peftClass
   return success();
 }
 
-static void cloneComputeBody(OpBuilder &builder, Block &source, IRMapping &mapper, SmallVectorImpl<Value> &yieldedValues) {
+static void cloneComputeBody(OpBuilder &builder, Block &source, IRMapping &mapper,
+                             SmallVectorImpl<Value> &yieldedValues,
+                             const llvm::SmallPtrSetImpl<Operation *> &absorbed) {
   for (Operation &op : source.without_terminator())
-    builder.clone(op, mapper);
+    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));
@@ -362,7 +359,6 @@ static LogicalResult materializeResultfulBatchChunkAsScalar(PatternRewriter &rew
                                                             ValueRange scheduledInputs,
                                                             Block &block,
                                                             const MergeScheduleResult &schedule,
-                                                            uint64_t &exchangeId,
                                                             SmallVectorImpl<Value> &yieldedValues) {
   SmallVector<Value> initResults;
   SmallVector<BatchFragmentSpec> fragmentSpecs;
@@ -373,7 +369,7 @@ static LogicalResult materializeResultfulBatchChunkAsScalar(PatternRewriter &rew
     fragmentSpecs.push_back(*spec);
     auto empty = createEmptyTensorForType(rewriter, batch.getLoc(), spec->resultType);
     if (failed(empty))
-      return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+      return batch.emitOpError("scheduled materialization requires a physical graph fragment publication");
     initResults.push_back(*empty);
   }
 
@@ -393,28 +389,35 @@ static LogicalResult materializeResultfulBatchChunkAsScalar(PatternRewriter &rew
       mapper.map(*batch.getLaneArgument(), originalLane);
       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())) {
-        Value mapped;
-        if (failed(mapSingleCpuInput(builder,
+        DeferredInputPlan plan;
+        if (failed(prepareSingleCpuInput(builder,
                                      input.getLoc(),
                                      input,
+                                     *batch.getInputArgument(index),
                                      instance,
                                      schedule,
                                      scheduledInputs,
                                      block,
                                      scheduledWeights.size(),
                                      ArrayRef<ProducerValueKey> {},
-                                     exchangeId,
-                                     mapped)))
+                                     *batch.getLaneArgument(),
+                                     lower,
+                                     plan)))
           return failure();
-        mapper.map(*batch.getInputArgument(index), mapped);
+        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())
-        builder.clone(op, mapper);
+        if (!absorbed.contains(&op))
+          builder.clone(op, mapper);
 
       auto inParallel = dyn_cast<SpatInParallelOp>(source.getTerminator());
       if (!inParallel)
@@ -432,13 +435,13 @@ static LogicalResult materializeResultfulBatchChunkAsScalar(PatternRewriter &rew
       for (Operation &op : inParallel.getRegion().front()) {
         auto insert = dyn_cast<tensor::ParallelInsertSliceOp>(&op);
         if (!insert)
-          return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+          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 only supports regular leading-lane graph_compute_batch fragments"), failure();
+          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[fragmentSpecs[resultIndex].laneDim] = localLane;
+        offsets.front() = localLane;
         current[resultIndex] = tensor::InsertSliceOp::create(builder,
                                                              insert.getLoc(),
                                                              mapper.lookup(insert.getSource()),
@@ -463,8 +466,7 @@ static LogicalResult materializeSingleCpuPeftClass(
   const PeftClassPlan &peftClassPlan,
   const MergeScheduleResult &schedule,
   DenseMap<GraphComputeBlockKey, Block *> &graphComputeToBlockMap,
-  ScheduledMaterializationRecord &record,
-  uint64_t &exchangeId) {
+  ScheduledMaterializationRecord &record) {
   size_t cpu = peftClassPlan.cpus.front();
   auto instancesIt = peftClassPlan.instancesByCpu.find(cpu);
   assert(instancesIt != peftClassPlan.instancesByCpu.end() && "missing single-cpu schedule");
@@ -487,23 +489,29 @@ static LogicalResult materializeSingleCpuPeftClass(
       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())) {
-        Value mapped;
-        if (failed(mapSingleCpuInput(rewriter,
+        DeferredInputPlan plan;
+        if (failed(prepareSingleCpuInput(rewriter,
                                      input.getLoc(),
                                      input,
+                                     *compute.getInputArgument(index),
                                      instance,
                                      schedule,
                                      scheduled.getInputs(),
                                      *block,
                                      scheduled.getWeights().size(),
                                      carriedKeys,
-                                     exchangeId,
-                                     mapped)))
+                                     {},
+                                     {},
+                                     plan)))
           return failure();
-        mapper.map(*compute.getInputArgument(index), mapped);
+        inputPlans.push_back(std::move(plan));
       }
-      cloneComputeBody(rewriter, compute.getBody().front(), mapper, 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);
     } else {
       auto batch = cast<SpatComputeBatch>(instance.op);
       if (failed(materializeResultfulBatchChunkAsScalar(rewriter,
@@ -513,7 +521,6 @@ static LogicalResult materializeSingleCpuPeftClass(
                                                         scheduled.getInputs(),
                                                         *block,
                                                         schedule,
-                                                        exchangeId,
                                                         yieldedValues)))
         return failure();
     }
@@ -556,8 +563,7 @@ static SmallVector<OpFoldResult> buildScheduledOutputInsertOffsets(OpBuilder &bu
                                                                    Location loc,
                                                                    Value scheduledLane,
                                                                    int64_t lanesPerScheduledLane,
-                                                                   RankedTensorType localFragmentType,
-                                                                   unsigned laneDim) {
+                                                                   RankedTensorType localFragmentType) {
   SmallVector<OpFoldResult> offsets;
   Value scheduledOutputLane = scheduledLane;
   if (lanesPerScheduledLane != 1) {
@@ -570,8 +576,8 @@ static SmallVector<OpFoldResult> buildScheduledOutputInsertOffsets(OpBuilder &bu
                                     ValueRange {scheduledLane})
         .getResult();
   }
-  for (unsigned dim = 0; dim < static_cast<unsigned>(localFragmentType.getRank()); ++dim)
-    offsets.push_back(dim == laneDim ? OpFoldResult(scheduledOutputLane) : OpFoldResult(builder.getIndexAttr(0)));
+  offsets.push_back(scheduledOutputLane);
+  offsets.append(localFragmentType.getRank() - 1, OpFoldResult(builder.getIndexAttr(0)));
   return offsets;
 }
 
@@ -581,8 +587,7 @@ static LogicalResult materializeMultiCpuPeftClass(
   const PeftClassPlan &peftClassPlan,
   const MergeScheduleResult &schedule,
   DenseMap<GraphComputeBlockKey, Block *> &graphComputeToBlockMap,
-  ScheduledMaterializationRecord &record,
-  uint64_t &exchangeId) {
+  ScheduledMaterializationRecord &record) {
   std::map<std::vector<uint32_t>, Value> laneStartTableCache;
   ArrayRef<ScheduledStepPlan> stepPlans = record.stepPlans;
   for (const ScheduledStepPlan &stepPlan : stepPlans) {
@@ -615,31 +620,37 @@ static LogicalResult materializeMultiCpuPeftClass(
     Value scheduledLane = block->getArgument(0);
     const ComputeInstance &representative = stepTuple.instances.front();
     SmallVector<Value> finalLocalFragments;
-    SmallVector<unsigned> finalLaneDims;
 
     if (auto compute = dyn_cast<SpatCompute>(representative.op)) {
       IRMapping mapper;
       for (auto [index, weight] : llvm::enumerate(compute.getWeights()))
         mapper.map(*compute.getWeightArgument(index), block->getArgument(1 + index));
       unsigned firstInputArg = 1 + scheduled.getWeights().size();
+      SmallVector<DeferredInputPlan> inputPlans;
       for (auto [index, input] : llvm::enumerate(compute.getInputs())) {
-        Value mapped;
-        if (failed(mapMultiCpuTupleInput(rewriter,
+        DeferredInputPlan plan;
+        if (failed(prepareMultiCpuTupleInput(rewriter,
                                          input.getLoc(),
                                          input,
+                                         *compute.getInputArgument(index),
                                          stepTuple,
                                          peftClassPlan,
                                          schedule,
                                          scheduled.getInputs(),
                                          *block,
                                          firstInputArg,
-                                         exchangeId,
-                                         mapped)))
+                                         {},
+                                         {},
+                                         scheduledLane,
+                                         plan)))
           return failure();
-        mapper.map(*compute.getInputArgument(index), mapped);
+        inputPlans.push_back(std::move(plan));
       }
       SmallVector<Value> yieldedValues;
-      cloneComputeBody(rewriter, compute.getBody().front(), mapper, 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)
@@ -656,7 +667,6 @@ static LogicalResult materializeMultiCpuPeftClass(
                                                                     yielded,
                                                                     reassociation)
                                        .getResult());
-        finalLaneDims.push_back(0);
       }
     } else {
       auto batch = cast<SpatComputeBatch>(representative.op);
@@ -702,26 +712,34 @@ static LogicalResult materializeMultiCpuPeftClass(
           for (auto [index, weight] : llvm::enumerate(batch.getWeights()))
             mapper.map(*batch.getWeightArgument(index), block->getArgument(1 + index));
           unsigned firstInputArg = 1 + scheduled.getWeights().size();
+          SmallVector<DeferredInputPlan> inputPlans;
           for (auto [index, input] : llvm::enumerate(batch.getInputs())) {
-            Value mapped;
-            if (failed(mapMultiCpuTupleInput(builder,
+            DeferredInputPlan plan;
+            if (failed(prepareMultiCpuTupleInput(builder,
                                              input.getLoc(),
                                              input,
+                                             *batch.getInputArgument(index),
                                              stepTuple,
                                              peftClassPlan,
                                              schedule,
                                              scheduled.getInputs(),
                                              *block,
                                              firstInputArg,
-                                             exchangeId,
-                                             mapped)))
+                                             *batch.getLaneArgument(),
+                                             sourceLane,
+                                             scheduledLane,
+                                             plan)))
               return failure();
-            mapper.map(*batch.getInputArgument(index), mapped);
+            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())
-            builder.clone(op, mapper);
+            if (!absorbed.contains(&op))
+              builder.clone(op, mapper);
 
           auto inParallel = dyn_cast<SpatInParallelOp>(batch.getBody().front().getTerminator());
           if (!inParallel)
@@ -735,13 +753,13 @@ static LogicalResult materializeMultiCpuPeftClass(
           for (Operation &op : inParallel.getRegion().front()) {
             auto insert = dyn_cast<tensor::ParallelInsertSliceOp>(&op);
             if (!insert)
-              return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+              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 only supports regular leading-lane graph_compute_batch fragments"), failure();
+              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[fragmentSpecs[resultIndex].laneDim] = innerLane;
+            offsets.front() = innerLane;
             current[resultIndex] = tensor::InsertSliceOp::create(builder,
                                                                  insert.getLoc(),
                                                                  mapper.lookup(insert.getSource()),
@@ -757,8 +775,6 @@ static LogicalResult materializeMultiCpuPeftClass(
       if (failed(loop))
         return failure();
       finalLocalFragments.assign(loop->results.begin(), loop->results.end());
-      for (const BatchFragmentSpec &spec : fragmentSpecs)
-        finalLaneDims.push_back(spec.laneDim);
     }
 
     auto inParallel = SpatInParallelOp::create(rewriter, scheduled.getLoc());
@@ -771,8 +787,7 @@ static LogicalResult materializeMultiCpuPeftClass(
         scheduled.getLoc(),
         scheduledLane,
         lanesPerScheduledLane,
-        localFragmentType,
-        finalLaneDims[resultIndex]);
+        localFragmentType);
       SmallVector<OpFoldResult> sizes;
       SmallVector<OpFoldResult> strides;
       for (int64_t dim : localFragmentType.getShape()) {
@@ -800,6 +815,14 @@ 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);
@@ -829,7 +852,6 @@ materializeScheduledCompute(func::FuncOp funcOp,
   DenseMap<size_t, SpatScheduledComputeBatch> scheduledComputeBatches;
   DenseMap<size_t, size_t> classToRecordIndex;
   std::vector<ScheduledMaterializationRecord> materializedSchedules;
-  uint64_t exchangeId = 0;
 
   for (auto &entry : peftClassPlans) {
     PeftClassPlan &peftClassPlan = entry.second;
@@ -855,9 +877,11 @@ materializeScheduledCompute(func::FuncOp funcOp,
       SmallVector<int64_t> stepResultCounts;
       SmallVector<int64_t> sourceLaneStarts;
       SmallVector<int64_t> sourceLaneCounts;
+      SmallVector<int64_t> stepSourceIds;
       size_t resultOffset = 0;
       for (const ComputeInstance &instance : peftClassPlan.instancesByCpu.lookup(peftClassPlan.cpus.front())) {
         stepSources.push_back(rewriter.getStringAttr(getInstanceName(instance)));
+        stepSourceIds.push_back(graphIds.lookup(instance.op));
         sourceLaneSelectors.push_back(rewriter.getStringAttr(isa<SpatCompute>(instance.op) ? "scalar" : "affine"));
         size_t resultCount = getComputeInstanceResultValueCount(instance);
         stepResultOffsets.push_back(static_cast<int64_t>(resultOffset));
@@ -872,6 +896,7 @@ materializeScheduledCompute(func::FuncOp funcOp,
         }
       }
       scheduled->setAttr("scheduled.step_sources", rewriter.getArrayAttr(stepSources));
+      scheduled->setAttr("scheduled.step_source_ids", rewriter.getDenseI64ArrayAttr(stepSourceIds));
       scheduled->setAttr("scheduled.step_result_offsets", rewriter.getDenseI64ArrayAttr(stepResultOffsets));
       scheduled->setAttr("scheduled.step_result_counts", rewriter.getDenseI64ArrayAttr(stepResultCounts));
       scheduled->setAttr("scheduled.source_lane_starts", rewriter.getDenseI64ArrayAttr(sourceLaneStarts));
@@ -894,8 +919,10 @@ materializeScheduledCompute(func::FuncOp funcOp,
       SmallVector<int64_t> resultCounts;
       SmallVector<int64_t> sourceLaneStarts;
       SmallVector<int64_t> sourceLaneCounts;
+      SmallVector<int64_t> stepSourceIds;
       for (const ScheduledStepPlan &stepPlan : record.stepPlans) {
         stepSources.push_back(rewriter.getStringAttr(getInstanceName(stepPlan.stepTuple.instances.front())));
+        stepSourceIds.push_back(graphIds.lookup(stepPlan.stepTuple.instances.front().op));
         sourceLaneSelectors.push_back(rewriter.getStringAttr(usesAffineSourceLaneMapping(stepPlan.stepTuple) ? "affine" : "table"));
         resultOffsets.push_back(static_cast<int64_t>(stepPlan.resultOffset));
         resultCounts.push_back(static_cast<int64_t>(stepPlan.resultCount));
@@ -908,6 +935,7 @@ materializeScheduledCompute(func::FuncOp funcOp,
         {static_cast<int64_t>(record.stepPlans.size()), static_cast<int64_t>(peftClassPlan.cpus.size())},
         rewriter.getI64Type());
       scheduled->setAttr("scheduled.step_sources", rewriter.getArrayAttr(stepSources));
+      scheduled->setAttr("scheduled.step_source_ids", rewriter.getDenseI64ArrayAttr(stepSourceIds));
       scheduled->setAttr("scheduled.step_result_offsets", rewriter.getDenseI64ArrayAttr(resultOffsets));
       scheduled->setAttr("scheduled.step_result_counts", rewriter.getDenseI64ArrayAttr(resultCounts));
       scheduled->setAttr("scheduled.source_lane_starts", DenseElementsAttr::get(sourceLaneTableType, ArrayRef<int64_t>(sourceLaneStarts)));
@@ -931,8 +959,7 @@ materializeScheduledCompute(func::FuncOp funcOp,
                                                peftClassPlan,
                                                schedule,
                                                graphComputeToBlockMap,
-                                               record,
-                                               exchangeId)))
+                                               record)))
         return failure();
     } else {
       if (failed(materializeMultiCpuPeftClass(rewriter,
@@ -940,8 +967,7 @@ materializeScheduledCompute(func::FuncOp funcOp,
                                               peftClassPlan,
                                               schedule,
                                               graphComputeToBlockMap,
-                                              record,
-                                              exchangeId)))
+                                              record)))
         return failure();
     }
   }

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

@@ -80,22 +80,6 @@ struct SourceLaneSelector {
   SmallVector<uint32_t> tableValues;
 };
 
-struct DeferredEndpoint {
-  ComputeInstance instance;
-  size_t cpu = 0;
-  size_t canonicalClassId = 0;
-  uint32_t laneStart = 0;
-  uint32_t laneCount = 1;
-};
-
-struct DeferredTransferMetadata {
-  SmallVector<DeferredEndpoint> producers;
-  SmallVector<DeferredEndpoint> consumers;
-  bool isBatched = false;
-  SmallVector<int64_t> sourceOperandForScheduledLane;
-  bool multiSourcePayload = false;
-};
-
 struct ScheduledMaterializationRecord {
   Operation *scheduledOp = nullptr;
   size_t canonicalPeftClassId = 0;

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

@@ -81,59 +81,18 @@ LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp) {
   pim::CappedDiagnosticReporter diagnostics;
   funcOp.walk([&](SpatDeferredCommunicationOp transfer) {
     for (Value source : transfer.getSources()) {
-      if (isa_and_nonnull<SpatScheduledCompute, SpatScheduledComputeBatch>(source.getDefiningOp())) {
+      auto result = dyn_cast<OpResult>(source);
+      if (!result || !isa<SpatGraphCompute, SpatGraphComputeBatch>(result.getOwner())) {
         diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) {
-          illegalOp->emitOpError("phase-check deferred communication source operand must remain an original graph SSA producer value");
+          illegalOp->emitOpError("phase-check deferred communication source operand must be an original graph SSA result");
         });
       }
     }
-
-    if (auto scheduled = dyn_cast<SpatScheduledComputeBatch>(transfer->getParentOp())) {
-      auto batchedAttr = transfer->getAttrOfType<BoolAttr>("batched");
-      auto sourceOperandForScheduledLane =
-        transfer->getAttrOfType<DenseI64ArrayAttr>("source_operand_for_scheduled_lane");
-      auto multiSourcePayload = transfer->getAttrOfType<BoolAttr>("multi_source_payload");
-      auto sourceCpus = transfer->getAttrOfType<DenseI64ArrayAttr>("source_cpus");
-      auto sourceClasses = transfer->getAttrOfType<DenseI64ArrayAttr>("source_classes");
-      auto sourceLaneRanges = transfer->getAttrOfType<DenseI64ArrayAttr>("source_lane_ranges");
-      auto targetCpus = transfer->getAttrOfType<DenseI64ArrayAttr>("target_cpus");
-      auto targetClasses = transfer->getAttrOfType<DenseI64ArrayAttr>("target_classes");
-      auto targetLaneRanges = transfer->getAttrOfType<DenseI64ArrayAttr>("target_lane_ranges");
-      int64_t laneCount = scheduled.getLaneCount();
-      auto hasExpectedSize = [&](DenseI64ArrayAttr attr, int64_t expected, StringRef name) {
-        if (!attr || attr.size() != expected)
-          diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) {
-            illegalOp->emitOpError() << "phase-check scheduled batch deferred communication requires " << name
-                                     << " length " << expected;
-          });
-      };
-      if (!batchedAttr || !batchedAttr.getValue())
-        diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) {
-          illegalOp->emitOpError("phase-check scheduled batch deferred communication requires batched metadata");
-        });
-      hasExpectedSize(sourceCpus, laneCount, "source_cpus");
-      hasExpectedSize(sourceClasses, laneCount, "source_classes");
-      hasExpectedSize(targetCpus, laneCount, "target_cpus");
-      hasExpectedSize(targetClasses, laneCount, "target_classes");
-      hasExpectedSize(sourceLaneRanges, laneCount * 2, "source_lane_ranges");
-      hasExpectedSize(targetLaneRanges, laneCount * 2, "target_lane_ranges");
-      hasExpectedSize(sourceOperandForScheduledLane, laneCount, "source_operand_for_scheduled_lane");
-      if (!multiSourcePayload || multiSourcePayload.getValue() != (transfer.getSources().size() > 1))
-        diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) {
-          illegalOp->emitOpError("phase-check scheduled batch deferred communication has inconsistent multi_source_payload metadata");
-        });
-      if (sourceOperandForScheduledLane) {
-        for (int64_t sourceOperandIndex : sourceOperandForScheduledLane.asArrayRef()) {
-          if (sourceOperandIndex < 0 || sourceOperandIndex >= static_cast<int64_t>(transfer.getSources().size())) {
-            diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) {
-              illegalOp->emitOpError(
-                "phase-check scheduled batch deferred communication source_operand_for_scheduled_lane index is out of range");
-            });
-            break;
-          }
-        }
-      }
-    }
+    if (!transfer->getParentOfType<SpatScheduledCompute>() &&
+        !transfer->getParentOfType<SpatScheduledComputeBatch>())
+      diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) {
+        illegalOp->emitOpError("phase-check deferred communication must be inside a scheduled compute");
+      });
   });
 
   diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial deferred communication verification failed");

tools/analyze_spatial_ir_cardinality.py

@@ -0,0 +1,253 @@
+#!/usr/bin/env python3.13
+
+import argparse
+import re
+from collections import Counter, defaultdict
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Iterable
+
+
+OP_PATTERNS = {
+    "tensor.extract_slice": re.compile(r"\btensor\.extract_slice\b"),
+    "tensor.insert_slice": re.compile(r"\btensor\.insert_slice\b"),
+    "spat.channel_send": re.compile(r"\bspat\.channel_send\b"),
+    "spat.channel_receive": re.compile(r"\bspat\.channel_receive\b"),
+    "scf.for": re.compile(r"\bscf\.for\b"),
+    "tensor.empty": re.compile(r"\btensor\.empty\b"),
+}
+
+VALUE_RE = re.compile(r"^\s*(%[\w.$-]+)\s*=\s*(.+)$")
+TYPE_RE = re.compile(r":\s*([^:]+?)\s*(?:to|into|$)")
+CHANNEL_RE = re.compile(r"channel\s+(%c[-\w.$]+)")
+FROM_TO_RE = re.compile(r"from\s+(%c[-\w.$]+)\s+to\s+(%c[-\w.$]+)")
+EXTRACT_SLICE_RE = re.compile(
+    r"^\s*(%[\w.$-]+)\s*=\s*tensor\.extract_slice\s+(%[\w.$-]+)\[(.*?)\]\s*\[(.*?)\]\s*\[(.*?)\]\s*:\s*(.*?)\s+to\s+(.*)$"
+)
+INSERT_SLICE_RE = re.compile(
+    r"^\s*(%[\w.$-]+)\s*=\s*tensor\.insert_slice\s+(%[\w.$-]+)\s+into\s+(%[\w.$-]+)\[(.*?)\]\s*\[(.*?)\]\s*\[(.*?)\]\s*:\s*(.*?)\s+into\s+(.*)$"
+)
+CHANNEL_RECEIVE_RE = re.compile(
+    r"^\s*(%[\w.$-]+)\s*=\s*spat\.channel_receive\s+channel\s+(%c[-\w.$]+)\s+from\s+(%c[-\w.$]+)\s+to\s+(%c[-\w.$]+)\s*:\s*(.*)$"
+)
+CHANNEL_SEND_RE = re.compile(
+    r"^\s*spat\.channel_send\s+(%[\w.$-]+)\s+channel\s+(%c[-\w.$]+)\s+from\s+(%c[-\w.$]+)\s+to\s+(%c[-\w.$]+)\s*:\s*(.*)$"
+)
+CONST_INDEX_RE = re.compile(r"^\s*(%c[\w.$-]+)\s*=\s*arith\.constant\s+(-?\d+)\s*:\s*index\b")
+
+
+@dataclass
+class ChainGroup:
+    kind: str
+    signature: str
+    count: int = 0
+    first_line: int = 0
+    last_line: int = 0
+    fragment_type: str = ""
+    dest_type: str = ""
+    varying_dims: set[int] = field(default_factory=set)
+    rows: list[int | None] = field(default_factory=list)
+    channels: list[int | None] = field(default_factory=list)
+    sources: list[int | None] = field(default_factory=list)
+    targets: list[int | None] = field(default_factory=list)
+
+    def add(self,
+            line_no: int,
+            fragment_type: str,
+            dest_type: str,
+            offsets: list[str],
+            channel: int | None = None,
+            source: int | None = None,
+            target: int | None = None) -> None:
+        self.count += 1
+        if self.first_line == 0:
+            self.first_line = line_no
+        self.last_line = line_no
+        self.fragment_type = fragment_type or self.fragment_type
+        self.dest_type = dest_type or self.dest_type
+        numeric_offsets = []
+        for idx, offset in enumerate(offsets):
+            try:
+                numeric_offsets.append(int(offset))
+            except ValueError:
+                self.varying_dims.add(idx)
+                numeric_offsets.append(None)
+        if self.rows is not None:
+            row = numeric_offsets[2] if len(numeric_offsets) > 2 else None
+            self.rows.append(row)
+            if len(self.rows) >= 2 and self.rows[-1] != self.rows[-2]:
+                self.varying_dims.add(2)
+        self.channels.append(channel)
+        self.sources.append(source)
+        self.targets.append(target)
+
+
+def parse_const_indices(lines: Iterable[str]) -> dict[str, int]:
+    constants: dict[str, int] = {}
+    for line in lines:
+        match = CONST_INDEX_RE.match(line)
+        if match:
+            constants[match.group(1)] = int(match.group(2))
+    return constants
+
+
+def split_index_list(value: str) -> list[str]:
+    return [piece.strip() for piece in value.split(",") if piece.strip()]
+
+
+def decode_const_index(token: str, constants: dict[str, int]) -> int | None:
+    token = token.strip()
+    if token in constants:
+        return constants[token]
+    try:
+        return int(token)
+    except ValueError:
+        return None
+
+
+def sequence_kind(values: list[int | None]) -> str:
+    concrete = [value for value in values if value is not None]
+    if not concrete:
+        return "dynamic"
+    if len(concrete) == len(values) and all(b - a == 1 for a, b in zip(concrete, concrete[1:])):
+        return "consecutive"
+    if len(set(concrete)) == 1:
+        return "constant"
+    return "table"
+
+
+def analyze_file(path: Path) -> tuple[Counter, dict[tuple[str, str], ChainGroup]]:
+    text = path.read_text()
+    lines = text.splitlines()
+    consts = parse_const_indices(lines)
+    counts = Counter()
+    groups: dict[tuple[str, str], ChainGroup] = {}
+
+    for line in lines:
+        for name, pattern in OP_PATTERNS.items():
+            if pattern.search(line):
+                counts[name] += 1
+
+    value_defs: dict[str, tuple[str, int, re.Match[str] | None]] = {}
+    for line_no, line in enumerate(lines, start=1):
+        if match := CHANNEL_RECEIVE_RE.match(line):
+            value_defs[match.group(1)] = ("receive", line_no, match)
+        elif match := EXTRACT_SLICE_RE.match(line):
+            value_defs[match.group(1)] = ("extract", line_no, match)
+        elif match := INSERT_SLICE_RE.match(line):
+            source = match.group(2)
+            producer = value_defs.get(source)
+            if not producer:
+                continue
+
+            offsets = split_index_list(match.group(4))
+            sizes = split_index_list(match.group(5))
+            strides = split_index_list(match.group(6))
+            dest = match.group(3)
+            dest_type = match.group(8).strip()
+
+            if producer[0] == "receive":
+                recv = producer[2]
+                assert recv is not None
+                channel = decode_const_index(recv.group(2), consts)
+                source_core = decode_const_index(recv.group(3), consts)
+                target_core = decode_const_index(recv.group(4), consts)
+                signature = f"recv_insert:{recv.group(5).strip()}->{dest_type}|sizes={','.join(sizes)}|strides={','.join(strides)}"
+                group = groups.setdefault(
+                    ("receive_to_insert", signature),
+                    ChainGroup("receive_to_insert", signature),
+                )
+                group.add(match.start() and producer[1] or line_no,
+                          recv.group(5).strip(),
+                          dest_type,
+                          offsets,
+                          channel=channel,
+                          source=source_core,
+                          target=target_core)
+            elif producer[0] == "extract":
+                extract = producer[2]
+                assert extract is not None
+                extract_offsets = split_index_list(extract.group(3))
+                signature = (
+                    f"extract_insert:{extract.group(6).strip()}->{dest_type}|"
+                    f"extract_sizes={extract.group(4).strip()}|insert_sizes={','.join(sizes)}|"
+                    f"src={extract.group(2)}"
+                )
+                group = groups.setdefault(
+                    ("extract_to_insert", signature),
+                    ChainGroup("extract_to_insert", signature),
+                )
+                group.add(producer[1], extract.group(7).strip(), dest_type, offsets)
+                if extract_offsets and decode_const_index(extract_offsets[0], consts) is None:
+                    group.varying_dims.add(0)
+
+        elif match := CHANNEL_SEND_RE.match(line):
+            source_value = match.group(1)
+            producer = value_defs.get(source_value)
+            if not producer or producer[0] != "extract":
+                continue
+            extract = producer[2]
+            assert extract is not None
+            signature = (
+                f"extract_send:{extract.group(6).strip()}->{match.group(5).strip()}|"
+                f"extract_sizes={extract.group(4).strip()}|src={extract.group(2)}"
+            )
+            group = groups.setdefault(
+                ("extract_to_send", signature),
+                ChainGroup("extract_to_send", signature),
+            )
+            group.add(
+                producer[1],
+                extract.group(7).strip(),
+                match.group(5).strip(),
+                split_index_list(extract.group(3)),
+                channel=decode_const_index(match.group(2), consts),
+                source=decode_const_index(match.group(3), consts),
+                target=decode_const_index(match.group(4), consts),
+            )
+
+    return counts, groups
+
+
+def print_report(path: Path, counts: Counter, groups: dict[tuple[str, str], ChainGroup], limit: int) -> None:
+    print(f"== {path} ==")
+    print("counts:")
+    for name in OP_PATTERNS:
+      print(f"  {name}: {counts[name]}")
+
+    ranked = sorted(groups.values(), key=lambda group: (-group.count, group.first_line))
+    print("hot chains:")
+    for group in ranked[:limit]:
+        varying = ",".join(str(dim) for dim in sorted(group.varying_dims)) or "none"
+        print(f"  - kind: {group.kind}")
+        print(f"    lines: {group.first_line}-{group.last_line}")
+        print(f"    fragments: {group.count}")
+        print(f"    fragment_type: {group.fragment_type}")
+        print(f"    dest_type: {group.dest_type}")
+        print(f"    varying_dims: {varying}")
+        if group.rows:
+            print(f"    row_sequence: {sequence_kind(group.rows)}")
+        if group.channels:
+            print(f"    channel_ids: {sequence_kind(group.channels)}")
+        if group.sources:
+            print(f"    source_ids: {sequence_kind(group.sources)}")
+        if group.targets:
+            print(f"    target_ids: {sequence_kind(group.targets)}")
+        print(f"    signature: {group.signature}")
+    print()
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Analyze repeated Spatial/PIM tensor IR cardinality patterns.")
+    parser.add_argument("paths", nargs="+", help="MLIR files to analyze.")
+    parser.add_argument("--limit", type=int, default=12, help="Maximum number of hot chains to print per file.")
+    args = parser.parse_args()
+
+    for path_arg in args.paths:
+        path = Path(path_arg)
+        counts, groups = analyze_file(path)
+        print_report(path, counts, groups, args.limit)
+
+
+if __name__ == "__main__":
+    main()