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

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