Now something work but not trust us (phase 1 + phase 2 of merge)

src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp

@@ -141,7 +141,8 @@ collectTopLevelFragmentAssemblyCopies(OpResult result, RankedTensorType packedRe
     std::optional<StringRef> mode = blueprint.getMode();
     std::optional<ArrayRef<int64_t>> operandIndicesAttr = blueprint.getFragmentOperandIndices();
     std::optional<ArrayRef<int64_t>> sourceOffsetsAttr = blueprint.getFragmentSourceOffsets();
-    if (!mode || *mode != "fragment_assembly" || !operandIndicesAttr || !sourceOffsetsAttr)
+    std::optional<ArrayRef<int64_t>> sourceSlotsAttr = blueprint.getFragmentSourceSlots();
+    if (!mode || *mode != "fragment_assembly" || !operandIndicesAttr || !sourceOffsetsAttr || !sourceSlotsAttr)
       return failure();
     if (!blueprint.getOutput().hasOneUse() || !isa<func::ReturnOp>(*blueprint.getOutput().getUsers().begin()))
       return failure();
@@ -153,6 +154,9 @@ collectTopLevelFragmentAssemblyCopies(OpResult result, RankedTensorType packedRe
 
     ArrayRef<int64_t> operandIndices = *operandIndicesAttr;
     ArrayRef<int64_t> sourceOffsets = *sourceOffsetsAttr;
+    ArrayRef<int64_t> sourceSlots = *sourceSlotsAttr;
+    if (sourceSlots.size() != operandIndices.size())
+      return failure();
     ArrayRef<int64_t> flatOffsets = blueprint.getFragmentOffsets();
     ArrayRef<int64_t> flatSizes = blueprint.getFragmentSizes();
     ArrayRef<int64_t> flatStrides = *stridesAttr;
@@ -174,7 +178,8 @@ collectTopLevelFragmentAssemblyCopies(OpResult result, RankedTensorType packedRe
       if (operandIndices[fragmentIndex] != static_cast<int64_t>(use.getOperandNumber()))
         continue;
 
-      int64_t sourceElementOffset = sourceOffsets[fragmentIndex];
+      int64_t sourceElementOffset =
+        sourceSlots[fragmentIndex] * payloadElementCount + sourceOffsets[fragmentIndex];
       int64_t lane = sourceElementOffset / payloadElementCount;
       if (lane < 0 || lane >= static_cast<int64_t>(laneCount))
         return failure();

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

@@ -204,15 +204,32 @@ analyzeCopyRewrite(Value target, Value source, Value targetOffset, Value sourceO
   if (!targetType || !sourceType || size <= 0)
     return failure();
 
-  auto logicalCopyShape = inferLogicalCopyShape(targetType, sourceType, size);
-  if (failed(logicalCopyShape))
-    return failure();
-
   auto targetPlan = analyzeCopyEndpoint(target, targetOffset, targetType);
   auto sourcePlan = analyzeCopyEndpoint(source, sourceOffset, sourceType);
   if (failed(targetPlan) || failed(sourcePlan))
     return failure();
 
+  auto targetBytes = getShapedByteSize(targetType);
+  auto sourceBytes = getShapedByteSize(sourceType);
+  if (targetType.getElementType() == sourceType.getElementType() && succeeded(targetBytes) && succeeded(sourceBytes)
+      && *targetBytes == size && *sourceBytes == size) {
+    auto targetSuffixRank = getContiguousSuffixRank(targetType, targetType.getShape());
+    auto sourceSuffixRank = getContiguousSuffixRank(sourceType, sourceType.getShape());
+    if (succeeded(targetSuffixRank) && succeeded(sourceSuffixRank)
+        && *targetSuffixRank == targetType.getRank() && *sourceSuffixRank == sourceType.getRank()) {
+      CopyRewritePlan plan;
+      plan.kind = CopyRewritePlan::Kind::Direct;
+      plan.target = *targetPlan;
+      plan.source = *sourcePlan;
+      plan.directBytes = size;
+      return plan;
+    }
+  }
+
+  auto logicalCopyShape = inferLogicalCopyShape(targetType, sourceType, size);
+  if (failed(logicalCopyShape))
+    return failure();
+
   auto targetSuffixRank = getContiguousSuffixRank(targetType, *logicalCopyShape);
   auto sourceSuffixRank = getContiguousSuffixRank(sourceType, *logicalCopyShape);
   if (failed(targetSuffixRank) || failed(sourceSuffixRank))

src/PIM/Dialect/Spatial/CMakeLists.txt

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

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

@@ -2,6 +2,7 @@
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/IRMapping.h"
 
@@ -24,6 +25,93 @@ static FailureOr<Value> getOriginalProducerValue(const ProducerValueRef &produce
   return outputs[producer.resultIndex];
 }
 
+static SmallVector<Value> getBlueprintFragments(SpatBlueprintOp blueprint) {
+  SmallVector<Value> fragments {blueprint.getInput()};
+  llvm::append_range(fragments, blueprint.getFragments());
+  return fragments;
+}
+
+static FailureOr<Value> buildBlueprintReconstruction(
+    OpBuilder &builder, Location loc, SpatBlueprintOp blueprint,
+    ValueRange sourceBlockArgs) {
+  auto resultType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
+  auto operandIndices = blueprint.getFragmentOperandIndices();
+  auto sourceSlots = blueprint.getFragmentSourceSlots();
+  auto sourceOffsets = blueprint.getFragmentSourceOffsets();
+  auto strides = blueprint.getFragmentStrides();
+  if (!resultType || !resultType.hasStaticShape() || !operandIndices ||
+      !sourceSlots || !sourceOffsets || !strides)
+    return blueprint.emitOpError("phase 1 requires complete static fragment assembly metadata"), failure();
+  int64_t rank = resultType.getRank();
+  ArrayRef<int64_t> offsets = blueprint.getFragmentOffsets();
+  ArrayRef<int64_t> sizes = blueprint.getFragmentSizes();
+  if (offsets.size() != sizes.size() || offsets.size() != strides->size() ||
+      offsets.size() != operandIndices->size() * rank ||
+      sourceSlots->size() != operandIndices->size() ||
+      sourceOffsets->size() != operandIndices->size())
+    return blueprint.emitOpError("phase 1 fragment assembly metadata has inconsistent sizes"), failure();
+
+  Value result = tensor::EmptyOp::create(builder, loc, resultType.getShape(),
+                                         resultType.getElementType());
+  for (auto [fragmentIndex, operandIndex] : llvm::enumerate(*operandIndices)) {
+    if (operandIndex < 0 || operandIndex >= static_cast<int64_t>(sourceBlockArgs.size()))
+      return blueprint.emitOpError("phase 1 fragment assembly operand index is out of range"), failure();
+    auto physicalType = dyn_cast<RankedTensorType>(sourceBlockArgs[operandIndex].getType());
+    if (!physicalType || !physicalType.hasStaticShape() || physicalType.getRank() != rank + 1)
+      return blueprint.emitOpError("phase 1 fragment assembly source is not a physical fragment batch"), failure();
+    SmallVector<int64_t> fragmentShape(physicalType.getShape().drop_front());
+    int64_t linearOffset = (*sourceOffsets)[fragmentIndex];
+    SmallVector<int64_t> sourceCoordinates(rank);
+    for (int64_t dim = rank - 1; dim >= 0; --dim) {
+      sourceCoordinates[dim] = linearOffset % fragmentShape[dim];
+      linearOffset /= fragmentShape[dim];
+    }
+    if (linearOffset != 0)
+      return blueprint.emitOpError("phase 1 fragment source offset is out of range"), failure();
+
+    SmallVector<OpFoldResult> sliceOffsets, sliceSizes, sliceStrides;
+    sliceOffsets.push_back(builder.getIndexAttr((*sourceSlots)[fragmentIndex]));
+    sliceSizes.push_back(builder.getIndexAttr(1));
+    sliceStrides.push_back(builder.getIndexAttr(1));
+    SmallVector<int64_t> selectedShape {1};
+    for (int64_t dim = 0; dim < rank; ++dim) {
+      int64_t index = fragmentIndex * rank + dim;
+      int64_t size = sizes[index];
+      if ((*strides)[index] != 1 || sourceCoordinates[dim] < 0 || size <= 0 ||
+          sourceCoordinates[dim] + size > fragmentShape[dim])
+        return blueprint.emitOpError("phase 1 fragment geometry is unsupported"), failure();
+      sliceOffsets.push_back(builder.getIndexAttr(sourceCoordinates[dim]));
+      sliceSizes.push_back(builder.getIndexAttr(size));
+      sliceStrides.push_back(builder.getIndexAttr(1));
+      selectedShape.push_back(size);
+    }
+    auto selectedType = RankedTensorType::get(selectedShape, resultType.getElementType());
+    Value selected = tensor::ExtractSliceOp::create(
+        builder, loc, selectedType, sourceBlockArgs[operandIndex], sliceOffsets,
+        sliceSizes, sliceStrides);
+    SmallVector<int64_t> fragmentResultShape(selectedShape.begin() + 1,
+                                             selectedShape.end());
+    auto fragmentType = RankedTensorType::get(fragmentResultShape,
+                                              resultType.getElementType());
+    SmallVector<ReassociationIndices> reassociation {{0, 1}};
+    for (int64_t dim = 1; dim < rank; ++dim)
+      reassociation.push_back({dim + 1});
+    Value fragment = tensor::CollapseShapeOp::create(
+        builder, loc, fragmentType, selected, reassociation);
+    SmallVector<OpFoldResult> targetOffsets, targetSizes, targetStrides;
+    for (int64_t dim = 0; dim < rank; ++dim) {
+      int64_t index = fragmentIndex * rank + dim;
+      targetOffsets.push_back(builder.getIndexAttr(offsets[index]));
+      targetSizes.push_back(builder.getIndexAttr(sizes[index]));
+      targetStrides.push_back(builder.getIndexAttr((*strides)[index]));
+    }
+    result = tensor::InsertSliceOp::create(builder, loc, fragment, result,
+                                           targetOffsets, targetSizes,
+                                           targetStrides);
+  }
+  return result;
+}
+
 static bool isSupportedDeferredShapingOp(Operation *op) {
   return isa<tensor::ExtractSliceOp, tensor::InsertSliceOp, tensor::CollapseShapeOp,
              tensor::ExpandShapeOp, tensor::CastOp, tensor::EmptyOp, tensor::ExtractOp,
@@ -99,7 +187,7 @@ static bool isEligible(Value value, Block &body, const DeferredInputPlan &plan,
 
 static FailureOr<Value> clonePayloadRoot(Value root, Block &body, const DeferredInputPlan &plan,
                                          OpBuilder &builder, SpatDeferredCommunicationOp transfer,
-                                         Value selectedSource) {
+                                         Value selectedSource, Value boundGraphLane) {
   IRMapping mapping;
   mapping.map(plan.graphInput, selectedSource);
   std::function<FailureOr<Value>(Value)> cloneScheduledLane = [&](Value value) -> FailureOr<Value> {
@@ -118,7 +206,7 @@ static FailureOr<Value> clonePayloadRoot(Value root, Block &body, const Deferred
   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);
+      auto mappedLane = cloneScheduledLane(boundGraphLane ? boundGraphLane : plan.scheduledGraphLane);
       if (failed(mappedLane)) return failure();
       mapping.map(value, *mappedLane);
       return *mappedLane;
@@ -136,6 +224,62 @@ static FailureOr<Value> clonePayloadRoot(Value root, Block &body, const Deferred
   return clone(root);
 }
 
+static bool dependsOnGraphLane(Value value, Value graphLane, Block &body,
+                               llvm::SmallPtrSetImpl<Operation *> &seen) {
+  if (value == graphLane)
+    return true;
+  Operation *op = value.getDefiningOp();
+  if (!op || !isTopLevelShaping(op, body) || !seen.insert(op).second)
+    return false;
+  return llvm::any_of(op->getOperands(), [&](Value operand) {
+    return dependsOnGraphLane(operand, graphLane, body, seen);
+  });
+}
+
+static FailureOr<Value> buildPayloadAggregate(OpBuilder &builder, Location loc,
+                                              ArrayRef<Value> payloads) {
+  auto payloadType = dyn_cast<RankedTensorType>(payloads.front().getType());
+  if (!payloadType || !payloadType.hasStaticShape())
+    return failure();
+  SmallVector<int64_t> shape {static_cast<int64_t>(payloads.size())};
+  llvm::append_range(shape, payloadType.getShape());
+  auto aggregateType = RankedTensorType::get(shape, payloadType.getElementType());
+  auto empty = createEmptyTensorForType(builder, loc, aggregateType);
+  if (failed(empty)) return failure();
+  Value aggregate = *empty;
+  SmallVector<OpFoldResult> sizes, strides(shape.size(), builder.getIndexAttr(1));
+  sizes.push_back(builder.getIndexAttr(1));
+  for (int64_t dim : payloadType.getShape()) sizes.push_back(builder.getIndexAttr(dim));
+  SmallVector<ReassociationIndices> reassociation {{0, 1}};
+  for (int64_t dim = 1; dim < payloadType.getRank(); ++dim) reassociation.push_back({dim + 1});
+  SmallVector<int64_t> expandedShape {1}; llvm::append_range(expandedShape, payloadType.getShape());
+  auto expandedType = RankedTensorType::get(expandedShape, payloadType.getElementType());
+  for (auto [index, payload] : llvm::enumerate(payloads)) {
+    Value expanded = tensor::ExpandShapeOp::create(builder, loc, expandedType, payload, reassociation).getResult();
+    SmallVector<OpFoldResult> offsets(shape.size(), builder.getIndexAttr(0));
+    offsets[0] = builder.getIndexAttr(index);
+    aggregate = tensor::InsertSliceOp::create(builder, loc, expanded, aggregate, offsets, sizes, strides).getResult();
+  }
+  return aggregate;
+}
+
+static FailureOr<Value> selectPayloadAggregate(OpBuilder &builder, Location loc, Value aggregate,
+                                               Value localLane) {
+  auto aggregateType = cast<RankedTensorType>(aggregate.getType());
+  SmallVector<int64_t> payloadShape(aggregateType.getShape().begin() + 1, aggregateType.getShape().end());
+  auto payloadType = RankedTensorType::get(payloadShape, aggregateType.getElementType());
+  SmallVector<OpFoldResult> offsets(aggregateType.getRank(), builder.getIndexAttr(0)); offsets[0] = localLane;
+  SmallVector<OpFoldResult> sizes, strides(aggregateType.getRank(), builder.getIndexAttr(1));
+  sizes.push_back(builder.getIndexAttr(1));
+  for (int64_t dim : payloadShape) sizes.push_back(builder.getIndexAttr(dim));
+  SmallVector<int64_t> unitShape {1}; llvm::append_range(unitShape, payloadShape);
+  Value unit = tensor::ExtractSliceOp::create(builder, loc,
+    RankedTensorType::get(unitShape, aggregateType.getElementType()), aggregate, offsets, sizes, strides).getResult();
+  SmallVector<ReassociationIndices> reassociation {{0, 1}};
+  for (int64_t dim = 1; dim < payloadType.getRank(); ++dim) reassociation.push_back({dim + 1});
+  return tensor::CollapseShapeOp::create(builder, loc, payloadType, unit, reassociation).getResult();
+}
+
 static void collectClosure(Value value, Block &body, const DeferredInputPlan &plan,
                            llvm::SmallPtrSetImpl<Operation *> &ops) {
   Operation *op = value.getDefiningOp();
@@ -146,12 +290,27 @@ static void collectClosure(Value value, Block &body, const DeferredInputPlan &pl
 
 } // namespace
 
+bool isDeferredFragmentAssemblyInput(
+    Value input, const ComputeInstance &consumerInstance) {
+  auto blueprint = input.getDefiningOp<SpatBlueprintOp>();
+  if (!blueprint || blueprint.getMode() != "fragment_assembly")
+    return false;
+  return llvm::all_of(getBlueprintFragments(blueprint), [&](Value fragment) {
+    return getProducerValueRef(fragment, &consumerInstance).has_value();
+  });
+}
+
 LogicalResult prepareSingleCpuInput(OpBuilder &, Location loc, Value input, BlockArgument graphInput,
                                     const ComputeInstance &consumerInstance, const MergeScheduleResult &,
                                     ValueRange scheduledInputs, Block &block, unsigned firstInputArgument,
                                     ArrayRef<ProducerValueKey> carriedKeys, Value graphLane, Value scheduledGraphLane,
                                     DeferredInputPlan &plan) {
-  plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, {}};
+  plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, {}, {}, {}, {}, 1, nullptr};
+  if (isDeferredFragmentAssemblyInput(input, consumerInstance)) {
+    plan.blueprint = input.getDefiningOp<SpatBlueprintOp>();
+    plan.originalSources = getBlueprintFragments(plan.blueprint);
+    return success();
+  }
   auto producer = getProducerValueRef(input, &consumerInstance);
   if (!producer) { plan.availableValue = getBlockOperand(block, scheduledInputs, input, firstInputArgument); return success(); }
   ProducerValueKey key {producer->instance, producer->resultIndex};
@@ -171,8 +330,13 @@ LogicalResult prepareMultiCpuTupleInput(OpBuilder &, Location loc, Value input,
                                         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();
+  plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, scheduledLane, {}, {}, {}, 1, nullptr};
+  if (isDeferredFragmentAssemblyInput(input, representative)) {
+    plan.blueprint = input.getDefiningOp<SpatBlueprintOp>();
+    plan.originalSources = getBlueprintFragments(plan.blueprint);
+    return success();
+  }
   auto producer = getProducerValueRef(input, &representative);
   if (!producer) { plan.availableValue = getBlockOperand(block, scheduledInputs, input, firstInputArgument); return success(); }
   auto inputs = getComputeInstanceInputs(representative);
@@ -222,19 +386,67 @@ LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc
     llvm::sort(roots, [](Value a, Value b) { return a.getAsOpaquePointer() < b.getAsOpaquePointer(); });
     roots.erase(std::unique(roots.begin(), roots.end()), roots.end());
     for (Value root : roots) {
+      llvm::SmallPtrSet<Operation *, 16> laneDependencies;
+      bool scalarize = plan.scalarizedGraphLaneBase
+                    && dependsOnGraphLane(root, plan.graphLane, body, laneDependencies);
+      OpBuilder::InsertPoint restore = builder.saveInsertionPoint();
+      Operation *loop = nullptr;
+      if (scalarize) {
+        loop = builder.getInsertionBlock()->getParentOp();
+        if (loop && !isa<scf::ForOp>(loop))
+          loop = loop->getParentOfType<scf::ForOp>();
+        if (loop)
+          builder.setInsertionPoint(loop);
+        else if (plan.scalarizedHoistBlock)
+          builder.setInsertionPointToEnd(plan.scalarizedHoistBlock);
+        else
+          return emitError(loc) << "phase 1 scalarized deferred payload is missing a hoist point";
+      }
+      SmallVector<Value> payloads;
+      unsigned count = scalarize ? plan.scalarizedLaneCount : 1;
+      for (unsigned offset = 0; offset < count; ++offset) {
       auto transfer = SpatDeferredCommunicationOp::create(builder, loc, root.getType(), plan.originalSources);
       Block *deferred = builder.createBlock(&transfer.getBody(), transfer.getBody().end(),
         TypeRange {transfer.getSources().getTypes()}, SmallVector<Location>(transfer.getSources().size(), loc));
       builder.setInsertionPointToStart(deferred);
-      auto selected = buildSelectedDeferredSource(builder, loc, transfer, plan.scheduledLane,
-                                                  deferred->getArguments(), plan.sourceOperandForScheduledLane);
+      auto selected = plan.blueprint
+                        ? buildBlueprintReconstruction(builder, loc, plan.blueprint,
+                                                       deferred->getArguments())
+                        : buildSelectedDeferredSource(builder, loc, transfer,
+                                                      plan.scheduledLane,
+                                                      deferred->getArguments(),
+                                                      plan.sourceOperandForScheduledLane);
       if (failed(selected)) return failure();
-      auto payload = clonePayloadRoot(root, body, plan, builder, transfer, *selected);
+      Value boundGraphLane;
+      if (scalarize) {
+        Value offsetValue = arith::ConstantIndexOp::create(builder, loc, offset);
+        boundGraphLane = offset ? arith::AddIOp::create(builder, loc, plan.scalarizedGraphLaneBase, offsetValue).getResult()
+                                : plan.scalarizedGraphLaneBase;
+      }
+      auto payload = clonePayloadRoot(root, body, plan, builder, transfer, *selected, boundGraphLane);
       if (failed(payload)) return failure();
       SpatYieldOp::create(builder, loc, *payload);
-      mapper.map(root, transfer.getOutput());
-      collectClosure(root, body, plan, absorbed);
+      payloads.push_back(transfer.getOutput());
       builder.setInsertionPointAfter(transfer);
+      }
+      if (scalarize) {
+        builder.restoreInsertionPoint(restore);
+        if (payloads.size() == 1) {
+          mapper.map(root, payloads.front());
+        } else {
+          if (loop) builder.setInsertionPoint(loop);
+          else builder.setInsertionPointToEnd(plan.scalarizedHoistBlock);
+          auto aggregate = buildPayloadAggregate(builder, loc, payloads);
+          if (failed(aggregate)) return failure();
+          builder.restoreInsertionPoint(restore);
+          auto selected = selectPayloadAggregate(builder, loc, *aggregate, plan.scalarizedLocalLane);
+          if (failed(selected)) return failure();
+          mapper.map(root, *selected);
+        }
+      } else {
+        mapper.map(root, payloads.front());
+      }
+      collectClosure(root, body, plan, absorbed);
     }
   }
   for (Operation *op : absorbed) {

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

@@ -15,8 +15,16 @@ struct DeferredInputPlan {
   Value graphLane;
   Value scheduledGraphLane;
   Value scheduledLane;
+  SpatBlueprintOp blueprint;
+  Value scalarizedLocalLane;
+  Value scalarizedGraphLaneBase;
+  int64_t scalarizedLaneCount = 1;
+  Block *scalarizedHoistBlock = nullptr;
 };
 
+bool isDeferredFragmentAssemblyInput(Value input,
+                                     const ComputeInstance &consumerInstance);
+
 LogicalResult prepareSingleCpuInput(OpBuilder &builder, Location loc, Value input,
                                     BlockArgument graphInput,
                                     const ComputeInstance &consumerInstance,

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

@@ -0,0 +1,287 @@
+#include "DeferredProjectionAnalysis.hpp"
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/IR/Matchers.h"
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+
+#include <limits>
+
+namespace onnx_mlir::spatial {
+using namespace mlir;
+namespace {
+
+static FailureOr<int64_t> evaluate(Value value, const StaticIndexEnvironment &environment,
+                                   llvm::SmallDenseSet<Value, 16> &visiting) {
+  if (auto it = environment.bindings.find(value); it != environment.bindings.end())
+    return it->second;
+  if (!visiting.insert(value).second)
+    return failure();
+  if (auto constant = value.getDefiningOp<arith::ConstantOp>())
+    if (auto integer = dyn_cast<IntegerAttr>(constant.getValue()))
+      { visiting.erase(value); return integer.getInt(); }
+  if (auto cast = value.getDefiningOp<arith::IndexCastOp>()) {
+    auto result = evaluate(cast.getIn(), environment, visiting); visiting.erase(value); return result;
+  }
+  auto binary = [&](auto op, auto fn) -> FailureOr<int64_t> {
+    auto lhs = evaluate(op.getLhs(), environment, visiting);
+    auto rhs = evaluate(op.getRhs(), environment, visiting);
+    if (failed(lhs) || failed(rhs)) return failure();
+    return fn(*lhs, *rhs);
+  };
+  if (auto op = value.getDefiningOp<arith::AddIOp>()) { auto result = binary(op, [](int64_t a, int64_t b) -> FailureOr<int64_t> { int64_t r; if (llvm::AddOverflow(a, b, r)) return failure(); return r; }); visiting.erase(value); return result; }
+  if (auto op = value.getDefiningOp<arith::SubIOp>()) { auto result = binary(op, [](int64_t a, int64_t b) -> FailureOr<int64_t> { int64_t r; if (llvm::SubOverflow(a, b, r)) return failure(); return r; }); visiting.erase(value); return result; }
+  if (auto op = value.getDefiningOp<arith::MulIOp>()) { auto result = binary(op, [](int64_t a, int64_t b) -> FailureOr<int64_t> { int64_t r; if (llvm::MulOverflow(a, b, r)) return failure(); return r; }); visiting.erase(value); return result; }
+  if (auto apply = value.getDefiningOp<affine::AffineApplyOp>()) { auto result = evaluateAffineApply(apply, [&](Value operand) { return evaluate(operand, environment, visiting); }); visiting.erase(value); return result; }
+  if (auto extract = value.getDefiningOp<tensor::ExtractOp>()) {
+    auto constant = extract.getTensor().getDefiningOp<arith::ConstantOp>();
+    auto elements = constant ? dyn_cast<DenseIntElementsAttr>(constant.getValue()) : DenseIntElementsAttr();
+    auto type = elements ? dyn_cast<RankedTensorType>(elements.getType()) : RankedTensorType();
+    if (!elements || !type || extract.getIndices().size() != static_cast<size_t>(type.getRank())) return failure();
+    int64_t linear = 0;
+    for (auto [index, dim] : llvm::zip(extract.getIndices(), type.getShape())) {
+      auto i = evaluate(index, environment, visiting);
+      if (failed(i) || *i < 0 || *i >= dim) return failure();
+      linear = linear * dim + *i;
+    }
+    visiting.erase(value); return elements.getValues<APInt>()[linear].getSExtValue();
+  }
+  visiting.erase(value);
+  return failure();
+}
+
+static FailureOr<std::optional<unsigned>> sourceArgument(Value value, SpatDeferredCommunicationOp deferred,
+                                                          const StaticIndexEnvironment &environment) {
+  while (auto cast = value.getDefiningOp<tensor::CastOp>()) value = cast.getSource();
+  if (auto argument = dyn_cast<BlockArgument>(value);
+      argument && argument.getOwner() == &deferred.getBody().front()
+      && argument.getArgNumber() < deferred.getSources().size())
+    return std::optional<unsigned>(argument.getArgNumber());
+  // Phase 1's selector ends in collapse(extract_slice(stacked, table[lane])).
+  auto collapse = value.getDefiningOp<tensor::CollapseShapeOp>();
+  if (!collapse) return std::optional<unsigned>();
+  value = collapse.getSrc();
+  auto slice = value.getDefiningOp<tensor::ExtractSliceOp>();
+  if (!slice) return std::optional<unsigned>();
+  auto sourceType = dyn_cast<RankedTensorType>(slice.getSourceType());
+  if (!sourceType || slice.getMixedOffsets().size() != static_cast<size_t>(sourceType.getRank()))
+    return std::optional<unsigned>();
+  for (unsigned dim = 1; dim < sourceType.getRank(); ++dim) {
+    auto offset = evaluateDeferredIndex(slice.getMixedOffsets()[dim], environment);
+    auto size = evaluateDeferredIndex(slice.getMixedSizes()[dim], environment);
+    auto stride = evaluateDeferredIndex(slice.getMixedStrides()[dim], environment);
+    if (failed(offset) || failed(size) || failed(stride) || *offset != 0 || *size != sourceType.getDimSize(dim) || *stride != 1)
+      return std::optional<unsigned>();
+  }
+  auto leadingSize = evaluateDeferredIndex(slice.getMixedSizes().front(), environment);
+  auto leadingStride = evaluateDeferredIndex(slice.getMixedStrides().front(), environment);
+  if (failed(leadingSize) || failed(leadingStride) || *leadingSize != 1 || *leadingStride != 1)
+    return std::optional<unsigned>();
+  auto selected = evaluateDeferredIndex(slice.getMixedOffsets().front(), environment);
+  if (failed(selected)) return failure();
+  Value stacked = slice.getSource();
+  while (auto cast = stacked.getDefiningOp<tensor::CastOp>()) stacked = cast.getSource();
+  while (auto insert = stacked.getDefiningOp<tensor::InsertSliceOp>()) stacked = insert.getDest();
+  // The stack is a chain. Find the insertion at the selected leading offset.
+  for (Value cursor = slice.getSource(); auto insert = cursor.getDefiningOp<tensor::InsertSliceOp>(); cursor = insert.getDest()) {
+    auto offset = evaluateDeferredIndex(insert.getMixedOffsets().front(), environment);
+    if (succeeded(offset) && *offset == *selected) {
+      Value source = insert.getSource();
+      if (auto expand = source.getDefiningOp<tensor::ExpandShapeOp>()) source = expand.getSrc();
+      if (auto arg = dyn_cast<BlockArgument>(source); arg && arg.getOwner() == &deferred.getBody().front())
+        return std::optional<unsigned>(arg.getArgNumber());
+    }
+  }
+  return std::optional<unsigned>();
+}
+
+static bool isResidual(Operation *op) {
+  return isa<tensor::CollapseShapeOp, tensor::ExpandShapeOp, tensor::CastOp,
+             tensor::ExtractSliceOp, tensor::InsertSliceOp, tensor::ConcatOp,
+             tensor::EmptyOp, tensor::ExtractOp, arith::ConstantOp,
+             arith::IndexCastOp, arith::AddIOp, arith::SubIOp, arith::MulIOp,
+             affine::AffineApplyOp>(op);
+}
+
+static SpatGraphComputeBatch graphBatchOwner(Value value) {
+  if (auto result = dyn_cast<OpResult>(value))
+    return dyn_cast<SpatGraphComputeBatch>(result.getOwner());
+  return {};
+}
+
+static Value getEnclosingScheduledLane(SpatDeferredCommunicationOp deferred,
+                                       SpatScheduledComputeBatch scheduled) {
+  Block *block = deferred->getBlock();
+  while (block && block->getParentOp() != scheduled) {
+    Operation *parent = block->getParentOp();
+    block = parent ? parent->getBlock() : nullptr;
+  }
+  return block && !block->empty() ? block->getArgument(0) : Value();
+}
+
+} // namespace
+
+FailureOr<int64_t> evaluateDeferredIndex(Value value, const StaticIndexEnvironment &environment) {
+  llvm::SmallDenseSet<Value, 16> visiting;
+  return evaluate(value, environment, visiting);
+}
+FailureOr<int64_t> evaluateDeferredIndex(OpFoldResult value, const StaticIndexEnvironment &environment) {
+  if (auto attr = dyn_cast<Attribute>(value))
+    if (auto integer = dyn_cast<IntegerAttr>(attr)) return integer.getInt();
+  if (auto dynamic = dyn_cast<Value>(value)) return evaluateDeferredIndex(dynamic, environment);
+  return failure();
+}
+
+FailureOr<std::optional<ResolvedDeferredSource>> tryResolveDeferredSource(Value value, SpatDeferredCommunicationOp deferred,
+                                                                           const StaticIndexEnvironment &environment) {
+  auto index = sourceArgument(value, deferred, environment);
+  if (failed(index))
+    return failure();
+  if (!*index)
+    return std::optional<ResolvedDeferredSource>();
+  return std::optional<ResolvedDeferredSource>(ResolvedDeferredSource {**index, deferred.getSources()[**index]});
+}
+
+FailureOr<ResolvedDeferredSource> requireResolvedDeferredSource(Value value, SpatDeferredCommunicationOp deferred,
+                                                                 const StaticIndexEnvironment &environment) {
+  auto source = tryResolveDeferredSource(value, deferred, environment);
+  if (failed(source) || !*source)
+    return deferred.emitOpError("cannot statically resolve deferred source selection"), failure();
+  return **source;
+}
+
+FailureOr<SpecializedDeferredProgram> analyzeDeferredProgram(SpatDeferredCommunicationOp deferred,
+                                                              std::optional<unsigned> targetScheduledLane) {
+  Block &body = deferred.getBody().front();
+  auto yield = dyn_cast<SpatYieldOp>(body.getTerminator());
+  if (!yield || yield.getOutputs().size() != 1)
+    return deferred.emitOpError("requires exactly one deferred yielded value"), failure();
+  StaticIndexEnvironment environment;
+  if (auto scheduled = deferred->getParentOfType<SpatScheduledComputeBatch>()) {
+    if (!targetScheduledLane) return deferred.emitOpError("scheduled-batch deferred program requires lane specialization"), failure();
+    Value scheduledLane = getEnclosingScheduledLane(deferred, scheduled);
+    if (!scheduledLane)
+      return deferred.emitOpError("cannot locate the enclosing scheduled lane"), failure();
+    environment.bindings[scheduledLane] = *targetScheduledLane;
+  } else if (targetScheduledLane) return deferred.emitOpError("scalar deferred program cannot have a target lane"), failure();
+  SpecializedDeferredProgram program;
+  program.deferred = deferred;
+  program.targetScheduledLane = targetScheduledLane;
+  if (auto scheduled = deferred->getParentOfType<SpatScheduledComputeBatch>())
+    program.scheduledLane = getEnclosingScheduledLane(deferred, scheduled);
+  program.yieldedValue = yield.getOutputs().front();
+  llvm::SmallDenseSet<Value, 32> visited;
+  std::function<LogicalResult(Value)> visit = [&](Value value) -> LogicalResult {
+    if (!visited.insert(value).second) return success();
+    // A graph-batch projection is semantically different from the canonical
+    // source-selector scaffold even though both contain extract/collapse ops.
+    if (auto slice = value.getDefiningOp<tensor::ExtractSliceOp>()) {
+      auto source = tryResolveDeferredSource(slice.getSource(), deferred, environment);
+      if (failed(source)) return failure();
+      if (*source && graphBatchOwner((*source)->selectedValue)) {
+        auto type = dyn_cast<RankedTensorType>((*source)->selectedValue.getType());
+        auto result = dyn_cast<RankedTensorType>(slice.getResult().getType());
+        if (!type || !result || type.getRank() == 0) return deferred.emitOpError("graph projection requires ranked tensors");
+        auto offset = evaluateDeferredIndex(slice.getMixedOffsets().front(), environment);
+        if (failed(offset)) return deferred.emitOpError("graph projection leading offset is not statically resolvable");
+        auto size = evaluateDeferredIndex(slice.getMixedSizes().front(), environment);
+        if (failed(size)) return deferred.emitOpError("graph projection leading size is not statically resolvable");
+        auto stride = evaluateDeferredIndex(slice.getMixedStrides().front(), environment);
+        if (failed(stride)) return deferred.emitOpError("graph projection leading stride is not statically resolvable");
+        if (*offset < 0) return deferred.emitOpError("graph projection leading offset is negative");
+        if (*size <= 0) return deferred.emitOpError("graph projection leading size must be positive");
+        if (*stride <= 0) return deferred.emitOpError("graph projection leading stride must be positive");
+        DeferredProjectionLeaf leaf;
+        leaf.kind = DeferredLeafKind::GraphBatchProjection; leaf.sourceOperandIndex = (*source)->sourceOperandIndex;
+        leaf.replacementRoot = value; leaf.leadingProjection = slice; leaf.reconstructedType = result;
+        for (int64_t i = 0; i < *size; ++i) {
+          int64_t slot;
+          if (llvm::MulOverflow(i, *stride, slot) || llvm::AddOverflow(*offset, slot, slot)
+              || slot >= type.getDimSize(0))
+            return deferred.emitOpError("graph projection selects a physical slot outside the batch");
+          leaf.physicalSlots.push_back(slot);
+        }
+        for (unsigned i = 1; i < type.getRank(); ++i) {
+          auto innerOffset = evaluateDeferredIndex(slice.getMixedOffsets()[i], environment);
+          auto innerSize = evaluateDeferredIndex(slice.getMixedSizes()[i], environment);
+          auto innerStride = evaluateDeferredIndex(slice.getMixedStrides()[i], environment);
+          if (failed(innerOffset) || failed(innerSize) || failed(innerStride))
+            return deferred.emitOpError("graph projection has unresolved inner geometry");
+          leaf.innerGeometry.offsets.push_back(*innerOffset); leaf.innerGeometry.sizes.push_back(*innerSize); leaf.innerGeometry.strides.push_back(*innerStride);
+        }
+        program.leaves.push_back(std::move(leaf)); return success();
+      }
+    }
+
+    auto source = tryResolveDeferredSource(value, deferred, environment);
+    if (failed(source)) return failure();
+    if (*source) {
+      auto type = dyn_cast<RankedTensorType>((*source)->selectedValue.getType());
+      if (!type) return deferred.emitOpError("deferred source is not a ranked tensor");
+      DeferredProjectionLeaf leaf;
+      leaf.sourceOperandIndex = (*source)->sourceOperandIndex;
+      leaf.replacementRoot = value;
+      leaf.reconstructedType = type;
+      if (graphBatchOwner((*source)->selectedValue)) {
+        leaf.kind = DeferredLeafKind::GraphBatchIdentity;
+        for (int64_t slot = 0; slot < type.getDimSize(0); ++slot) leaf.physicalSlots.push_back(slot);
+      } else leaf.kind = DeferredLeafKind::ScalarSource;
+      program.leaves.push_back(std::move(leaf));
+      return success();
+    }
+    Operation *op = value.getDefiningOp();
+    if (!op || op->getBlock() != &body || !isResidual(op))
+      return deferred.emitOpError("deferred residual contains an unsupported operation");
+    for (Value operand : op->getOperands()) if (failed(visit(operand))) return failure();
+    program.residualOps.push_back(op); return success();
+  };
+  if (failed(visit(program.yieldedValue))) return failure();
+  return std::move(program);
+}
+
+FailureOr<const GraphBatchPublicationMap *> getGraphBatchPublicationMap(
+  SpatGraphComputeBatch graphBatch, unsigned resultIndex, GraphBatchPublicationCache &cache) {
+  GraphBatchPublicationKey key {graphBatch.getOperation(), resultIndex};
+  if (auto it = cache.find(key); it != cache.end()) return &it->second;
+  auto resultType = dyn_cast<RankedTensorType>(graphBatch.getResult(resultIndex).getType());
+  auto output = graphBatch.getOutputArgument(resultIndex);
+  auto lane = graphBatch.getLaneArgument();
+  if (!resultType || !output || !lane || resultType.getRank() == 0)
+    return graphBatch.emitOpError("graph batch publication is malformed"), failure();
+  tensor::ParallelInsertSliceOp publication;
+  auto parallel = dyn_cast<SpatInParallelOp>(graphBatch.getBody().front().getTerminator());
+  if (!parallel) return graphBatch.emitOpError("graph batch lacks publication region"), failure();
+  for (Operation &op : parallel.getRegion().front()) if (auto insert = dyn_cast<tensor::ParallelInsertSliceOp>(op); insert && insert.getDest() == *output) {
+    if (publication) return graphBatch.emitOpError("graph result has multiple publications"), failure();
+    publication = insert;
+  }
+  if (!publication) return graphBatch.emitOpError("graph result lacks publication"), failure();
+  auto fragment = dyn_cast<RankedTensorType>(publication.getSource().getType());
+  if (!fragment || resultType.getRank() != fragment.getRank() + 1) return graphBatch.emitOpError("graph publication fragment type is invalid"), failure();
+  GraphBatchPublicationMap map;
+  map.physicalResultType = resultType;
+  map.publicationFragmentType = fragment;
+  map.graphLaneToPhysicalSlot.resize(graphBatch.getLaneCount(), -1);
+  map.physicalSlotToGraphLane.resize(resultType.getDimSize(0), -1);
+  for (int64_t index = 0; index < graphBatch.getLaneCount(); ++index) {
+    StaticIndexEnvironment environment; environment.bindings[*lane] = index;
+    auto slot = evaluateDeferredIndex(publication.getMixedOffsets().front(), environment);
+    auto size = evaluateDeferredIndex(publication.getMixedSizes().front(), environment);
+    auto stride = evaluateDeferredIndex(publication.getMixedStrides().front(), environment);
+    if (failed(slot) || failed(size) || failed(stride) || *size != 1 || *stride != 1 || *slot < 0 || *slot >= resultType.getDimSize(0))
+      return graphBatch.emitOpError("graph publication leading geometry is invalid"), failure();
+    for (unsigned dim = 1; dim < resultType.getRank(); ++dim) {
+      auto offset = evaluateDeferredIndex(publication.getMixedOffsets()[dim], environment);
+      auto extent = evaluateDeferredIndex(publication.getMixedSizes()[dim], environment);
+      auto step = evaluateDeferredIndex(publication.getMixedStrides()[dim], environment);
+      if (failed(offset) || failed(extent) || failed(step) || *offset != 0 || *extent != fragment.getDimSize(dim - 1) || *step != 1)
+        return graphBatch.emitOpError("graph publication inner geometry is invalid"), failure();
+    }
+    if (map.physicalSlotToGraphLane[*slot] != -1) return graphBatch.emitOpError("graph publication has duplicate physical slot"), failure();
+    map.graphLaneToPhysicalSlot[index] = *slot; map.physicalSlotToGraphLane[*slot] = index;
+  }
+  if (llvm::is_contained(map.physicalSlotToGraphLane, -1)) return graphBatch.emitOpError("graph publication has missing physical slot"), failure();
+  return &cache.try_emplace(key, std::move(map)).first->second;
+}
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,102 @@
+#pragma once
+
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/Operation.h"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+#include "llvm/ADT/DenseMap.h"
+
+namespace onnx_mlir::spatial {
+
+struct StaticIndexEnvironment {
+  llvm::DenseMap<mlir::Value, int64_t> bindings;
+};
+
+mlir::FailureOr<int64_t> evaluateDeferredIndex(
+  mlir::Value value, const StaticIndexEnvironment &environment);
+mlir::FailureOr<int64_t> evaluateDeferredIndex(
+  mlir::OpFoldResult value, const StaticIndexEnvironment &environment);
+
+enum class DeferredLeafKind { ScalarSource, GraphBatchProjection, GraphBatchIdentity };
+
+struct StaticSliceGeometry {
+  llvm::SmallVector<int64_t> offsets;
+  llvm::SmallVector<int64_t> sizes;
+  llvm::SmallVector<int64_t> strides;
+};
+
+struct DeferredProjectionLeaf {
+  DeferredLeafKind kind = DeferredLeafKind::ScalarSource;
+  unsigned sourceOperandIndex = 0;
+  mlir::Value replacementRoot;
+  mlir::tensor::ExtractSliceOp leadingProjection;
+  llvm::SmallVector<int64_t> physicalSlots;
+  StaticSliceGeometry innerGeometry;
+  mlir::RankedTensorType reconstructedType;
+};
+
+struct SpecializedDeferredProgram {
+  SpatDeferredCommunicationOp deferred;
+  std::optional<unsigned> targetScheduledLane;
+  mlir::Value scheduledLane;
+  mlir::Value yieldedValue;
+  llvm::SmallVector<DeferredProjectionLeaf, 0> leaves;
+  llvm::SmallVector<mlir::Operation *> residualOps;
+};
+
+struct ResolvedDeferredSource {
+  unsigned sourceOperandIndex = 0;
+  mlir::Value selectedValue;
+};
+
+mlir::FailureOr<std::optional<ResolvedDeferredSource>> tryResolveDeferredSource(
+  mlir::Value value, SpatDeferredCommunicationOp deferred,
+  const StaticIndexEnvironment &environment);
+mlir::FailureOr<ResolvedDeferredSource> requireResolvedDeferredSource(
+  mlir::Value value, SpatDeferredCommunicationOp deferred,
+  const StaticIndexEnvironment &environment);
+mlir::FailureOr<SpecializedDeferredProgram> analyzeDeferredProgram(
+  SpatDeferredCommunicationOp deferred,
+  std::optional<unsigned> targetScheduledLane);
+
+struct GraphBatchPublicationMap {
+  mlir::RankedTensorType physicalResultType;
+  mlir::RankedTensorType publicationFragmentType;
+  llvm::SmallVector<int64_t> graphLaneToPhysicalSlot;
+  llvm::SmallVector<int64_t> physicalSlotToGraphLane;
+};
+
+struct GraphBatchPublicationKey {
+  mlir::Operation *graphBatch = nullptr;
+  unsigned resultIndex = 0;
+  bool operator==(const GraphBatchPublicationKey &other) const {
+    return graphBatch == other.graphBatch && resultIndex == other.resultIndex;
+  }
+};
+
+using GraphBatchPublicationCache =
+  llvm::DenseMap<GraphBatchPublicationKey, GraphBatchPublicationMap>;
+
+mlir::FailureOr<const GraphBatchPublicationMap *> getGraphBatchPublicationMap(
+  SpatGraphComputeBatch graphBatch, unsigned resultIndex,
+  GraphBatchPublicationCache &cache);
+
+} // namespace onnx_mlir::spatial
+
+namespace llvm {
+template <> struct DenseMapInfo<onnx_mlir::spatial::GraphBatchPublicationKey> {
+  static onnx_mlir::spatial::GraphBatchPublicationKey getEmptyKey() {
+    return {DenseMapInfo<mlir::Operation *>::getEmptyKey(), 0};
+  }
+  static onnx_mlir::spatial::GraphBatchPublicationKey getTombstoneKey() {
+    return {DenseMapInfo<mlir::Operation *>::getTombstoneKey(), 0};
+  }
+  static unsigned getHashValue(const onnx_mlir::spatial::GraphBatchPublicationKey &key) {
+    return hash_combine(key.graphBatch, key.resultIndex);
+  }
+  static bool isEqual(const onnx_mlir::spatial::GraphBatchPublicationKey &lhs,
+                      const onnx_mlir::spatial::GraphBatchPublicationKey &rhs) {
+    return lhs == rhs;
+  }
+};
+} // namespace llvm

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

@@ -9,6 +9,7 @@
 #include "Scheduling/MergeSchedulingAnalysis.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
 #include "src/Accelerators/PIM/Pass/PIMPasses.h"
 
 using namespace mlir;
@@ -44,6 +45,10 @@ struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, Operatio
       return;
     }
 
+    // Phase 1 is intentionally dumped before its verifier: malformed deferred
+    // payloads must be diagnosed from the producer-owned body.
+    dumpModule(moduleOp, "spatial2_scheduled_no_comm", /*assumeVerified=*/true);
+
     if (failed(verifyMaterializedScheduleMapping(funcOp,
                                                  schedule,
                                                  materialization->peftClassPlans,

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

@@ -299,7 +299,8 @@ static LogicalResult collectPeftClassOperandsAndResults(PeftClassPlan &peftClass
       for (Value weight : getComputeInstanceWeights(instance))
         appendUnique(peftClassPlan.weights, weight);
       for (Value input : getComputeInstanceInputs(instance))
-        if (!getProducerValueRef(input, &instance))
+        if (!getProducerValueRef(input, &instance) &&
+            !isDeferredFragmentAssemblyInput(input, instance))
           appendUnique(peftClassPlan.inputs, input);
     }
     return success();
@@ -333,7 +334,8 @@ static LogicalResult collectPeftClassOperandsAndResults(PeftClassPlan &peftClass
       for (Value weight : getComputeInstanceWeights(instance))
         appendUnique(peftClassPlan.weights, weight);
       for (Value input : getComputeInstanceInputs(instance))
-        if (!getProducerValueRef(input, &instance))
+        if (!getProducerValueRef(input, &instance) &&
+            !isDeferredFragmentAssemblyInput(input, instance))
           appendUnique(peftClassPlan.inputs, input);
     }
   }
@@ -387,6 +389,11 @@ static LogicalResult materializeResultfulBatchChunkAsScalar(PatternRewriter &rew
 
       IRMapping mapper;
       mapper.map(*batch.getLaneArgument(), originalLane);
+      Value localLane = arith::SubIOp::create(builder,
+                                              bodyLoc,
+                                              originalLane,
+                                              getOrCreateIndexConstant(rewriter, batch.getOperation(), instance.laneStart))
+                          .getResult();
       for (auto [index, weight] : llvm::enumerate(batch.getWeights()))
         mapper.map(*batch.getWeightArgument(index), getBlockOperand(block, scheduledWeights, weight));
       SmallVector<DeferredInputPlan> inputPlans;
@@ -403,9 +410,13 @@ static LogicalResult materializeResultfulBatchChunkAsScalar(PatternRewriter &rew
                                      scheduledWeights.size(),
                                      ArrayRef<ProducerValueKey> {},
                                      *batch.getLaneArgument(),
-                                     lower,
+                                     originalLane,
                                      plan)))
           return failure();
+        plan.scalarizedLocalLane = localLane;
+        plan.scalarizedGraphLaneBase = lower;
+        plan.scalarizedLaneCount = instance.laneCount;
+        plan.scalarizedHoistBlock = &block;
         inputPlans.push_back(std::move(plan));
       }
       for (auto [index, outputArg] : llvm::enumerate(batch.getOutputs()))
@@ -422,11 +433,6 @@ static LogicalResult materializeResultfulBatchChunkAsScalar(PatternRewriter &rew
       auto inParallel = dyn_cast<SpatInParallelOp>(source.getTerminator());
       if (!inParallel)
           return batch.emitOpError("expected spat.in_parallel in resultful spat.graph_compute_batch"), failure();
-      Value localLane = arith::SubIOp::create(builder,
-                                              bodyLoc,
-                                              originalLane,
-                                              getOrCreateIndexConstant(rewriter, batch.getOperation(), instance.laneStart))
-                          .getResult();
       DenseMap<BlockArgument, size_t> outputIndexByArg;
       for (size_t index = 0; index < batch.getNumResults(); ++index)
         outputIndexByArg[*batch.getOutputArgument(index)] = index;
@@ -624,7 +630,8 @@ static LogicalResult materializeMultiCpuPeftClass(
     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));
+        mapper.map(*compute.getWeightArgument(index),
+                   getBlockOperand(*block, scheduled.getWeights(), weight, 1));
       unsigned firstInputArg = 1 + scheduled.getWeights().size();
       SmallVector<DeferredInputPlan> inputPlans;
       for (auto [index, input] : llvm::enumerate(compute.getInputs())) {
@@ -686,6 +693,10 @@ static LogicalResult materializeMultiCpuPeftClass(
       Value lower = getOrCreateIndexConstant(rewriter, scheduled.getOperation(), 0);
       Value upper = getOrCreateIndexConstant(rewriter, scheduled.getOperation(), representative.laneCount);
       Value step = getOrCreateIndexConstant(rewriter, scheduled.getOperation(), 1);
+      FailureOr<Value> sourceLaneStart =
+        buildSourceLaneStartForScheduledLane(rewriter, batch.getLoc(), scheduledLane, sourceLaneSelector, scheduled.getOperation());
+      if (failed(sourceLaneStart))
+        return failure();
       auto loop = buildNormalizedScfFor(
         rewriter,
         batch.getLoc(),
@@ -696,10 +707,6 @@ static LogicalResult materializeMultiCpuPeftClass(
         [&](OpBuilder &builder, Location bodyLoc, Value innerLane, ValueRange iterArgs, SmallVectorImpl<Value> &yielded) -> LogicalResult {
 
           IRMapping mapper;
-          FailureOr<Value> sourceLaneStart =
-            buildSourceLaneStartForScheduledLane(builder, bodyLoc, scheduledLane, sourceLaneSelector, scheduled.getOperation());
-          if (failed(sourceLaneStart))
-            return failure();
           Value sourceLane =
             affine::AffineApplyOp::create(builder,
                                           bodyLoc,
@@ -710,7 +717,8 @@ static LogicalResult materializeMultiCpuPeftClass(
               .getResult();
           mapper.map(*batch.getLaneArgument(), sourceLane);
           for (auto [index, weight] : llvm::enumerate(batch.getWeights()))
-            mapper.map(*batch.getWeightArgument(index), block->getArgument(1 + index));
+            mapper.map(*batch.getWeightArgument(index),
+                       getBlockOperand(*block, scheduled.getWeights(), weight, 1));
           unsigned firstInputArg = 1 + scheduled.getWeights().size();
           SmallVector<DeferredInputPlan> inputPlans;
           for (auto [index, input] : llvm::enumerate(batch.getInputs())) {
@@ -730,6 +738,10 @@ static LogicalResult materializeMultiCpuPeftClass(
                                              scheduledLane,
                                              plan)))
               return failure();
+            plan.scalarizedLocalLane = innerLane;
+            plan.scalarizedGraphLaneBase = *sourceLaneStart;
+            plan.scalarizedLaneCount = representative.laneCount;
+            plan.scalarizedHoistBlock = block;
             inputPlans.push_back(std::move(plan));
           }
           for (unsigned index = 0; index < batch.getNumResults(); ++index)
@@ -838,7 +850,10 @@ materializeScheduledCompute(func::FuncOp funcOp,
     llvm::sort(peftClassPlan.cpus);
     for (size_t cpu : peftClassPlan.cpus)
       llvm::sort(peftClassPlan.instancesByCpu[cpu], [&](const ComputeInstance &lhs, const ComputeInstance &rhs) {
-        return schedule.computeToCpuSlotMap.lookup(lhs) < schedule.computeToCpuSlotMap.lookup(rhs);
+        return std::tie(graphIds.find(lhs.op)->second,
+                        schedule.computeToCpuSlotMap.find(lhs)->second) <
+               std::tie(graphIds.find(rhs.op)->second,
+                        schedule.computeToCpuSlotMap.find(rhs)->second);
       });
     if (failed(verifyPeftClassPlan(funcOp.getOperation(), peftClassPlan, schedule)))
       return failure();

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

@@ -64,8 +64,6 @@ struct PeftMaterializationReportSummary {
   size_t scheduledCompute = 0;
   size_t scheduledComputeBatch = 0;
   size_t deferredCommunication = 0;
-  size_t deferredCommunicationScalarMetadata = 0;
-  size_t deferredCommunicationScheduledLaneMetadata = 0;
   size_t deferredCommunicationMultiSourcePayloads = 0;
 };
 
@@ -95,10 +93,6 @@ static PeftMaterializationReportSummary buildPeftMaterializationReportSummary(
   }
   funcOp.walk([&](SpatDeferredCommunicationOp transfer) {
     summary.deferredCommunication++;
-    if (auto batchedAttr = transfer->getAttrOfType<BoolAttr>("batched"); batchedAttr && batchedAttr.getValue())
-      summary.deferredCommunicationScheduledLaneMetadata++;
-    else
-      summary.deferredCommunicationScalarMetadata++;
     if (transfer.getSources().size() > 1)
       summary.deferredCommunicationMultiSourcePayloads++;
   });
@@ -200,8 +194,6 @@ static void dumpPeftMaterializationReport(ModuleOp moduleOp,
   os << "  scheduled_compute_batch: " << summary.scheduledComputeBatch << "\n";
   os << "Deferred communications:\n";
   os << "  total: " << summary.deferredCommunication << "\n";
-  os << "  scalar metadata: " << summary.deferredCommunicationScalarMetadata << "\n";
-  os << "  scheduled-lane metadata: " << summary.deferredCommunicationScheduledLaneMetadata << "\n";
   os << "  multi-source payloads: " << summary.deferredCommunicationMultiSourcePayloads << "\n\n";
 
   os << "PEFT Classes\n";

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

@@ -1,4 +1,5 @@
 #include "ScheduledComputeVerification.hpp"
+#include "DeferredProjectionAnalysis.hpp"
 
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
@@ -79,6 +80,7 @@ LogicalResult verifyMaterializedScheduleMapping(
 
 LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp) {
   pim::CappedDiagnosticReporter diagnostics;
+  GraphBatchPublicationCache publicationCache;
   funcOp.walk([&](SpatDeferredCommunicationOp transfer) {
     for (Value source : transfer.getSources()) {
       auto result = dyn_cast<OpResult>(source);
@@ -92,7 +94,24 @@ LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp) {
         !transfer->getParentOfType<SpatScheduledComputeBatch>())
       diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) {
         illegalOp->emitOpError("phase-check deferred communication must be inside a scheduled compute");
-      });
+        });
+    if (auto scheduled = transfer->getParentOfType<SpatScheduledComputeBatch>()) {
+      for (unsigned lane = 0; lane < static_cast<unsigned>(scheduled.getLaneCount()); ++lane) {
+        auto program = analyzeDeferredProgram(transfer, lane);
+        if (failed(program))
+          continue;
+        for (const DeferredProjectionLeaf &leaf : program->leaves) {
+          if (leaf.kind == DeferredLeafKind::ScalarSource)
+            continue;
+          auto source = dyn_cast<OpResult>(transfer.getSources()[leaf.sourceOperandIndex]);
+          auto graph = source ? dyn_cast<SpatGraphComputeBatch>(source.getOwner()) : SpatGraphComputeBatch();
+          if (!graph) continue;
+          if (failed(getGraphBatchPublicationMap(graph, source.getResultNumber(), publicationCache)))
+            diagnostics.report(transfer.getOperation(), [&](Operation *) {});
+        }
+      }
+    } else
+      (void)analyzeDeferredProgram(transfer, std::nullopt);
   });
 
   diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial deferred communication verification failed");