a new new beginning phase 1

src/PIM/Common/Support/DebugDump.cpp

@@ -24,7 +24,7 @@ void dumpModule(mlir::ModuleOp moduleOp, const std::string& name, bool assumeVer
 
   llvm::raw_os_ostream os(file);
   mlir::OpPrintingFlags flags;
-  flags.elideLargeElementsAttrs().enableDebugInfo(true, false);
+  flags.elideLargeElementsAttrs().enableDebugInfo(false, false);
   if (assumeVerified)
     flags.assumeVerified();
   moduleOp.print(os, flags);

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

@@ -60,6 +60,56 @@ struct SpatComputeBatchBodyArgs {
   mlir::ValueRange outputs;
 };
 
+inline mlir::SmallVector<mlir::Type> getGraphComputeBlockArgTypes(mlir::ValueRange weights, mlir::ValueRange inputs) {
+  mlir::SmallVector<mlir::Type> blockArgTypes;
+  blockArgTypes.reserve(weights.size() + inputs.size());
+  for (mlir::Value weight : weights)
+    blockArgTypes.push_back(weight.getType());
+  for (mlir::Value input : inputs)
+    blockArgTypes.push_back(input.getType());
+  return blockArgTypes;
+}
+
+inline mlir::SmallVector<mlir::Location> getGraphComputeBlockArgLocs(mlir::Location defaultLoc,
+                                                                     mlir::ValueRange weights,
+                                                                     mlir::ValueRange inputs) {
+  mlir::SmallVector<mlir::Location> blockArgLocs;
+  blockArgLocs.reserve(weights.size() + inputs.size());
+  for (mlir::Value weight : weights)
+    blockArgLocs.push_back(weight.getLoc());
+  for (mlir::Value input : inputs)
+    blockArgLocs.push_back(input.getLoc());
+  return blockArgLocs;
+}
+
+inline mlir::SmallVector<mlir::Type> getGraphComputeBatchBlockArgTypes(mlir::OpBuilder& builder,
+                                                                       mlir::TypeRange resultTypes,
+                                                                       mlir::ValueRange weights,
+                                                                       mlir::ValueRange inputs) {
+  mlir::SmallVector<mlir::Type> blockArgTypes {builder.getIndexType()};
+  blockArgTypes.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
+  for (mlir::Value weight : weights)
+    blockArgTypes.push_back(weight.getType());
+  for (mlir::Value input : inputs)
+    blockArgTypes.push_back(input.getType());
+  llvm::append_range(blockArgTypes, resultTypes);
+  return blockArgTypes;
+}
+
+inline mlir::SmallVector<mlir::Location> getGraphComputeBatchBlockArgLocs(mlir::Location defaultLoc,
+                                                                          mlir::TypeRange resultTypes,
+                                                                          mlir::ValueRange weights,
+                                                                          mlir::ValueRange inputs) {
+  mlir::SmallVector<mlir::Location> blockArgLocs {defaultLoc};
+  blockArgLocs.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
+  for (mlir::Value weight : weights)
+    blockArgLocs.push_back(weight.getLoc());
+  for (mlir::Value input : inputs)
+    blockArgLocs.push_back(input.getLoc());
+  blockArgLocs.append(resultTypes.size(), defaultLoc);
+  return blockArgLocs;
+}
+
 } // namespace detail
 
 template <typename RewriterT>
@@ -87,6 +137,31 @@ inline mlir::Value createSpatConcat(RewriterT& rewriter, mlir::Location loc, int
   return spatial::SpatConcatOp::create(rewriter, loc, outputType, rewriter.getI64IntegerAttr(axis), inputs).getOutput();
 }
 
+template <typename RewriterT>
+spatial::SpatGraphCompute createEmptySpatGraphCompute(RewriterT& rewriter,
+                                                      mlir::Location loc,
+                                                      mlir::TypeRange resultTypes,
+                                                      mlir::ValueRange weights,
+                                                      mlir::ValueRange inputs,
+                                                      mlir::TypeRange blockArgTypes,
+                                                      llvm::ArrayRef<mlir::Location> blockArgLocs) {
+  auto computeOp = spatial::SpatGraphCompute::create(rewriter, loc, resultTypes, weights, inputs);
+  rewriter.createBlock(&computeOp.getBody(), computeOp.getBody().end(), blockArgTypes, blockArgLocs);
+  rewriter.setInsertionPointToStart(&computeOp.getBody().front());
+  return computeOp;
+}
+
+template <typename RewriterT>
+spatial::SpatGraphCompute createEmptySpatGraphCompute(RewriterT& rewriter,
+                                                      mlir::Location loc,
+                                                      mlir::TypeRange resultTypes,
+                                                      mlir::ValueRange weights,
+                                                      mlir::ValueRange inputs) {
+  auto blockArgTypes = detail::getGraphComputeBlockArgTypes(weights, inputs);
+  auto blockArgLocs = detail::getGraphComputeBlockArgLocs(loc, weights, inputs);
+  return createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs, blockArgTypes, blockArgLocs);
+}
+
 /// Builds a `spat.graph_compute` with a fixed number of SSA inputs and erases it if
 /// the body callback reports failure.
 template <size_t NumInputs, typename RewriterT, typename BodyFn>
@@ -97,16 +172,8 @@ auto createSpatGraphCompute(RewriterT& rewriter,
                             mlir::ValueRange inputs,
                             BodyFn&& body) {
   assert(inputs.size() == NumInputs && "NumInputs must match the number of input values");
-  auto computeOp = spatial::SpatGraphCompute::create(rewriter, loc, resultTypes, weights, inputs);
-
-  auto* block = new mlir::Block();
-  for (mlir::Value weight : weights)
-    block->addArgument(weight.getType(), loc);
-  for (mlir::Value input : inputs)
-    block->addArgument(input.getType(), loc);
-
-  computeOp.getBody().push_back(block);
-  rewriter.setInsertionPointToStart(block);
+  auto computeOp = createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs);
+  auto* block = &computeOp.getBody().front();
 
   using BodyResult = detail::InvokeWithBlockArgsResultT<std::decay_t<BodyFn>, std::make_index_sequence<NumInputs>>;
   if constexpr (std::is_same_v<BodyResult, void>) {
@@ -140,16 +207,8 @@ auto createSpatGraphCompute(RewriterT& rewriter,
                             mlir::ValueRange weights,
                             mlir::ValueRange inputs,
                             BodyFn&& body) {
-  auto computeOp = spatial::SpatGraphCompute::create(rewriter, loc, resultTypes, weights, inputs);
-
-  auto* block = new mlir::Block();
-  for (mlir::Value weight : weights)
-    block->addArgument(weight.getType(), loc);
-  for (mlir::Value input : inputs)
-    block->addArgument(input.getType(), loc);
-
-  computeOp.getBody().push_back(block);
-  rewriter.setInsertionPointToStart(block);
+  auto computeOp = createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs);
+  auto* block = &computeOp.getBody().front();
 
   using BodyResult = detail::InvokeWithValueRangeResultT<std::decay_t<BodyFn>>;
   if constexpr (std::is_same_v<BodyResult, void>) {
@@ -170,14 +229,15 @@ auto createSpatGraphCompute(RewriterT& rewriter,
   }
 }
 
-template <typename RewriterT, typename BodyFn>
-auto createSpatGraphComputeBatch(RewriterT& rewriter,
-                                 mlir::Location loc,
-                                 mlir::TypeRange resultTypes,
-                                 int64_t laneCount,
-                                 mlir::ValueRange weights,
-                                 mlir::ValueRange inputs,
-                                 BodyFn&& body) {
+template <typename RewriterT>
+auto createEmptySpatGraphComputeBatch(RewriterT& rewriter,
+                                      mlir::Location loc,
+                                      mlir::TypeRange resultTypes,
+                                      int64_t laneCount,
+                                      mlir::ValueRange weights,
+                                      mlir::ValueRange inputs,
+                                      mlir::TypeRange blockArgTypes,
+                                      llvm::ArrayRef<mlir::Location> blockArgLocs) {
   if (laneCount <= 0 || laneCount > std::numeric_limits<int32_t>::max())
     return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
 
@@ -186,27 +246,36 @@ auto createSpatGraphComputeBatch(RewriterT& rewriter,
     return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
 
   auto batchOp = spatial::SpatGraphComputeBatch::create(rewriter, loc, resultTypes, *laneCountAttr, weights, inputs);
+  rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), blockArgTypes, blockArgLocs);
+  rewriter.setInsertionPointToStart(&batchOp.getBody().front());
+  return mlir::FailureOr<spatial::SpatGraphComputeBatch>(batchOp);
+}
 
-  mlir::SmallVector<mlir::Type> blockArgTypes {rewriter.getIndexType()};
-  mlir::SmallVector<mlir::Location> blockArgLocs {loc};
-  blockArgTypes.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
-  blockArgLocs.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
-  for (mlir::Value weight : weights) {
-    blockArgTypes.push_back(weight.getType());
-    blockArgLocs.push_back(weight.getLoc());
-  }
-  for (mlir::Value input : inputs) {
-    blockArgTypes.push_back(input.getType());
-    blockArgLocs.push_back(input.getLoc());
-  }
-  for (mlir::Type resultType : resultTypes) {
-    blockArgTypes.push_back(resultType);
-    blockArgLocs.push_back(loc);
-  }
+template <typename RewriterT>
+auto createEmptySpatGraphComputeBatch(RewriterT& rewriter,
+                                      mlir::Location loc,
+                                      mlir::TypeRange resultTypes,
+                                      int64_t laneCount,
+                                      mlir::ValueRange weights,
+                                      mlir::ValueRange inputs) {
+  auto blockArgTypes = detail::getGraphComputeBatchBlockArgTypes(rewriter, resultTypes, weights, inputs);
+  auto blockArgLocs = detail::getGraphComputeBatchBlockArgLocs(loc, resultTypes, weights, inputs);
+  return createEmptySpatGraphComputeBatch(
+    rewriter, loc, resultTypes, laneCount, weights, inputs, blockArgTypes, blockArgLocs);
+}
 
-  auto* block =
-    rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), mlir::TypeRange(blockArgTypes), blockArgLocs);
-  rewriter.setInsertionPointToStart(block);
+template <typename RewriterT, typename BodyFn>
+auto createSpatGraphComputeBatch(RewriterT& rewriter,
+                                 mlir::Location loc,
+                                 mlir::TypeRange resultTypes,
+                                 int64_t laneCount,
+                                 mlir::ValueRange weights,
+                                 mlir::ValueRange inputs,
+                                 BodyFn&& body) {
+  auto batchOp = createEmptySpatGraphComputeBatch(rewriter, loc, resultTypes, laneCount, weights, inputs);
+  if (failed(batchOp))
+    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
+  auto* block = &(*batchOp).getBody().front();
 
   detail::SpatComputeBatchBodyArgs args {
     block->getArgument(0),
@@ -217,18 +286,18 @@ auto createSpatGraphComputeBatch(RewriterT& rewriter,
   using BodyResult = std::invoke_result_t<BodyFn, detail::SpatComputeBatchBodyArgs>;
   if constexpr (std::is_same_v<BodyResult, void>) {
     std::forward<BodyFn>(body)(args);
-    rewriter.setInsertionPointAfter(batchOp);
-    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(batchOp);
+    rewriter.setInsertionPointAfter(*batchOp);
+    return batchOp;
   }
   else {
     auto bodyResult = std::forward<BodyFn>(body)(args);
     if (mlir::failed(bodyResult)) {
-      rewriter.setInsertionPointAfter(batchOp);
-      rewriter.eraseOp(batchOp);
+      rewriter.setInsertionPointAfter(*batchOp);
+      rewriter.eraseOp(*batchOp);
       return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
     }
-    rewriter.setInsertionPointAfter(batchOp);
-    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(batchOp);
+    rewriter.setInsertionPointAfter(*batchOp);
+    return batchOp;
   }
 }
 

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -11,6 +11,7 @@
 #include "llvm/ADT/SmallPtrSet.h"
 
 #include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
+#include "mlir/Transforms/Passes.h"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
@@ -375,6 +376,12 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
         hasIllegalOps = true;
       }
     });
+
+    PassManager canonicalizationPM(ctx);
+    canonicalizationPM.addPass(createCanonicalizerPass());
+    if (failed(canonicalizationPM.run(moduleOp)))
+      moduleOp.emitWarning("failed to run LowerSpatialPlansPass canonicalization; continuing");
+
     if (hasIllegalOps) {
       signalPassFailure();
     } else {

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -13,6 +13,7 @@
 
 #include "Common/Common.hpp"
 #include "Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
@@ -66,9 +67,9 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
     sourceLocs.push_back(source.getLoc());
   }
 
-  auto newCompute = spatial::SpatGraphCompute::create(
-    rewriter, returnOp.getLoc(), returnOp.getOperandTypes(), funcOp.getArguments(), {}, {});
-  auto* newBlock = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), sourceTypes, sourceLocs);
+  auto newCompute = createEmptySpatGraphCompute(
+    rewriter, returnOp.getLoc(), returnOp.getOperandTypes(), {}, funcOp.getArguments(), sourceTypes, sourceLocs);
+  auto* newBlock = &newCompute.getBody().front();
   for (auto [blockArg, computeArg] : llvm::zip(newBlock->getArguments(), newCompute.getOperands()))
     mapper.map(computeArg, blockArg);
   newCompute.getProperties().setOperandSegmentSizes({0, static_cast<int>(sourceTypes.size())});
@@ -175,11 +176,6 @@ void ONNXToSpatialPass::runOnOperation() {
                                   arith::ArithDialect,
                                   scf::SCFDialect>();
 
-  PassManager cleanupPM(ctx);
-  cleanupPM.addPass(createCanonicalizerPass());
-  if (failed(cleanupPM.run(moduleOp)))
-    moduleOp.emitWarning("failed to run ONNX-to-Spatial canonicalization cleanup; continuing");
-
   annotateWeightsConstants(*entryFunc);
 
     if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
@@ -215,13 +211,18 @@ void ONNXToSpatialPass::runOnOperation() {
 
   populateEmptyFunction(*entryFunc);
 
-    if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
+  PassManager canonicalizationPM(ctx);
+  canonicalizationPM.addPass(createCanonicalizerPass());
+  if (failed(canonicalizationPM.run(moduleOp)))
+    moduleOp.emitWarning("failed to run ONNXToSpatial canonicalization; continuing");
+
+  dumpModule(moduleOp, "spatial0");
+
+  if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
       moduleOp.emitError("logical Spatial graph verification failed after ONNX-to-Spatial");
       signalPassFailure();
       return;
     }
-  dumpModule(moduleOp, "spatial0");
-
   if (failed(verifyONNXToSpatial(*entryFunc))) {
     moduleOp.emitError("ONNX-to-Spatial host legality verification failed");
     signalPassFailure();

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

@@ -917,9 +917,7 @@ struct MatMulToGemm : OpRewritePattern<ONNXMatMulOp> {
     if (failed(shapeInfo) || shapeInfo->lhsWasVector || shapeInfo->rhsWasVector)
       return failure();
 
-    const bool hasNonSingletonOutputBatch =
-      !shapeInfo->outputBatchShape.empty() && getStaticShapeElementCount(shapeInfo->outputBatchShape) != 1;
-    if (hasNonSingletonOutputBatch)
+    if (!shapeInfo->outputBatchShape.empty())
       return failure();
 
     Location loc = matmulOp.getLoc();

src/PIM/Conversion/ONNXToSpatial/Patterns/Post.cpp

@@ -10,6 +10,7 @@
 
 #include "src/Accelerators/PIM/Common/IR/WeightUtils.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/WeightMaterialization.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -128,8 +129,6 @@ struct PromoteWeightLikeComputeInputsPattern : OpRewritePattern<spatial::SpatGra
     Block& oldBlock = compute.getBody().front();
 
     rewriter.setInsertionPointAfter(compute);
-    auto newCompute = spatial::SpatGraphCompute::create(
-      rewriter, compute.getLoc(), compute.getResultTypes(), promoted->newWeights, promoted->newInputs);
     SmallVector<Type> newBlockArgTypes;
     SmallVector<Location> newBlockArgLocs;
     for (Value weight : promoted->newWeights) {
@@ -138,10 +137,14 @@ struct PromoteWeightLikeComputeInputsPattern : OpRewritePattern<spatial::SpatGra
     }
     llvm::append_range(newBlockArgTypes, promoted->newInputTypes);
     llvm::append_range(newBlockArgLocs, promoted->newInputLocs);
-    auto* newBlock = rewriter.createBlock(
-      &newCompute.getBody(), newCompute.getBody().end(), TypeRange(newBlockArgTypes), newBlockArgLocs);
-    newCompute.getProperties().setOperandSegmentSizes(
-      {static_cast<int>(promoted->newWeights.size()), static_cast<int>(promoted->newInputs.size())});
+    auto newCompute = createEmptySpatGraphCompute(rewriter,
+                                                  compute.getLoc(),
+                                                  compute.getResultTypes(),
+                                                  promoted->newWeights,
+                                                  promoted->newInputs,
+                                                  TypeRange(newBlockArgTypes),
+                                                  newBlockArgLocs);
+    auto* newBlock = &newCompute.getBody().front();
     rewriter.setInsertionPointToStart(newBlock);
 
     IRRewriter bodyRewriter(rewriter.getContext());
@@ -193,12 +196,6 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::Sp
 
     rewriter.setInsertionPointAfter(compute);
 
-    auto laneCountAttr = pim::getCheckedI32Attr(
-      rewriter, compute, static_cast<uint64_t>(compute.getLaneCount()), "promoted compute_batch lane count");
-    if (failed(laneCountAttr))
-      return failure();
-    auto newCompute = spatial::SpatGraphComputeBatch::create(
-      rewriter, compute.getLoc(), compute.getResultTypes(), *laneCountAttr, promoted->newWeights, promoted->newInputs);
     auto laneArg = compute.getLaneArgument();
     if (!laneArg)
       return rewriter.notifyMatchFailure(compute, "missing compute_batch lane block argument");
@@ -223,23 +220,30 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::Sp
       newBlockArgLocs.push_back(outputArg->getLoc());
     }
 
-    auto* newBlock = rewriter.createBlock(
-      &newCompute.getBody(), newCompute.getBody().end(), TypeRange(newBlockArgTypes), newBlockArgLocs);
-    newCompute.getProperties().setOperandSegmentSizes(
-      {static_cast<int>(promoted->newWeights.size()), static_cast<int>(promoted->newInputs.size())});
+    auto newCompute = createEmptySpatGraphComputeBatch(rewriter,
+                                                       compute.getLoc(),
+                                                       compute.getResultTypes(),
+                                                       compute.getLaneCount(),
+                                                       promoted->newWeights,
+                                                       promoted->newInputs,
+                                                       TypeRange(newBlockArgTypes),
+                                                       newBlockArgLocs);
+    if (failed(newCompute))
+      return failure();
+    auto* newBlock = &(*newCompute).getBody().front();
     rewriter.setInsertionPointToStart(newBlock);
 
     IRRewriter bodyRewriter(rewriter.getContext());
     bodyRewriter.setInsertionPointToStart(newBlock);
 
     IRMapping mapper;
-    auto newLaneArg = newCompute.getLaneArgument();
+    auto newLaneArg = (*newCompute).getLaneArgument();
     if (!newLaneArg)
       return rewriter.notifyMatchFailure(compute, "missing rewritten compute_batch lane block argument");
     mapper.map(*laneArg, *newLaneArg);
     for (auto [weightIndex, weight] : llvm::enumerate(compute.getWeights())) {
       auto oldWeightArg = compute.getWeightArgument(weightIndex);
-      auto newWeightArg = newCompute.getWeightArgument(weightIndex);
+      auto newWeightArg = (*newCompute).getWeightArgument(weightIndex);
       if (!oldWeightArg || !newWeightArg)
         return rewriter.notifyMatchFailure(compute, "missing compute_batch weight block argument during rewrite");
       mapper.map(*oldWeightArg, *newWeightArg);
@@ -249,7 +253,7 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::Sp
           *promoted,
           bodyRewriter,
           mapper,
-          [&](size_t index) { return newCompute.getInputArgument(index); },
+          [&](size_t index) { return (*newCompute).getInputArgument(index); },
           rewriter)))
       return failure();
     for (auto resultIndex : llvm::seq<size_t>(0, compute.getNumResults())) {
@@ -263,7 +267,7 @@ struct PromoteWeightLikeComputeBatchInputsPattern : OpRewritePattern<spatial::Sp
     for (Operation& op : oldBlock)
       rewriter.clone(op, mapper);
 
-    rewriter.replaceOp(compute, newCompute.getResults());
+    rewriter.replaceOp(compute, (*newCompute).getResults());
     return success();
   }
 };

src/PIM/Dialect/Spatial/CMakeLists.txt

@@ -9,7 +9,11 @@ add_pim_library(SpatialOps
   ${PIM_SRC_ROOT}/Conversion/ONNXToSpatial/CompileTime.cpp
   Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp
   Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp
-  Transforms/MergeComputeNodes/Scheduling/MergeComputeNodesPass.cpp
+  Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp
+  Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp
+  Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp
+  Transforms/MergeComputeNodes/ScheduledComputeReport.cpp
+  Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp
   Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp
   Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp
 

src/PIM/Dialect/Spatial/Spatial.td

@@ -98,6 +98,7 @@ class SpatComputeBatchLikeBase<string mnemonic> : SpatOp<mnemonic,
 }
 
 def SpatGraphComputeBatch : SpatComputeBatchLikeBase<"graph_compute_batch"> {
+  let hasCanonicalizer = 1;
   let extraClassDeclaration = [{
     std::optional<::mlir::BlockArgument> getLaneArgument();
     std::optional<::mlir::BlockArgument> getWeightArgument(unsigned idx);
@@ -114,6 +115,7 @@ def SpatGraphComputeBatch : SpatComputeBatchLikeBase<"graph_compute_batch"> {
 }
 
 def SpatScheduledComputeBatch : SpatComputeBatchLikeBase<"scheduled_compute_batch"> {
+  let hasCanonicalizer = 1;
   let extraClassDeclaration = [{
     std::optional<::mlir::BlockArgument> getLaneArgument();
     std::optional<::mlir::BlockArgument> getWeightArgument(unsigned idx);

src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp

@@ -50,64 +50,6 @@ static void printBlockArgumentList(OpAsmPrinter& printer, ArrayRef<BlockArgument
   printer << ")";
 }
 
-static void printBlockHeaderWithoutArgLocs(OpAsmPrinter& printer, Block& block) {
-  printer.printSuccessor(&block);
-  if (block.getNumArguments() == 0) {
-    printer << ":";
-    return;
-  }
-
-  printer << "(";
-  for (auto [index, argument] : llvm::enumerate(block.getArguments())) {
-    if (index != 0)
-      printer << ", ";
-    printer.printOperand(argument);
-    printer << ": ";
-    printer.printType(argument.getType());
-  }
-  printer << "):";
-}
-
-static void printRegionWithoutBlockArgLocs(OpAsmPrinter& printer,
-                                           Region& region,
-                                           bool printEntryBlockArgs = true,
-                                           bool printBlockTerminators = true,
-                                           bool printEmptyBlock = false,
-                                           bool printEntryBlockHeaderWhenMultiblock = false) {
-  printer << " {";
-  if (region.empty()) {
-    if (printEmptyBlock)
-      printer << "\n";
-    printer << "}";
-    return;
-  }
-
-  bool hasMultipleBlocks = std::next(region.begin()) != region.end();
-  for (auto [blockIndex, block] : llvm::enumerate(region)) {
-    bool printBlockHeader = blockIndex != 0 || printEntryBlockArgs
-                            || (blockIndex == 0 && printEntryBlockHeaderWhenMultiblock && hasMultipleBlocks)
-                            || (printEmptyBlock && block.empty());
-    unsigned indent = printBlockHeader ? 4u : 2u;
-
-    if (printBlockHeader) {
-      printer.getStream() << "\n";
-      printer.getStream().indent(2);
-      printBlockHeaderWithoutArgLocs(printer, block);
-    }
-
-    for (Operation& nestedOp : block) {
-      if (!printBlockTerminators && nestedOp.hasTrait<OpTrait::IsTerminator>())
-        continue;
-      printer.getStream() << "\n";
-      printer.getStream().indent(indent);
-      printer.printCustomOrGenericOp(&nestedOp);
-    }
-  }
-
-  printer.getStream() << "\n";
-  printer << "}";
-}
-
 static ParseResult parseBlockArgumentList(OpAsmParser& parser, SmallVectorImpl<OpAsmParser::Argument>& arguments) {
   if (parser.parseLParen())
     return failure();
@@ -218,9 +160,7 @@ void printComputeLikeOp(ComputeOpTy op, OpAsmPrinter& printer) {
   printer << " -> ";
   printCompressedTypeSequence(printer, op.getResultTypes());
   printer << " ";
-  printRegionWithoutBlockArgLocs(
-    printer, op.getBody(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/true,
-    /*printEmptyBlock=*/false, /*printEntryBlockHeaderWhenMultiblock=*/true);
+  printer.printRegion(op.getBody(), /*printEntryBlockArgs=*/!op.getBody().hasOneBlock());
 }
 
 template <typename ComputeOpTy>
@@ -350,9 +290,7 @@ void printComputeBatchLikeOp(ComputeBatchOpTy op, OpAsmPrinter& printer) {
   printer << " -> ";
   printCompressedTypeSequence(printer, op.getResultTypes());
   printer << " ";
-  printRegionWithoutBlockArgLocs(
-    printer, op.getBody(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/true,
-    /*printEmptyBlock=*/false, /*printEntryBlockHeaderWhenMultiblock=*/true);
+  printer.printRegion(op.getBody(), /*printEntryBlockArgs=*/!op.getBody().hasOneBlock());
 }
 
 template <typename ComputeBatchOpTy>
@@ -521,7 +459,7 @@ void SpatDeferredCommunicationOp::print(OpAsmPrinter& printer) {
   printer << " : ";
   printer.printFunctionalType(getSources().getTypes(), getOperation()->getResultTypes());
   printer << " ";
-  printRegionWithoutBlockArgLocs(printer, getBody(), /*printEntryBlockArgs=*/false);
+  printer.printRegion(getBody(), /*printEntryBlockArgs=*/false);
 }
 
 ParseResult SpatDeferredCommunicationOp::parse(OpAsmParser& parser, OperationState& result) {
@@ -763,8 +701,7 @@ ParseResult SpatScheduledComputeBatch::parse(OpAsmParser& parser, OperationState
 
 void SpatInParallelOp::print(OpAsmPrinter& printer) {
   printer << " ";
-  printRegionWithoutBlockArgLocs(
-    printer, getRegion(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/false);
+  printer.printRegion(getRegion(), /*printEntryBlockArgs=*/false, /*printBlockTerminators=*/false);
   printer.printOptionalAttrDict((*this)->getAttrs());
 }
 

src/PIM/Dialect/Spatial/SpatialOpsCanonicalization.cpp

@@ -1,8 +1,13 @@
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Block.h"
+#include "mlir/IR/IRMapping.h"
 
+#include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Support/LogicalResult.h"
 
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
@@ -40,5 +45,177 @@ LogicalResult SpatScheduledCompute::fold(FoldAdaptor adaptor, ::llvm::SmallVecto
   return foldComputeLike(*this, results);
 }
 
+template <typename ScalarComputeOpTy>
+static ScalarComputeOpTy createEmptyScalarCompute(PatternRewriter& rewriter,
+                                                  Location loc,
+                                                  TypeRange resultTypes,
+                                                  ValueRange weights,
+                                                  ValueRange inputs) {
+  auto computeOp = ScalarComputeOpTy::create(rewriter, loc, resultTypes, weights, inputs);
+  SmallVector<Type> blockArgTypes;
+  SmallVector<Location> blockArgLocs;
+  blockArgTypes.reserve(weights.size() + inputs.size());
+  blockArgLocs.reserve(weights.size() + inputs.size());
+  for (Value weight : weights) {
+    blockArgTypes.push_back(weight.getType());
+    blockArgLocs.push_back(weight.getLoc());
+  }
+  for (Value input : inputs) {
+    blockArgTypes.push_back(input.getType());
+    blockArgLocs.push_back(input.getLoc());
+  }
+  rewriter.createBlock(&computeOp.getBody(), computeOp.getBody().end(), blockArgTypes, blockArgLocs);
+  rewriter.setInsertionPointToStart(&computeOp.getBody().front());
+  return computeOp;
+}
+
+static SmallVector<OpFoldResult> remapMixedOffsets(ArrayRef<OpFoldResult> mixedOffsets, IRMapping& mapper) {
+  SmallVector<OpFoldResult> remapped;
+  remapped.reserve(mixedOffsets.size());
+  for (OpFoldResult ofr : mixedOffsets) {
+    if (auto value = dyn_cast<Value>(ofr))
+      remapped.push_back(cast<Value>(mapper.lookupOrDefault(value)));
+    else
+      remapped.push_back(cast<Attribute>(ofr));
+  }
+  return remapped;
+}
+
+static SmallVector<Value> createEmptyResults(PatternRewriter& rewriter, Location loc, TypeRange resultTypes) {
+  SmallVector<Value> resultValues;
+  resultValues.reserve(resultTypes.size());
+  for (Type resultType : resultTypes) {
+    auto tensorType = dyn_cast<RankedTensorType>(resultType);
+    if (!tensorType || !tensorType.hasStaticShape())
+      return {};
+    resultValues.push_back(tensor::EmptyOp::create(rewriter, loc, tensorType.getShape(), tensorType.getElementType()));
+  }
+  return resultValues;
+}
+
+template <typename ScalarComputeOpTy, typename ComputeBatchOpTy>
+static void copyCanonicalizedBatchAttrs(ScalarComputeOpTy compute, ComputeBatchOpTy batch, PatternRewriter& rewriter) {
+  for (NamedAttribute attr : batch->getAttrs()) {
+    if (attr.getName() == batch.getOperandSegmentSizesAttrName() || attr.getName() == batch.getLaneCountAttrName()
+        || attr.getName() == onnx_mlir::kCoreIdsAttrName)
+      continue;
+    compute->setAttr(attr.getName(), attr.getValue());
+  }
+  if constexpr (std::is_same_v<ComputeBatchOpTy, SpatScheduledComputeBatch>) {
+    if (auto coreIds = batch->template getAttrOfType<DenseI32ArrayAttr>(onnx_mlir::kCoreIdsAttrName)) {
+      assert(coreIds.size() == 1 && "single-lane scheduled compute_batch canonicalization expects exactly one core id");
+      compute->setAttr(onnx_mlir::kCoreIdAttrName, rewriter.getI32IntegerAttr(coreIds.asArrayRef().front()));
+    }
+  }
+}
+
+template <typename ComputeBatchOpTy, typename ScalarComputeOpTy>
+struct CanonicalizeSingleLaneComputeBatchPattern : OpRewritePattern<ComputeBatchOpTy> {
+  using OpRewritePattern<ComputeBatchOpTy>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ComputeBatchOpTy compute, PatternRewriter& rewriter) const override {
+    if (compute.getLaneCount() != 1)
+      return rewriter.notifyMatchFailure(compute, "lane count is not 1");
+
+    Block& oldBlock = compute.getBody().front();
+    auto oldLaneArg = compute.getLaneArgument();
+    if (!oldLaneArg)
+      return rewriter.notifyMatchFailure(compute, "missing compute_batch lane block argument");
+
+    rewriter.setInsertionPointAfter(compute);
+    auto newCompute =
+      createEmptyScalarCompute<ScalarComputeOpTy>(rewriter, compute.getLoc(), compute.getResultTypes(), compute.getWeights(), compute.getInputs());
+    copyCanonicalizedBatchAttrs(newCompute, compute, rewriter);
+    auto* newBlock = &newCompute.getBody().front();
+    rewriter.setInsertionPointToStart(newBlock);
+
+    IRMapping mapper;
+    Value zero = arith::ConstantIndexOp::create(rewriter, compute.getLoc(), 0);
+    mapper.map(*oldLaneArg, zero);
+    for (auto [index, weight] : llvm::enumerate(compute.getWeights())) {
+      auto oldArg = compute.getWeightArgument(index);
+      auto newArg = newCompute.getWeightArgument(index);
+      if (!oldArg || !newArg)
+        return rewriter.notifyMatchFailure(compute, "missing rewritten compute weight block argument");
+      mapper.map(*oldArg, *newArg);
+    }
+    for (auto [index, input] : llvm::enumerate(compute.getInputs())) {
+      auto oldArg = compute.getInputArgument(index);
+      auto newArg = newCompute.getInputArgument(index);
+      if (!oldArg || !newArg)
+        return rewriter.notifyMatchFailure(compute, "missing rewritten compute input block argument");
+      mapper.map(*oldArg, *newArg);
+    }
+
+    SmallVector<Value> resultValues = createEmptyResults(rewriter, compute.getLoc(), compute.getResultTypes());
+    if (resultValues.size() != compute.getNumResults())
+      return rewriter.notifyMatchFailure(compute, "single-lane compute_batch canonicalization requires static ranked results");
+    for (auto [index, resultValue] : llvm::enumerate(resultValues)) {
+      auto oldOutputArg = compute.getOutputArgument(index);
+      if (!oldOutputArg)
+        return rewriter.notifyMatchFailure(compute, "missing compute_batch output block argument");
+      mapper.map(*oldOutputArg, resultValue);
+    }
+
+    auto oldInParallel = dyn_cast<SpatInParallelOp>(oldBlock.getTerminator());
+    auto oldYield = dyn_cast<SpatYieldOp>(oldBlock.getTerminator());
+    for (Operation& op : oldBlock.without_terminator())
+      rewriter.clone(op, mapper);
+
+    if (oldYield) {
+      SpatYieldOp::create(rewriter, oldYield.getLoc(), ValueRange {});
+      rewriter.replaceOp(compute, newCompute.getResults());
+      return success();
+    }
+    if (!oldInParallel)
+      return rewriter.notifyMatchFailure(compute, "expected spat.in_parallel or empty spat.yield terminator");
+
+    DenseMap<BlockArgument, size_t> outputIndexByArg;
+    for (size_t index = 0; index < compute.getNumResults(); ++index) {
+      auto oldOutputArg = compute.getOutputArgument(index);
+      if (!oldOutputArg)
+        return rewriter.notifyMatchFailure(compute, "missing compute_batch output block argument");
+      outputIndexByArg[*oldOutputArg] = index;
+    }
+
+    for (Operation& op : oldInParallel.getRegion().front()) {
+      auto insertSlice = dyn_cast<tensor::ParallelInsertSliceOp>(&op);
+      if (!insertSlice)
+        return rewriter.notifyMatchFailure(compute, "expected only tensor.parallel_insert_slice in spat.in_parallel");
+      auto oldDest = dyn_cast<BlockArgument>(insertSlice.getDest());
+      if (!oldDest)
+        return rewriter.notifyMatchFailure(compute, "expected tensor.parallel_insert_slice destination to be a block argument");
+      auto resultIndexIt = outputIndexByArg.find(oldDest);
+      if (resultIndexIt == outputIndexByArg.end())
+        return rewriter.notifyMatchFailure(compute, "unexpected tensor.parallel_insert_slice destination");
+      size_t resultIndex = resultIndexIt->second;
+      Value remappedSource = mapper.lookupOrDefault(insertSlice.getSource());
+      auto remappedOffsets = remapMixedOffsets(insertSlice.getMixedOffsets(), mapper);
+      auto remappedSizes = remapMixedOffsets(insertSlice.getMixedSizes(), mapper);
+      auto remappedStrides = remapMixedOffsets(insertSlice.getMixedStrides(), mapper);
+      resultValues[resultIndex] = tensor::InsertSliceOp::create(rewriter,
+                                                                insertSlice.getLoc(),
+                                                                remappedSource,
+                                                                resultValues[resultIndex],
+                                                                remappedOffsets,
+                                                                remappedSizes,
+                                                                remappedStrides)
+                                    .getResult();
+    }
+
+    SpatYieldOp::create(rewriter, oldInParallel.getLoc(), resultValues);
+    rewriter.replaceOp(compute, newCompute.getResults());
+    return success();
+  }
+};
+
+void SpatGraphComputeBatch::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) {
+  results.add<CanonicalizeSingleLaneComputeBatchPattern<SpatGraphComputeBatch, SpatGraphCompute>>(context);
+}
+
+void SpatScheduledComputeBatch::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) {
+  results.add<CanonicalizeSingleLaneComputeBatchPattern<SpatScheduledComputeBatch, SpatScheduledCompute>>(context);
+}
+
 } // namespace spatial
 } // namespace onnx_mlir

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -847,30 +847,33 @@ LogicalResult verifyComputeBatchLikeOp(ComputeBatchOpTy batch, StringRef opName)
         return batch.emitOpError("compute_batch coreIds values must be unique");
   }
 
-  Block& block = batch.getBody().front();
-  if (block.getNumArguments() == 0)
-    return batch.emitOpError("compute_batch body must have exactly one lane block argument");
-  unsigned expectedArgCount = 1 + batch.getWeights().size() + batch.getInputs().size() + batch.getNumResults();
-  if (block.getNumArguments() != expectedArgCount)
-    return batch.emitOpError("compute_batch body block arguments must match lane, weight, input, and output operands/results");
-  auto laneArg = batch.getLaneArgument();
-  if (!laneArg || !laneArg->getType().isIndex())
-    return batch.emitOpError("compute_batch first block argument must have index type");
+  if (batch.getBody().empty())
+    return batch.emitOpError("compute_batch body must have at least one block");
 
-  for (auto [weightIndex, weight] : llvm::enumerate(batch.getWeights())) {
-    auto blockArg = batch.getWeightArgument(weightIndex);
-    if (!blockArg || blockArg->getType() != weight.getType())
-      return batch.emitOpError("compute_batch weight block argument types must match weight operand types exactly");
-  }
-  for (auto [inputIndex, input] : llvm::enumerate(batch.getInputs())) {
-    auto blockArg = batch.getInputArgument(inputIndex);
-    if (!blockArg || blockArg->getType() != input.getType())
-      return batch.emitOpError("compute_batch input block argument types must match input operand types exactly");
-  }
-  for (auto [resultIndex, resultType] : llvm::enumerate(batch.getResultTypes())) {
-    auto blockArg = batch.getOutputArgument(resultIndex);
-    if (!blockArg || blockArg->getType() != resultType)
-      return batch.emitOpError("compute_batch output block argument types must match result types exactly");
+  unsigned expectedArgCount = 1 + batch.getWeights().size() + batch.getInputs().size() + batch.getNumResults();
+  bool verifyLaneSliceOffsets = !isa<SpatScheduledComputeBatch>(batch.getOperation());
+  for (Block& block : batch.getBody()) {
+    if (block.getNumArguments() == 0)
+      return batch.emitOpError("compute_batch body must have exactly one lane block argument");
+    if (block.getNumArguments() != expectedArgCount)
+      return batch.emitOpError(
+        "compute_batch body block arguments must match lane, weight, input, and output operands/results");
+    if (!block.getArgument(0).getType().isIndex())
+      return batch.emitOpError("compute_batch first block argument must have index type");
+
+    for (auto [weightIndex, weight] : llvm::enumerate(batch.getWeights()))
+      if (block.getArgument(1 + weightIndex).getType() != weight.getType())
+        return batch.emitOpError("compute_batch weight block argument types must match weight operand types exactly");
+    for (auto [inputIndex, input] : llvm::enumerate(batch.getInputs()))
+      if (block.getArgument(1 + batch.getWeights().size() + inputIndex).getType() != input.getType())
+        return batch.emitOpError("compute_batch input block argument types must match input operand types exactly");
+    for (auto [resultIndex, resultType] : llvm::enumerate(batch.getResultTypes()))
+      if (block.getArgument(1 + batch.getWeights().size() + batch.getInputs().size() + resultIndex).getType()
+          != resultType)
+        return batch.emitOpError("compute_batch output block argument types must match result types exactly");
+
+    if (failed(verifyBatchBody(batch, block, verifyLaneSliceOffsets)))
+      return failure();
   }
 
   if (failed(verifyComputeResultsUses(batch.getOperation())))
@@ -879,7 +882,7 @@ LogicalResult verifyComputeBatchLikeOp(ComputeBatchOpTy batch, StringRef opName)
     return failure();
   if (failed(verifyOnlyConstantExternalValues(batch.getOperation(), batch.getBody(), opName)))
     return failure();
-  return verifyBatchBody(batch, block, !isa<SpatScheduledComputeBatch>(batch.getOperation()));
+  return success();
 }
 
 LogicalResult SpatGraphComputeBatch::verify() { return verifyComputeBatchLikeOp(*this, "spat.graph_compute_batch"); }

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

@@ -0,0 +1,426 @@
+#include "DeferredCommunicationPlanning.hpp"
+
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+
+namespace onnx_mlir {
+namespace 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");
+  return block.getArgument(firstArgument + std::distance(operands.begin(), it));
+}
+
+static FailureOr<Value> getOriginalProducerValue(const ProducerValueRef &producer) {
+  auto outputs = getComputeInstanceOutputValues(producer.instance);
+  if (producer.resultIndex >= outputs.size())
+    return failure();
+  return outputs[producer.resultIndex];
+}
+
+static 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 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,
+                                                    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();
+  }
+  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();
+}
+
+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);
+
+  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(),finalizeDeferredCommunication
+          sourceOperandForScheduledLaneAttr);
+        if (failed(multiSourceSelection))
+          return failure();
+        selectedSource = *multiSourceSelection;
+      } else {
+        selectedSource = block->getArgument(sourceOperandForScheduledLaneAttr.asArrayRef().front());
+      }
+    }
+  }
+
+  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();
+}
+
+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);
+}
+
+} // 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);
+    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);
+}
+
+static FailureOr<unsigned> getComputeInstanceInputIndex(const ComputeInstance &instance, Value input) {
+  auto inputs = getComputeInstanceInputs(instance);
+  auto it = llvm::find(inputs, input);
+  if (it == inputs.end())
+    return failure();
+  return static_cast<unsigned>(std::distance(inputs.begin(), it));
+}
+
+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";
+    }
+    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 (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);
+}
+
+} // namespace spatial
+} // namespace onnx_mlir

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

@@ -0,0 +1,33 @@
+#pragma once
+
+#include "ScheduledComputePlan.hpp"
+
+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);
+
+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);
+
+} // namespace spatial
+} // namespace onnx_mlir

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

@@ -0,0 +1,94 @@
+#include "ScheduledComputeMaterialization.hpp"
+#include "ScheduledComputeReport.hpp"
+#include "ScheduledComputeVerification.hpp"
+
+#include "mlir/Pass/Pass.h"
+
+#include "Scheduling/MergeSchedulingAnalysis.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Pass/PIMPasses.h"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace spatial {
+namespace {
+
+struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MergeComputeNodesPass)
+
+  StringRef getArgument() const override { return "pim-merge-compute-nodes"; }
+  StringRef getDescription() const override {
+    return "Materialize scheduled Spatial compute with deferred communication placeholders.";
+  }
+
+  void runOnOperation() override {
+    ModuleOp moduleOp = getOperation();
+    auto entryFunc = getPimEntryFunc(moduleOp);
+    if (failed(entryFunc)) {
+      moduleOp.emitError("failed to locate the PIM entry function during MergeComputeNodes");
+      signalPassFailure();
+      return;
+    }
+
+    func::FuncOp funcOp = *entryFunc;
+    MergeScheduleResult schedule = MergeSchedulingAnalysis(funcOp).getResult();
+    PatternRewriter rewriter(moduleOp.getContext());
+    FailureOr<ScheduledComputeMaterializationResult> materialization =
+      materializeScheduledCompute(funcOp, schedule, rewriter);
+    if (failed(materialization)) {
+      signalPassFailure();
+      return;
+    }
+
+    if (failed(verifyMaterializedScheduleMapping(funcOp,
+                                                 schedule,
+                                                 materialization->peftClassPlans,
+                                                 materialization->graphComputeToBlockMap,
+                                                 materialization->materializedSchedules))) {
+      moduleOp.emitError("scheduled Spatial materialization mapping verification failed");
+      signalPassFailure();
+      return;
+    }
+    if (failed(verifyDeferredTransferPhase1Invariants(funcOp))) {
+      moduleOp.emitError("scheduled Spatial deferred communication verification failed");
+      signalPassFailure();
+      return;
+    }
+    if (failed(verifyScheduledMaterializationRecords(materialization->materializedSchedules))) {
+      moduleOp.emitError("scheduled Spatial materialization record verification failed");
+      signalPassFailure();
+      return;
+    }
+    if (failed(verifyPeftMaterializationReportSummary(funcOp,
+                                                      schedule,
+                                                      materialization->peftClassPlans,
+                                                      materialization->materializedSchedules))) {
+      moduleOp.emitError("scheduled Spatial report verification failed");
+      signalPassFailure();
+      return;
+    }
+    if (failed(verifyScheduledSpatialInvariants(funcOp))) {
+      moduleOp.emitError("scheduled Spatial phase 1 verification failed");
+      signalPassFailure();
+      return;
+    }
+
+    dumpScheduledComputeReportAndModule(moduleOp,
+                                        funcOp,
+                                        schedule,
+                                        materialization->peftClassPlans,
+                                        materialization->materializedSchedules);
+    moduleOp.emitError("MergeComputeNodes stopped after spatial1_scheduled_no_comm; "
+                       "Phase 2 communication realization is not implemented");
+    signalPassFailure();
+  }
+};
+
+} // namespace
+} // namespace spatial
+
+std::unique_ptr<Pass> createMergeComputeNodesPass() { return std::make_unique<spatial::MergeComputeNodesPass>(); }
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,954 @@
+#include "ScheduledComputeMaterialization.hpp"
+#include "DeferredCommunicationPlanning.hpp"
+
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+
+#include <map>
+
+namespace onnx_mlir {
+namespace spatial {
+
+using namespace mlir;
+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) {
+  SmallVector<OpFoldResult> remapped;
+  remapped.reserve(mixedOffsets.size());
+  for (OpFoldResult ofr : mixedOffsets) {
+    if (auto value = dyn_cast<Value>(ofr))
+      remapped.push_back(cast<Value>(mapper.lookupOrDefault(value)));
+    else
+      remapped.push_back(cast<Attribute>(ofr));
+  }
+  return remapped;
+}
+
+static void appendUnique(SmallVectorImpl<Value> &values, Value value) {
+  if (!llvm::is_contained(values, value))
+    values.push_back(value);
+}
+
+static Value getBlockOperand(Block &block, ValueRange operands, Value value, unsigned firstArgument = 0) {
+  auto it = llvm::find(operands, value);
+  assert(it != operands.end() && "missing scheduled operand");
+  return block.getArgument(firstArgument + std::distance(operands.begin(), it));
+}
+
+static Value getScheduledComputeOutputArgument(Block &block, ValueRange scheduledWeights, ValueRange scheduledInputs,
+                                               ArrayRef<ProducerValueKey> carriedKeys, ProducerValueKey key) {
+  unsigned base = scheduledWeights.size() + scheduledInputs.size();
+  auto it = llvm::find(carriedKeys, key);
+  assert(it != carriedKeys.end() && "missing carried output");
+  return block.getArgument(base + std::distance(carriedKeys.begin(), it));
+}
+
+static unsigned getScheduledComputeResultArgBase(SpatScheduledCompute scheduled) {
+  return scheduled.getWeights().size() + scheduled.getInputs().size();
+}
+
+static void appendComputeBlockArguments(SmallVectorImpl<Type> &argTypes,
+                                        SmallVectorImpl<Location> &argLocs,
+                                        ValueRange weights,
+                                        ValueRange inputs,
+                                        ArrayRef<ProducerValueKey> carriedKeys,
+                                        Location loc) {
+  for (Value weight : weights)
+    argTypes.push_back(weight.getType());
+  for (Value input : inputs)
+    argTypes.push_back(input.getType());
+  for (ProducerValueKey key : carriedKeys) {
+    auto outputs = getComputeInstanceOutputValues(key.instance);
+    assert(key.resultIndex < outputs.size() && "missing carried result");
+    argTypes.push_back(outputs[key.resultIndex].getType());
+  }
+  argLocs.append(argTypes.size(), loc);
+}
+
+static Block *createScheduledComputeBlock(PatternRewriter &rewriter,
+                                          SpatScheduledCompute scheduled,
+                                          ArrayRef<ProducerValueKey> carriedKeys,
+                                          Location loc) {
+  SmallVector<Type> argTypes;
+  SmallVector<Location> argLocs;
+  appendComputeBlockArguments(argTypes, argLocs, scheduled.getWeights(), scheduled.getInputs(), carriedKeys, loc);
+  return rewriter.createBlock(&scheduled.getBody(), scheduled.getBody().end(), TypeRange(argTypes), argLocs);
+}
+
+static void appendBlockYieldBaseAndCarriedOperands(Block &block,
+                                                   unsigned baseArgCount,
+                                                   size_t carriedCount,
+                                                   SmallVectorImpl<Value> &operands) {
+  for (unsigned index = 0; index < baseArgCount; ++index)
+    operands.push_back(block.getArgument(index));
+  for (size_t index = 0; index < carriedCount; ++index)
+    operands.push_back(block.getArgument(baseArgCount + index));
+}
+
+static void createBlockYield(PatternRewriter &rewriter, Location loc, ValueRange outputs, Block *next = nullptr) {
+  OperationState state(loc, SpatBlockYieldOp::getOperationName());
+  state.addOperands(outputs);
+  if (next)
+    state.addSuccessors(next);
+  rewriter.create(state);
+}
+
+static FailureOr<BatchFragmentSpec> getBatchFragmentSpec(SpatComputeBatch batch,
+                                                         unsigned resultIndex,
+                                                         uint32_t fragmentLaneCount) {
+  auto inParallel = dyn_cast<SpatInParallelOp>(batch.getBody().front().getTerminator());
+  if (!inParallel)
+    return batch.emitOpError("scheduled materialization only supports resultful spat.graph_compute_batch");
+
+  auto outputArg = batch.getOutputArgument(resultIndex);
+  if (!outputArg)
+    return batch.emitOpError("scheduled materialization could not locate batch output block argument");
+
+  for (Operation &op : inParallel.getRegion().front()) {
+    auto insert = dyn_cast<tensor::ParallelInsertSliceOp>(&op);
+    if (!insert)
+      return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+    if (insert.getDest() != *outputArg)
+      continue;
+
+    RankedTensorType destType = insert.getDestType();
+    RankedTensorType sourceType = insert.getSourceType();
+    if (!destType || !sourceType || !destType.hasStaticShape() || !sourceType.hasStaticShape())
+      return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+    if (destType.getRank() == 0 || sourceType.getRank() == 0 || destType.getRank() != sourceType.getRank())
+      return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+    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 (!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())
+    };
+  }
+
+  return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+}
+
+
+static SourceLaneSelector buildSourceLaneSelector(PatternRewriter &rewriter,
+                                                  const ComputeStepTuple &stepTuple,
+                                                  Operation *constantAnchor,
+                                                  std::map<std::vector<uint32_t>, Value> &laneStartTableCache) {
+  if (std::optional<SourceLaneAffineMapping> affineMapping = getSourceLaneAffineMapping(stepTuple)) {
+    SourceLaneSelector selector;
+    selector.kind = SourceLaneSelector::Kind::Affine;
+    selector.affine = *affineMapping;
+    return selector;
+  }
+
+  SmallVector<uint32_t> tableValues = collectSourceLaneStarts(stepTuple);
+  std::vector<uint32_t> cacheKey(tableValues.begin(), tableValues.end());
+  auto cacheIt = laneStartTableCache.find(cacheKey);
+  if (cacheIt != laneStartTableCache.end()) {
+    SourceLaneSelector selector;
+    selector.kind = SourceLaneSelector::Kind::Table;
+    selector.table = cacheIt->second;
+    selector.tableValues = tableValues;
+    return selector;
+  }
+
+  SmallVector<int64_t> tableValuesI64;
+  tableValuesI64.reserve(tableValues.size());
+  for (uint32_t value : tableValues)
+    tableValuesI64.push_back(value);
+  Value table = createI64LookupTableConstant(rewriter, constantAnchor, tableValuesI64);
+  laneStartTableCache.emplace(std::move(cacheKey), table);
+  SourceLaneSelector selector;
+  selector.kind = SourceLaneSelector::Kind::Table;
+  selector.table = table;
+  selector.tableValues = tableValues;
+  return selector;
+}
+
+static FailureOr<Value> buildSourceLaneStartForScheduledLane(OpBuilder &builder,
+                                                             Location loc,
+                                                             Value scheduledLane,
+                                                             const SourceLaneSelector &selector,
+                                                             Operation *constantAnchor) {
+  (void)constantAnchor;
+  if (selector.kind == SourceLaneSelector::Kind::Affine) {
+    if (selector.affine.baseLaneStart == 0 && selector.affine.laneCount == 1)
+      return scheduledLane;
+    AffineExpr d0 = builder.getAffineDimExpr(0);
+    AffineExpr expr = d0;
+    if (selector.affine.laneCount != 1)
+      expr = d0 * selector.affine.laneCount;
+    if (selector.affine.baseLaneStart != 0)
+      expr = expr + selector.affine.baseLaneStart;
+    return affine::AffineApplyOp::create(
+             builder, loc, AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0, expr), ValueRange {scheduledLane})
+      .getResult();
+  }
+
+  if (!selector.table)
+    return failure();
+  Value sourceLaneStartI64 =
+    tensor::ExtractOp::create(builder, loc, selector.table, ValueRange {scheduledLane}).getResult();
+  return arith::IndexCastOp::create(builder, loc, builder.getIndexType(), sourceLaneStartI64).getResult();
+}
+
+static LogicalResult verifyPeftClassPlan(Operation *diagnosticAnchor,
+                                         const PeftClassPlan &peftClassPlan,
+                                         const MergeScheduleResult &schedule) {
+  if (peftClassPlan.cpus.empty())
+    return diagnosticAnchor->emitOpError("PEFT materialization class has no CPUs");
+
+  SmallVector<const SmallVector<ComputeInstance> *> schedules;
+  for (size_t cpu : peftClassPlan.cpus) {
+    auto it = peftClassPlan.instancesByCpu.find(cpu);
+    if (it == peftClassPlan.instancesByCpu.end())
+      return diagnosticAnchor->emitOpError("PEFT materialization class is missing a per-CPU schedule");
+    schedules.push_back(&it->second);
+    for (const ComputeInstance &instance : it->second)
+      if (!schedule.computeToCpuSlotMap.count(instance))
+        return diagnosticAnchor->emitOpError("PEFT materialization class references a compute instance without a scheduler position");
+  }
+
+  if (peftClassPlan.cpus.size() == 1)
+    return success();
+
+  auto emitNonIso = [&](size_t stepPosition) -> LogicalResult {
+    std::string cpus;
+    llvm::raw_string_ostream os(cpus);
+    llvm::interleaveComma(peftClassPlan.cpus, os, [&](size_t cpu) { os << cpu; });
+    diagnosticAnchor->emitOpError("PEFT equivalence class has non-isomorphic per-CPU schedules")
+      << " class " << peftClassPlan.canonicalClassId << " cpus [" << os.str() << "] step " << stepPosition;
+    return failure();
+  };
+
+  size_t tupleCount = schedules.front()->size();
+  for (const SmallVector<ComputeInstance> *cpuSchedule : schedules)
+    if (cpuSchedule->size() != tupleCount)
+      return emitNonIso(0);
+
+  for (size_t stepPosition = 0; stepPosition < tupleCount; ++stepPosition) {
+    const ComputeInstance &reference = (*schedules.front())[stepPosition];
+    bool refIsScalar = isa<SpatCompute>(reference.op);
+    for (size_t cpuIndex = 1; cpuIndex < schedules.size(); ++cpuIndex) {
+      const ComputeInstance &instance = (*schedules[cpuIndex])[stepPosition];
+      if (instance.op != reference.op || instance.laneCount != reference.laneCount)
+        return emitNonIso(stepPosition);
+      if (isa<SpatCompute>(instance.op) != refIsScalar)
+        return emitNonIso(stepPosition);
+    }
+  }
+
+  return success();
+}
+
+static LogicalResult collectPeftClassOperandsAndResults(PeftClassPlan &peftClassPlan,
+                                                        const MergeScheduleResult &schedule) {
+  peftClassPlan.weights.clear();
+  peftClassPlan.inputs.clear();
+  peftClassPlan.resultTypes.clear();
+
+  if (peftClassPlan.cpus.size() == 1) {
+    size_t cpu = peftClassPlan.cpus.front();
+    for (const ComputeInstance &instance : peftClassPlan.instancesByCpu.lookup(cpu)) {
+      if (auto compute = dyn_cast<SpatCompute>(instance.op)) {
+        llvm::append_range(peftClassPlan.resultTypes, compute.getResultTypes());
+      } else {
+        auto batch = cast<SpatComputeBatch>(instance.op);
+        for (unsigned resultIndex = 0; resultIndex < batch.getNumResults(); ++resultIndex) {
+          auto spec = getBatchFragmentSpec(batch, resultIndex, instance.laneCount);
+          if (failed(spec))
+            return failure();
+          peftClassPlan.resultTypes.push_back(spec->resultType);
+        }
+      }
+
+      for (Value weight : getComputeInstanceWeights(instance))
+        appendUnique(peftClassPlan.weights, weight);
+      for (Value input : getComputeInstanceInputs(instance))
+        if (!getProducerValueRef(input, &instance))
+          appendUnique(peftClassPlan.inputs, input);
+    }
+    return success();
+  }
+
+  for (const ScheduledStepPlan &stepPlan : buildScheduledStepPlans(peftClassPlan)) {
+    const ComputeStepTuple &stepTuple = stepPlan.stepTuple;
+    const ComputeInstance &representative = stepTuple.instances.front();
+    if (auto compute = dyn_cast<SpatCompute>(representative.op)) {
+      for (Type type : compute.getResultTypes()) {
+        auto tensorType = dyn_cast<RankedTensorType>(type);
+        if (!tensorType || !tensorType.hasStaticShape())
+          return compute.emitOpError("scheduled materialization only supports static ranked tensor scalar results");
+        SmallVector<int64_t> shape;
+        shape.push_back(static_cast<int64_t>(peftClassPlan.cpus.size()));
+        llvm::append_range(shape, tensorType.getShape());
+        peftClassPlan.resultTypes.push_back(RankedTensorType::get(shape, tensorType.getElementType()));
+      }
+    } else {
+      auto batch = cast<SpatComputeBatch>(representative.op);
+      uint32_t totalLanes = static_cast<uint32_t>(peftClassPlan.cpus.size()) * representative.laneCount;
+      for (unsigned resultIndex = 0; resultIndex < batch.getNumResults(); ++resultIndex) {
+        auto spec = getBatchFragmentSpec(batch, resultIndex, totalLanes);
+        if (failed(spec))
+          return failure();
+        peftClassPlan.resultTypes.push_back(spec->resultType);
+      }
+    }
+
+    for (const ComputeInstance &instance : stepTuple.instances) {
+      for (Value weight : getComputeInstanceWeights(instance))
+        appendUnique(peftClassPlan.weights, weight);
+      for (Value input : getComputeInstanceInputs(instance))
+        if (!getProducerValueRef(input, &instance))
+          appendUnique(peftClassPlan.inputs, input);
+    }
+  }
+
+  return success();
+}
+
+static void cloneComputeBody(OpBuilder &builder, Block &source, IRMapping &mapper, SmallVectorImpl<Value> &yieldedValues) {
+  for (Operation &op : source.without_terminator())
+    builder.clone(op, mapper);
+  auto yield = cast<SpatYieldOp>(source.getTerminator());
+  for (Value output : yield.getOutputs())
+    yieldedValues.push_back(mapper.lookup(output));
+}
+
+static LogicalResult materializeResultfulBatchChunkAsScalar(PatternRewriter &rewriter,
+                                                            SpatComputeBatch batch,
+                                                            const ComputeInstance &instance,
+                                                            ValueRange scheduledWeights,
+                                                            ValueRange scheduledInputs,
+                                                            Block &block,
+                                                            const MergeScheduleResult &schedule,
+                                                            uint64_t &exchangeId,
+                                                            SmallVectorImpl<Value> &yieldedValues) {
+  SmallVector<Value> initResults;
+  SmallVector<BatchFragmentSpec> fragmentSpecs;
+  for (unsigned resultIndex = 0; resultIndex < batch.getNumResults(); ++resultIndex) {
+    auto spec = getBatchFragmentSpec(batch, resultIndex, instance.laneCount);
+    if (failed(spec))
+      return failure();
+    fragmentSpecs.push_back(*spec);
+    auto empty = createEmptyTensorForType(rewriter, batch.getLoc(), spec->resultType);
+    if (failed(empty))
+      return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+    initResults.push_back(*empty);
+  }
+
+  Value lower = getOrCreateIndexConstant(rewriter, batch.getOperation(), instance.laneStart);
+  Value upper = getOrCreateIndexConstant(rewriter, batch.getOperation(), instance.laneStart + instance.laneCount);
+  Value step = getOrCreateIndexConstant(rewriter, batch.getOperation(), 1);
+  auto loop = buildNormalizedScfFor(
+    rewriter,
+    batch.getLoc(),
+    lower,
+    upper,
+    step,
+    initResults,
+    [&](OpBuilder &builder, Location bodyLoc, Value originalLane, ValueRange iterArgs, SmallVectorImpl<Value> &yielded) -> LogicalResult {
+
+      IRMapping mapper;
+      mapper.map(*batch.getLaneArgument(), originalLane);
+      for (auto [index, weight] : llvm::enumerate(batch.getWeights()))
+        mapper.map(*batch.getWeightArgument(index), getBlockOperand(block, scheduledWeights, weight));
+      for (auto [index, input] : llvm::enumerate(batch.getInputs())) {
+        Value mapped;
+        if (failed(mapSingleCpuInput(builder,
+                                     input.getLoc(),
+                                     input,
+                                     instance,
+                                     schedule,
+                                     scheduledInputs,
+                                     block,
+                                     scheduledWeights.size(),
+                                     ArrayRef<ProducerValueKey> {},
+                                     exchangeId,
+                                     mapped)))
+          return failure();
+        mapper.map(*batch.getInputArgument(index), mapped);
+      }
+      for (auto [index, outputArg] : llvm::enumerate(batch.getOutputs()))
+        (void)outputArg, mapper.map(*batch.getOutputArgument(index), iterArgs[index]);
+
+      Block &source = batch.getBody().front();
+      for (Operation &op : source.without_terminator())
+        builder.clone(op, mapper);
+
+      auto inParallel = dyn_cast<SpatInParallelOp>(source.getTerminator());
+      if (!inParallel)
+          return batch.emitOpError("expected spat.in_parallel in resultful spat.graph_compute_batch"), failure();
+      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;
+
+      SmallVector<Value> current(iterArgs.begin(), iterArgs.end());
+      for (Operation &op : inParallel.getRegion().front()) {
+        auto insert = dyn_cast<tensor::ParallelInsertSliceOp>(&op);
+        if (!insert)
+          return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+        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();
+        size_t resultIndex = outputIndexByArg.lookup(oldDest);
+        SmallVector<OpFoldResult> offsets = remapMixedOffsets(insert.getMixedOffsets(), mapper);
+        offsets[fragmentSpecs[resultIndex].laneDim] = localLane;
+        current[resultIndex] = tensor::InsertSliceOp::create(builder,
+                                                             insert.getLoc(),
+                                                             mapper.lookup(insert.getSource()),
+                                                             current[resultIndex],
+                                                             offsets,
+                                                             remapMixedOffsets(insert.getMixedSizes(), mapper),
+                                                             remapMixedOffsets(insert.getMixedStrides(), mapper))
+                                 .getResult();
+      }
+      llvm::append_range(yielded, current);
+      return success();
+    });
+  if (failed(loop))
+    return failure();
+  llvm::append_range(yieldedValues, loop->results);
+  return success();
+}
+
+static LogicalResult materializeSingleCpuPeftClass(
+  PatternRewriter &rewriter,
+  SpatScheduledCompute scheduled,
+  const PeftClassPlan &peftClassPlan,
+  const MergeScheduleResult &schedule,
+  DenseMap<GraphComputeBlockKey, Block *> &graphComputeToBlockMap,
+  ScheduledMaterializationRecord &record,
+  uint64_t &exchangeId) {
+  size_t cpu = peftClassPlan.cpus.front();
+  auto instancesIt = peftClassPlan.instancesByCpu.find(cpu);
+  assert(instancesIt != peftClassPlan.instancesByCpu.end() && "missing single-cpu schedule");
+  const SmallVector<ComputeInstance> &instances = instancesIt->second;
+
+  SmallVector<ProducerValueKey> carriedKeys;
+  Block *block = nullptr;
+  for (auto [ordinal, instance] : llvm::enumerate(instances)) {
+    if (!block)
+      block = createScheduledComputeBlock(rewriter, scheduled, carriedKeys, instance.op->getLoc());
+
+    GraphComputeBlockKey key = getGraphComputeBlockKey(instance);
+    graphComputeToBlockMap[key] = block;
+    record.computeKeys.push_back(key);
+    record.blocks.push_back(block);
+
+    rewriter.setInsertionPointToStart(block);
+    SmallVector<Value> yieldedValues;
+    if (auto compute = dyn_cast<SpatCompute>(instance.op)) {
+      IRMapping mapper;
+      for (auto [index, weight] : llvm::enumerate(compute.getWeights()))
+        mapper.map(*compute.getWeightArgument(index), getBlockOperand(*block, scheduled.getWeights(), weight));
+      for (auto [index, input] : llvm::enumerate(compute.getInputs())) {
+        Value mapped;
+        if (failed(mapSingleCpuInput(rewriter,
+                                     input.getLoc(),
+                                     input,
+                                     instance,
+                                     schedule,
+                                     scheduled.getInputs(),
+                                     *block,
+                                     scheduled.getWeights().size(),
+                                     carriedKeys,
+                                     exchangeId,
+                                     mapped)))
+          return failure();
+        mapper.map(*compute.getInputArgument(index), mapped);
+      }
+      cloneComputeBody(rewriter, compute.getBody().front(), mapper, yieldedValues);
+    } else {
+      auto batch = cast<SpatComputeBatch>(instance.op);
+      if (failed(materializeResultfulBatchChunkAsScalar(rewriter,
+                                                        batch,
+                                                        instance,
+                                                        scheduled.getWeights(),
+                                                        scheduled.getInputs(),
+                                                        *block,
+                                                        schedule,
+                                                        exchangeId,
+                                                        yieldedValues)))
+        return failure();
+    }
+
+    SmallVector<ProducerValueKey> currentKeys;
+    for (size_t index = 0; index < yieldedValues.size(); ++index)
+      currentKeys.push_back({instance, index});
+    unsigned baseArgCount = getScheduledComputeResultArgBase(scheduled);
+    SmallVector<Value> blockYieldOperands;
+    bool hasNextBlock = ordinal + 1 < instances.size();
+    if (hasNextBlock) {
+      SmallVector<ProducerValueKey> nextCarriedKeys(carriedKeys);
+      llvm::append_range(nextCarriedKeys, currentKeys);
+      Block *nextBlock = createScheduledComputeBlock(rewriter, scheduled, nextCarriedKeys, instance.op->getLoc());
+      appendBlockYieldBaseAndCarriedOperands(*block, baseArgCount, carriedKeys.size(), blockYieldOperands);
+      llvm::append_range(blockYieldOperands, yieldedValues);
+      rewriter.setInsertionPointToEnd(block);
+      createBlockYield(rewriter, instance.op->getLoc(), blockYieldOperands, nextBlock);
+      carriedKeys = std::move(nextCarriedKeys);
+      block = nextBlock;
+    } else {
+      for (ProducerValueKey carried : carriedKeys)
+        blockYieldOperands.push_back(getScheduledComputeOutputArgument(*block,
+                                                                      scheduled.getWeights(),
+                                                                      scheduled.getInputs(),
+                                                                      carriedKeys,
+                                                                      carried));
+      llvm::append_range(blockYieldOperands, yieldedValues);
+      rewriter.setInsertionPointToEnd(block);
+      createBlockYield(rewriter, instance.op->getLoc(), blockYieldOperands);
+    }
+  }
+  return success();
+}
+
+// Builds offsets for inserting one per-CPU local fragment into the
+// scheduled_compute_batch output. The lane offset is in scheduled-output
+// lane space, not local fragment lane space.
+static SmallVector<OpFoldResult> buildScheduledOutputInsertOffsets(OpBuilder &builder,
+                                                                   Location loc,
+                                                                   Value scheduledLane,
+                                                                   int64_t lanesPerScheduledLane,
+                                                                   RankedTensorType localFragmentType,
+                                                                   unsigned laneDim) {
+  SmallVector<OpFoldResult> offsets;
+  Value scheduledOutputLane = scheduledLane;
+  if (lanesPerScheduledLane != 1) {
+    scheduledOutputLane =
+      affine::AffineApplyOp::create(builder,
+                                    loc,
+                                    AffineMap::get(/*dimCount=*/1,
+                                                   /*symbolCount=*/0,
+                                                   builder.getAffineDimExpr(0) * lanesPerScheduledLane),
+                                    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)));
+  return offsets;
+}
+
+static LogicalResult materializeMultiCpuPeftClass(
+  PatternRewriter &rewriter,
+  SpatScheduledComputeBatch scheduled,
+  const PeftClassPlan &peftClassPlan,
+  const MergeScheduleResult &schedule,
+  DenseMap<GraphComputeBlockKey, Block *> &graphComputeToBlockMap,
+  ScheduledMaterializationRecord &record,
+  uint64_t &exchangeId) {
+  std::map<std::vector<uint32_t>, Value> laneStartTableCache;
+  ArrayRef<ScheduledStepPlan> stepPlans = record.stepPlans;
+  for (const ScheduledStepPlan &stepPlan : stepPlans) {
+    const ComputeStepTuple &stepTuple = stepPlan.stepTuple;
+    SourceLaneSelector sourceLaneSelector =
+      buildSourceLaneSelector(rewriter, stepTuple, scheduled.getOperation(), laneStartTableCache);
+    SmallVector<Type> blockArgTypes {rewriter.getIndexType()};
+    SmallVector<Location> blockArgLocs {scheduled.getLoc()};
+    for (Value weight : scheduled.getWeights()) {
+      blockArgTypes.push_back(weight.getType());
+      blockArgLocs.push_back(weight.getLoc());
+    }
+    for (Value input : scheduled.getInputs()) {
+      blockArgTypes.push_back(input.getType());
+      blockArgLocs.push_back(input.getLoc());
+    }
+    for (Type resultType : scheduled.getResultTypes()) {
+      blockArgTypes.push_back(resultType);
+      blockArgLocs.push_back(scheduled.getLoc());
+    }
+    Block *block = rewriter.createBlock(&scheduled.getBody(), scheduled.getBody().end(), blockArgTypes, blockArgLocs);
+    for (const ComputeInstance &instance : stepTuple.instances) {
+      GraphComputeBlockKey key = getGraphComputeBlockKey(instance);
+      graphComputeToBlockMap[key] = block;
+      record.computeKeys.push_back(key);
+      record.blocks.push_back(block);
+    }
+
+    rewriter.setInsertionPointToStart(block);
+    Value scheduledLane = block->getArgument(0);
+    const ComputeInstance &representative = stepTuple.instances.front();
+    SmallVector<Value> finalLocalFragments;
+    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();
+      for (auto [index, input] : llvm::enumerate(compute.getInputs())) {
+        Value mapped;
+        if (failed(mapMultiCpuTupleInput(rewriter,
+                                         input.getLoc(),
+                                         input,
+                                         stepTuple,
+                                         peftClassPlan,
+                                         schedule,
+                                         scheduled.getInputs(),
+                                         *block,
+                                         firstInputArg,
+                                         exchangeId,
+                                         mapped)))
+          return failure();
+        mapper.map(*compute.getInputArgument(index), mapped);
+      }
+      SmallVector<Value> yieldedValues;
+      cloneComputeBody(rewriter, compute.getBody().front(), mapper, yieldedValues);
+      for (Value yielded : yieldedValues) {
+        auto tensorType = dyn_cast<RankedTensorType>(yielded.getType());
+        if (!tensorType || !tensorType.hasStaticShape() || tensorType.getRank() == 0)
+          return compute.emitOpError("scheduled materialization only supports static ranked tensor scalar step results");
+        SmallVector<ReassociationIndices> reassociation;
+        reassociation.push_back({0, 1});
+        for (int64_t dim = 1; dim < tensorType.getRank(); ++dim)
+          reassociation.push_back({static_cast<int64_t>(dim + 1)});
+        SmallVector<int64_t> expandedShape {1};
+        llvm::append_range(expandedShape, tensorType.getShape());
+        finalLocalFragments.push_back(tensor::ExpandShapeOp::create(rewriter,
+                                                                    scheduled.getLoc(),
+                                                                    RankedTensorType::get(expandedShape, tensorType.getElementType()),
+                                                                    yielded,
+                                                                    reassociation)
+                                       .getResult());
+        finalLaneDims.push_back(0);
+      }
+    } else {
+      auto batch = cast<SpatComputeBatch>(representative.op);
+      SmallVector<Value> localFragments;
+      SmallVector<BatchFragmentSpec> fragmentSpecs;
+      for (unsigned resultIndex = 0; resultIndex < batch.getNumResults(); ++resultIndex) {
+        auto spec = getBatchFragmentSpec(batch, resultIndex, representative.laneCount);
+        if (failed(spec))
+          return failure();
+        fragmentSpecs.push_back(*spec);
+        auto empty = createEmptyTensorForType(rewriter, batch.getLoc(), spec->resultType);
+        if (failed(empty))
+          return failure();
+        localFragments.push_back(*empty);
+      }
+
+      Value lower = getOrCreateIndexConstant(rewriter, scheduled.getOperation(), 0);
+      Value upper = getOrCreateIndexConstant(rewriter, scheduled.getOperation(), representative.laneCount);
+      Value step = getOrCreateIndexConstant(rewriter, scheduled.getOperation(), 1);
+      auto loop = buildNormalizedScfFor(
+        rewriter,
+        batch.getLoc(),
+        lower,
+        upper,
+        step,
+        localFragments,
+        [&](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,
+                                          AffineMap::get(/*dimCount=*/2,
+                                                         /*symbolCount=*/0,
+                                                         builder.getAffineDimExpr(0) + builder.getAffineDimExpr(1)),
+                                          ValueRange {*sourceLaneStart, innerLane})
+              .getResult();
+          mapper.map(*batch.getLaneArgument(), sourceLane);
+          for (auto [index, weight] : llvm::enumerate(batch.getWeights()))
+            mapper.map(*batch.getWeightArgument(index), block->getArgument(1 + index));
+          unsigned firstInputArg = 1 + scheduled.getWeights().size();
+          for (auto [index, input] : llvm::enumerate(batch.getInputs())) {
+            Value mapped;
+            if (failed(mapMultiCpuTupleInput(builder,
+                                             input.getLoc(),
+                                             input,
+                                             stepTuple,
+                                             peftClassPlan,
+                                             schedule,
+                                             scheduled.getInputs(),
+                                             *block,
+                                             firstInputArg,
+                                             exchangeId,
+                                             mapped)))
+              return failure();
+            mapper.map(*batch.getInputArgument(index), mapped);
+          }
+          for (unsigned index = 0; index < batch.getNumResults(); ++index)
+            mapper.map(*batch.getOutputArgument(index), iterArgs[index]);
+          for (Operation &op : batch.getBody().front().without_terminator())
+            builder.clone(op, mapper);
+
+          auto inParallel = dyn_cast<SpatInParallelOp>(batch.getBody().front().getTerminator());
+          if (!inParallel)
+            return batch.emitOpError("expected spat.in_parallel in resultful spat.graph_compute_batch"), failure();
+
+          DenseMap<BlockArgument, size_t> outputIndexByArg;
+          for (size_t index = 0; index < batch.getNumResults(); ++index)
+            outputIndexByArg[*batch.getOutputArgument(index)] = index;
+
+          SmallVector<Value> current(iterArgs.begin(), iterArgs.end());
+          for (Operation &op : inParallel.getRegion().front()) {
+            auto insert = dyn_cast<tensor::ParallelInsertSliceOp>(&op);
+            if (!insert)
+              return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+            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();
+            size_t resultIndex = outputIndexByArg.lookup(oldDest);
+            SmallVector<OpFoldResult> offsets = remapMixedOffsets(insert.getMixedOffsets(), mapper);
+            offsets[fragmentSpecs[resultIndex].laneDim] = innerLane;
+            current[resultIndex] = tensor::InsertSliceOp::create(builder,
+                                                                 insert.getLoc(),
+                                                                 mapper.lookup(insert.getSource()),
+                                                                 current[resultIndex],
+                                                                 offsets,
+                                                                 remapMixedOffsets(insert.getMixedSizes(), mapper),
+                                                                 remapMixedOffsets(insert.getMixedStrides(), mapper))
+                                     .getResult();
+          }
+          llvm::append_range(yielded, current);
+          return success();
+        });
+      if (failed(loop))
+        return failure();
+      finalLocalFragments.assign(loop->results.begin(), loop->results.end());
+      for (const BatchFragmentSpec &spec : fragmentSpecs)
+        finalLaneDims.push_back(spec.laneDim);
+    }
+
+    auto inParallel = SpatInParallelOp::create(rewriter, scheduled.getLoc());
+    rewriter.setInsertionPointToStart(&inParallel.getRegion().front());
+    for (auto [resultIndex, localFragment] : llvm::enumerate(finalLocalFragments)) {
+      auto localFragmentType = cast<RankedTensorType>(localFragment.getType());
+      int64_t lanesPerScheduledLane = isa<SpatCompute>(representative.op) ? 1 : representative.laneCount;
+      SmallVector<OpFoldResult> offsets = buildScheduledOutputInsertOffsets(
+        rewriter,
+        scheduled.getLoc(),
+        scheduledLane,
+        lanesPerScheduledLane,
+        localFragmentType,
+        finalLaneDims[resultIndex]);
+      SmallVector<OpFoldResult> sizes;
+      SmallVector<OpFoldResult> strides;
+      for (int64_t dim : localFragmentType.getShape()) {
+        sizes.push_back(rewriter.getIndexAttr(dim));
+        strides.push_back(rewriter.getIndexAttr(1));
+      }
+      tensor::ParallelInsertSliceOp::create(
+        rewriter,
+        scheduled.getLoc(),
+        localFragment,
+        block->getArgument(getScheduledBatchResultArgBase(scheduled) + stepPlan.resultOffset + resultIndex),
+        offsets,
+        sizes,
+        strides);
+    }
+  }
+  return success();
+}
+
+
+} // namespace
+
+
+FailureOr<ScheduledComputeMaterializationResult>
+materializeScheduledCompute(func::FuncOp funcOp,
+                            const MergeScheduleResult &schedule,
+                            PatternRewriter &rewriter) {
+  llvm::MapVector<size_t, PeftClassPlan> peftClassPlans;
+  for (const ComputeInstance &instance : schedule.dominanceOrderCompute) {
+    size_t cpu = schedule.computeToCpuMap.lookup(instance);
+    size_t canonicalPeftClassId = getCanonicalPeftClassId(cpu, schedule);
+    auto &peftClassPlan = peftClassPlans[canonicalPeftClassId];
+    peftClassPlan.canonicalClassId = canonicalPeftClassId;
+    if (!llvm::is_contained(peftClassPlan.cpus, cpu))
+      peftClassPlan.cpus.push_back(cpu);
+    peftClassPlan.instancesByCpu[cpu].push_back(instance);
+  }
+  for (auto &entry : peftClassPlans) {
+    PeftClassPlan &peftClassPlan = entry.second;
+    llvm::sort(peftClassPlan.cpus);
+    for (size_t cpu : peftClassPlan.cpus)
+      llvm::sort(peftClassPlan.instancesByCpu[cpu], [&](const ComputeInstance &lhs, const ComputeInstance &rhs) {
+        return schedule.computeToCpuSlotMap.lookup(lhs) < schedule.computeToCpuSlotMap.lookup(rhs);
+      });
+    if (failed(verifyPeftClassPlan(funcOp.getOperation(), peftClassPlan, schedule)))
+      return failure();
+    if (failed(collectPeftClassOperandsAndResults(peftClassPlan, schedule)))
+      return failure();
+  }
+
+  Operation *insertionPoint = funcOp.getBody().front().getTerminator();
+  DenseMap<GraphComputeBlockKey, Block *> graphComputeToBlockMap;
+  DenseMap<size_t, SpatScheduledCompute> scheduledComputes;
+  DenseMap<size_t, SpatScheduledComputeBatch> scheduledComputeBatches;
+  DenseMap<size_t, size_t> classToRecordIndex;
+  std::vector<ScheduledMaterializationRecord> materializedSchedules;
+  uint64_t exchangeId = 0;
+
+  for (auto &entry : peftClassPlans) {
+    PeftClassPlan &peftClassPlan = entry.second;
+    rewriter.setInsertionPoint(insertionPoint);
+
+    ScheduledMaterializationRecord record;
+    record.canonicalPeftClassId = peftClassPlan.canonicalClassId;
+    record.cpus = peftClassPlan.cpus;
+    record.stepPlans = buildScheduledStepPlans(peftClassPlan);
+
+    if (peftClassPlan.cpus.size() == 1) {
+      auto scheduled = SpatScheduledCompute::create(
+        rewriter,
+        peftClassPlan.instancesByCpu.lookup(peftClassPlan.cpus.front()).front().op->getLoc(),
+        TypeRange(peftClassPlan.resultTypes),
+        peftClassPlan.weights,
+        peftClassPlan.inputs);
+      scheduled->setAttr(kCoreIdAttrName, rewriter.getI32IntegerAttr(static_cast<int32_t>(peftClassPlan.cpus.front())));
+      scheduled->setAttr("scheduled.peft_cpus", rewriter.getDenseI64ArrayAttr(toI64Array(peftClassPlan.cpus)));
+      SmallVector<Attribute> stepSources;
+      SmallVector<Attribute> sourceLaneSelectors;
+      SmallVector<int64_t> stepResultOffsets;
+      SmallVector<int64_t> stepResultCounts;
+      SmallVector<int64_t> sourceLaneStarts;
+      SmallVector<int64_t> sourceLaneCounts;
+      size_t resultOffset = 0;
+      for (const ComputeInstance &instance : peftClassPlan.instancesByCpu.lookup(peftClassPlan.cpus.front())) {
+        stepSources.push_back(rewriter.getStringAttr(getInstanceName(instance)));
+        sourceLaneSelectors.push_back(rewriter.getStringAttr(isa<SpatCompute>(instance.op) ? "scalar" : "affine"));
+        size_t resultCount = getComputeInstanceResultValueCount(instance);
+        stepResultOffsets.push_back(static_cast<int64_t>(resultOffset));
+        stepResultCounts.push_back(static_cast<int64_t>(resultCount));
+        resultOffset += resultCount;
+        if (isa<SpatCompute>(instance.op)) {
+          sourceLaneStarts.push_back(0);
+          sourceLaneCounts.push_back(0);
+        } else {
+          sourceLaneStarts.push_back(instance.laneStart);
+          sourceLaneCounts.push_back(instance.laneCount);
+        }
+      }
+      scheduled->setAttr("scheduled.step_sources", rewriter.getArrayAttr(stepSources));
+      scheduled->setAttr("scheduled.step_result_offsets", rewriter.getDenseI64ArrayAttr(stepResultOffsets));
+      scheduled->setAttr("scheduled.step_result_counts", rewriter.getDenseI64ArrayAttr(stepResultCounts));
+      scheduled->setAttr("scheduled.source_lane_starts", rewriter.getDenseI64ArrayAttr(sourceLaneStarts));
+      scheduled->setAttr("scheduled.source_lane_counts", rewriter.getDenseI64ArrayAttr(sourceLaneCounts));
+      scheduled->setAttr("scheduled.source_lane_selector", rewriter.getArrayAttr(sourceLaneSelectors));
+      record.scheduledOp = scheduled.getOperation();
+      scheduledComputes[peftClassPlan.canonicalClassId] = scheduled;
+    } else {
+      auto scheduled = SpatScheduledComputeBatch::create(rewriter,
+                                                         peftClassPlan.instancesByCpu.lookup(peftClassPlan.cpus.front()).front().op->getLoc(),
+                                                         TypeRange(peftClassPlan.resultTypes),
+                                                         rewriter.getI32IntegerAttr(static_cast<int32_t>(peftClassPlan.cpus.size())),
+                                                         peftClassPlan.weights,
+                                                         peftClassPlan.inputs);
+      scheduled->setAttr(kCoreIdsAttrName, rewriter.getDenseI32ArrayAttr(toI32Array(peftClassPlan.cpus)));
+      scheduled->setAttr("scheduled.peft_cpus", rewriter.getDenseI64ArrayAttr(toI64Array(peftClassPlan.cpus)));
+      SmallVector<Attribute> stepSources;
+      SmallVector<Attribute> sourceLaneSelectors;
+      SmallVector<int64_t> resultOffsets;
+      SmallVector<int64_t> resultCounts;
+      SmallVector<int64_t> sourceLaneStarts;
+      SmallVector<int64_t> sourceLaneCounts;
+      for (const ScheduledStepPlan &stepPlan : record.stepPlans) {
+        stepSources.push_back(rewriter.getStringAttr(getInstanceName(stepPlan.stepTuple.instances.front())));
+        sourceLaneSelectors.push_back(rewriter.getStringAttr(usesAffineSourceLaneMapping(stepPlan.stepTuple) ? "affine" : "table"));
+        resultOffsets.push_back(static_cast<int64_t>(stepPlan.resultOffset));
+        resultCounts.push_back(static_cast<int64_t>(stepPlan.resultCount));
+        for (const ComputeInstance &instance : stepPlan.stepTuple.instances) {
+          sourceLaneStarts.push_back(instance.laneStart);
+          sourceLaneCounts.push_back(instance.laneCount);
+        }
+      }
+      RankedTensorType sourceLaneTableType = RankedTensorType::get(
+        {static_cast<int64_t>(record.stepPlans.size()), static_cast<int64_t>(peftClassPlan.cpus.size())},
+        rewriter.getI64Type());
+      scheduled->setAttr("scheduled.step_sources", rewriter.getArrayAttr(stepSources));
+      scheduled->setAttr("scheduled.step_result_offsets", rewriter.getDenseI64ArrayAttr(resultOffsets));
+      scheduled->setAttr("scheduled.step_result_counts", rewriter.getDenseI64ArrayAttr(resultCounts));
+      scheduled->setAttr("scheduled.source_lane_starts", DenseElementsAttr::get(sourceLaneTableType, ArrayRef<int64_t>(sourceLaneStarts)));
+      scheduled->setAttr("scheduled.source_lane_counts", DenseElementsAttr::get(sourceLaneTableType, ArrayRef<int64_t>(sourceLaneCounts)));
+      scheduled->setAttr("scheduled.source_lane_selector", rewriter.getArrayAttr(sourceLaneSelectors));
+      record.scheduledOp = scheduled.getOperation();
+      scheduledComputeBatches[peftClassPlan.canonicalClassId] = scheduled;
+    }
+
+    classToRecordIndex[peftClassPlan.canonicalClassId] = materializedSchedules.size();
+    materializedSchedules.push_back(std::move(record));
+  }
+
+  for (auto &entry : peftClassPlans) {
+    PeftClassPlan &peftClassPlan = entry.second;
+    ScheduledMaterializationRecord &record =
+      materializedSchedules[classToRecordIndex.lookup(peftClassPlan.canonicalClassId)];
+    if (peftClassPlan.cpus.size() == 1) {
+      if (failed(materializeSingleCpuPeftClass(rewriter,
+                                               scheduledComputes.lookup(peftClassPlan.canonicalClassId),
+                                               peftClassPlan,
+                                               schedule,
+                                               graphComputeToBlockMap,
+                                               record,
+                                               exchangeId)))
+        return failure();
+    } else {
+      if (failed(materializeMultiCpuPeftClass(rewriter,
+                                              scheduledComputeBatches.lookup(peftClassPlan.canonicalClassId),
+                                              peftClassPlan,
+                                              schedule,
+                                              graphComputeToBlockMap,
+                                              record,
+                                              exchangeId)))
+        return failure();
+    }
+  }
+
+  return ScheduledComputeMaterializationResult {std::move(peftClassPlans), std::move(materializedSchedules), std::move(graphComputeToBlockMap)};
+}
+
+
+} // namespace spatial
+} // namespace onnx_mlir

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

@@ -0,0 +1,20 @@
+#pragma once
+
+#include "ScheduledComputePlan.hpp"
+
+namespace onnx_mlir {
+namespace spatial {
+
+struct ScheduledComputeMaterializationResult {
+  llvm::MapVector<size_t, PeftClassPlan> peftClassPlans;
+  std::vector<ScheduledMaterializationRecord> materializedSchedules;
+  DenseMap<GraphComputeBlockKey, Block *> graphComputeToBlockMap;
+};
+
+FailureOr<ScheduledComputeMaterializationResult>
+materializeScheduledCompute(func::FuncOp funcOp,
+                            const MergeScheduleResult &schedule,
+                            PatternRewriter &rewriter);
+
+} // namespace spatial
+} // namespace onnx_mlir

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

@@ -0,0 +1,337 @@
+#pragma once
+
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/AsmState.h"
+#include "mlir/IR/IRMapping.h"
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/raw_ostream.h"
+
+#include <limits>
+#include <optional>
+#include <string>
+#include <vector>
+
+#include "Scheduling/ComputeInstanceUtils.hpp"
+#include "Scheduling/MergeSchedulingAnalysis.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+
+namespace onnx_mlir {
+namespace spatial {
+
+using namespace mlir;
+
+struct ProducerValueKey {
+  ComputeInstance instance;
+  size_t resultIndex = 0;
+
+  bool operator==(const ProducerValueKey &other) const {
+    return instance == other.instance && resultIndex == other.resultIndex;
+  }
+};
+
+struct GraphComputeBlockKey {
+  Operation *op = nullptr;
+  uint32_t laneStart = 0;
+  uint32_t laneCount = 1;
+
+  bool operator==(const GraphComputeBlockKey &other) const {
+    return op == other.op && laneStart == other.laneStart && laneCount == other.laneCount;
+  }
+};
+
+struct PeftClassPlan {
+  size_t canonicalClassId = 0;
+  SmallVector<size_t> cpus;
+  llvm::MapVector<size_t, SmallVector<ComputeInstance>> instancesByCpu;
+
+  SmallVector<Value> weights;
+  SmallVector<Value> inputs;
+  SmallVector<Type> resultTypes;
+};
+
+struct ComputeStepTuple {
+  SmallVector<ComputeInstance> instances;
+};
+
+struct ScheduledStepPlan {
+  ComputeStepTuple stepTuple;
+  size_t stepIndex = 0;
+  size_t resultOffset = 0;
+  size_t resultCount = 0;
+};
+
+struct SourceLaneAffineMapping {
+  uint32_t baseLaneStart = 0;
+  uint32_t laneCount = 1;
+};
+
+struct SourceLaneSelector {
+  enum class Kind { Affine, Table } kind = Kind::Affine;
+  SourceLaneAffineMapping affine;
+  Value table;
+  SmallVector<uint32_t> tableValues;
+};
+
+struct 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;
+  SmallVector<size_t> cpus;
+  SmallVector<ScheduledStepPlan> stepPlans;
+  SmallVector<GraphComputeBlockKey> computeKeys;
+  SmallVector<Block *> blocks;
+};
+
+struct ScheduledComputePrintContext {
+  mlir::AsmState asmState;
+  explicit ScheduledComputePrintContext(ModuleOp module, const OpPrintingFlags &flags = OpPrintingFlags())
+    : asmState(module.getOperation(), flags) {}
+};
+
+inline GraphComputeBlockKey getGraphComputeBlockKey(const ComputeInstance &instance) {
+  return {instance.op, instance.laneStart, instance.laneCount};
+}
+
+inline SmallVector<size_t> getPeftClassCpus(size_t cpu, const MergeScheduleResult &schedule) {
+  llvm::SmallDenseSet<size_t, 8> seen;
+  SmallVector<size_t> cpus;
+  auto append = [&](size_t value) {
+    if (seen.insert(value).second)
+      cpus.push_back(value);
+  };
+  append(cpu);
+  if (auto it = schedule.equivalentClass.find(cpu); it != schedule.equivalentClass.end())
+    for (size_t equivalentCpu : it->second)
+      append(equivalentCpu);
+  llvm::sort(cpus);
+  return cpus;
+}
+
+inline size_t getCanonicalPeftClassId(size_t cpu, const MergeScheduleResult &schedule) {
+  SmallVector<size_t> cpus = getPeftClassCpus(cpu, schedule);
+  return cpus.empty() ? cpu : cpus.front();
+}
+
+inline size_t getScheduledCpuForComputeInstance(const ComputeInstance &instance, const MergeScheduleResult &schedule) {
+  if (auto it = schedule.computeToCpuMap.find(instance); it != schedule.computeToCpuMap.end())
+    return it->second;
+
+  auto batch = dyn_cast<SpatComputeBatch>(instance.op);
+  assert(batch && instance.laneCount != 0 && "missing scheduled CPU for non-batch compute instance");
+  assert(instance.laneStart < static_cast<uint32_t>(batch.getLaneCount()) && "batch lane start out of range");
+  ComputeInstance chunk = getBatchChunkForLane(batch, instance.laneStart);
+  auto it = schedule.computeToCpuMap.find(chunk);
+  assert(it != schedule.computeToCpuMap.end() && "missing scheduled CPU for batch chunk");
+  return it->second;
+}
+
+inline std::string getInstanceName(const ComputeInstance &instance) {
+  return llvm::formatv("{0}[lanes={1}:{2}]",
+                       instance.op->getName().getStringRef(),
+                       instance.laneStart,
+                       instance.laneStart + instance.laneCount)
+    .str();
+}
+
+inline SmallVector<int64_t> toI64Array(ArrayRef<size_t> values) {
+  SmallVector<int64_t> converted;
+  converted.reserve(values.size());
+  for (size_t value : values)
+    converted.push_back(static_cast<int64_t>(value));
+  return converted;
+}
+
+inline SmallVector<int32_t> toI32Array(ArrayRef<size_t> values) {
+  SmallVector<int32_t> converted;
+  converted.reserve(values.size());
+  for (size_t value : values)
+    converted.push_back(static_cast<int32_t>(value));
+  return converted;
+}
+
+inline unsigned getScheduledBatchResultArgBase(SpatScheduledComputeBatch scheduled) {
+  unsigned weightArgBase = 1;
+  unsigned inputArgBase = weightArgBase + scheduled.getWeights().size();
+  return inputArgBase + scheduled.getInputs().size();
+}
+
+inline SmallVector<GraphComputeBlockKey> collectExpectedGraphComputeBlockKeys(func::FuncOp funcOp) {
+  SmallVector<GraphComputeBlockKey> keys;
+  for (Operation &op : funcOp.getOps()) {
+    if (auto compute = dyn_cast<SpatGraphCompute>(&op))
+      keys.push_back(getGraphComputeBlockKey({compute.getOperation(), 0, 1}));
+    else if (auto batch = dyn_cast<SpatGraphComputeBatch>(&op))
+      for (ComputeInstance chunk : getBatchChunksForRange(batch, 0, static_cast<uint32_t>(batch.getLaneCount())))
+        keys.push_back(getGraphComputeBlockKey(chunk));
+  }
+  return keys;
+}
+
+inline size_t countPeftEquivalenceClasses(const MergeScheduleResult &schedule) {
+  llvm::SmallDenseSet<size_t, 16> classes;
+  for (const ComputeInstance &instance : schedule.dominanceOrderCompute) {
+    auto cpuIt = schedule.computeToCpuMap.find(instance);
+    if (cpuIt == schedule.computeToCpuMap.end())
+      continue;
+    classes.insert(getCanonicalPeftClassId(cpuIt->second, schedule));
+  }
+  return classes.size();
+}
+
+inline SmallVector<ComputeStepTuple> buildComputeStepTuples(const PeftClassPlan &peftClassPlan) {
+  SmallVector<ComputeStepTuple> stepTuples;
+  if (peftClassPlan.cpus.empty())
+    return stepTuples;
+  size_t stepCount = peftClassPlan.instancesByCpu.lookup(peftClassPlan.cpus.front()).size();
+  stepTuples.resize(stepCount);
+  for (size_t stepIndex = 0; stepIndex < stepCount; ++stepIndex)
+    for (size_t cpu : peftClassPlan.cpus)
+      stepTuples[stepIndex].instances.push_back(peftClassPlan.instancesByCpu.lookup(cpu)[stepIndex]);
+  return stepTuples;
+}
+
+inline size_t getComputeInstanceResultValueCount(const ComputeInstance &instance) {
+  return TypeSwitch<Operation *, size_t>(instance.op)
+    .Case<SpatCompute>([](SpatCompute compute) { return compute.getNumResults(); })
+    .Case<SpatComputeBatch>([](SpatComputeBatch batch) { return batch.getNumResults(); })
+    .Default([](Operation *) -> size_t {
+      llvm_unreachable("expected graph compute or graph compute batch");
+    });
+}
+
+inline SmallVector<ScheduledStepPlan> buildScheduledStepPlans(const PeftClassPlan &peftClassPlan) {
+  SmallVector<ScheduledStepPlan> stepPlans;
+  size_t resultOffset = 0;
+  for (auto [stepIndex, stepTuple] : llvm::enumerate(buildComputeStepTuples(peftClassPlan))) {
+    assert(!stepTuple.instances.empty() && "expected non-empty step tuple");
+    size_t resultCount = getComputeInstanceResultValueCount(stepTuple.instances.front());
+    stepPlans.push_back(ScheduledStepPlan {stepTuple, stepIndex, resultOffset, resultCount});
+    resultOffset += resultCount;
+  }
+  return stepPlans;
+}
+
+inline bool valueTransitivelyDependsOn(Value value, Value dependency) {
+  SmallVector<Value> worklist {value};
+  DenseSet<Value> visited;
+  while (!worklist.empty()) {
+    Value current = worklist.pop_back_val();
+    if (!visited.insert(current).second)
+      continue;
+    if (current == dependency)
+      return true;
+    Operation *def = current.getDefiningOp();
+    if (!def)
+      continue;
+    for (Value operand : def->getOperands())
+      worklist.push_back(operand);
+  }
+  return false;
+}
+
+inline std::optional<SourceLaneAffineMapping> getSourceLaneAffineMapping(const ComputeStepTuple &stepTuple) {
+  if (stepTuple.instances.empty())
+    return std::nullopt;
+  const ComputeInstance &reference = stepTuple.instances.front();
+  for (const ComputeInstance &instance : stepTuple.instances) {
+    if (instance.op != reference.op || instance.laneCount != reference.laneCount)
+      return std::nullopt;
+  }
+
+  for (size_t index = 0; index < stepTuple.instances.size(); ++index) {
+    uint32_t expectedLaneStart = reference.laneStart + static_cast<uint32_t>(index) * reference.laneCount;
+    if (stepTuple.instances[index].laneStart != expectedLaneStart)
+      return std::nullopt;
+  }
+
+  return SourceLaneAffineMapping {reference.laneStart, reference.laneCount};
+}
+
+inline bool usesAffineSourceLaneMapping(const ComputeStepTuple &stepTuple) {
+  return getSourceLaneAffineMapping(stepTuple).has_value();
+}
+
+inline SmallVector<uint32_t> collectSourceLaneStarts(const ComputeStepTuple &stepTuple) {
+  SmallVector<uint32_t> sourceLaneStarts;
+  sourceLaneStarts.reserve(stepTuple.instances.size());
+  for (const ComputeInstance &instance : stepTuple.instances)
+    sourceLaneStarts.push_back(instance.laneStart);
+  return sourceLaneStarts;
+}
+
+inline Value createI64LookupTableConstant(OpBuilder &builder, Operation *constantAnchor, ArrayRef<int64_t> values) {
+  OpBuilder::InsertionGuard guard(builder);
+  builder.setInsertionPoint(constantAnchor);
+  RankedTensorType tableType = RankedTensorType::get({static_cast<int64_t>(values.size())}, builder.getI64Type());
+  DenseElementsAttr tableAttr = DenseElementsAttr::get(tableType, values);
+  return arith::ConstantOp::create(builder, constantAnchor->getLoc(), tableType, tableAttr).getResult();
+}
+
+inline FailureOr<Value> createEmptyTensorForType(OpBuilder &builder, Location loc, Type type) {
+  auto tensorType = dyn_cast<RankedTensorType>(type);
+  if (!tensorType || !tensorType.hasStaticShape())
+    return failure();
+  return tensor::EmptyOp::create(builder, loc, tensorType.getShape(), tensorType.getElementType()).getResult();
+}
+
+} // namespace spatial
+} // namespace onnx_mlir
+
+namespace llvm {
+template <>
+struct DenseMapInfo<onnx_mlir::spatial::ProducerValueKey> {
+  static onnx_mlir::spatial::ProducerValueKey getEmptyKey() {
+    return {DenseMapInfo<onnx_mlir::spatial::ComputeInstance>::getEmptyKey(), std::numeric_limits<size_t>::max()};
+  }
+  static onnx_mlir::spatial::ProducerValueKey getTombstoneKey() {
+    return {DenseMapInfo<onnx_mlir::spatial::ComputeInstance>::getTombstoneKey(), std::numeric_limits<size_t>::max()};
+  }
+  static unsigned getHashValue(const onnx_mlir::spatial::ProducerValueKey &key) {
+    return hash_combine(DenseMapInfo<onnx_mlir::spatial::ComputeInstance>::getHashValue(key.instance), key.resultIndex);
+  }
+  static bool isEqual(const onnx_mlir::spatial::ProducerValueKey &lhs,
+                      const onnx_mlir::spatial::ProducerValueKey &rhs) {
+    return lhs == rhs;
+  }
+};
+
+template <>
+struct DenseMapInfo<onnx_mlir::spatial::GraphComputeBlockKey> {
+  static onnx_mlir::spatial::GraphComputeBlockKey getEmptyKey() {
+    return {DenseMapInfo<mlir::Operation *>::getEmptyKey(), UINT32_MAX, UINT32_MAX};
+  }
+  static onnx_mlir::spatial::GraphComputeBlockKey getTombstoneKey() {
+    return {DenseMapInfo<mlir::Operation *>::getTombstoneKey(), UINT32_MAX, UINT32_MAX - 1};
+  }
+  static unsigned getHashValue(const onnx_mlir::spatial::GraphComputeBlockKey &key) {
+    return hash_combine(key.op, key.laneStart, key.laneCount);
+  }
+  static bool isEqual(const onnx_mlir::spatial::GraphComputeBlockKey &lhs,
+                      const onnx_mlir::spatial::GraphComputeBlockKey &rhs) {
+    return lhs == rhs;
+  }
+};
+} // namespace llvm

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

@@ -0,0 +1,312 @@
+#include "ScheduledComputeReport.hpp"
+
+#include "llvm/Support/raw_os_ostream.h"
+
+#include <fstream>
+
+#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
+#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
+
+namespace onnx_mlir {
+namespace spatial {
+
+using namespace mlir;
+namespace {
+
+static std::string formatValueLabel(Value value, AsmState &asmState) {
+  std::string storage;
+  llvm::raw_string_ostream os(storage);
+  value.printAsOperand(os, asmState);
+  return storage;
+}
+
+static std::string formatOperationLabel(Operation *op, AsmState &asmState) {
+  if (op->getNumResults() == 0)
+    return op->getName().getStringRef().str();
+  std::string storage;
+  llvm::raw_string_ostream os(storage);
+  llvm::interleaveComma(op->getResults(), os, [&](Value result) { os << formatValueLabel(result, asmState); });
+  return os.str();
+}
+
+static std::string formatGraphComputeBlockKey(const GraphComputeBlockKey &key, AsmState &asmState) {
+  return llvm::formatv("{0} {1}", formatOperationLabel(key.op, asmState), key.op->getName().getStringRef()).str();
+}
+
+static std::string formatComputeInstanceForReport(const ComputeInstance &instance, AsmState &asmState) {
+  std::string opLabel = formatGraphComputeBlockKey(getGraphComputeBlockKey(instance), asmState);
+  if (isa<SpatCompute>(instance.op))
+    return opLabel;
+  return llvm::formatv("{0} sourceLanes [{1}:{2}]",
+                       opLabel,
+                       instance.laneStart,
+                       instance.laneStart + instance.laneCount)
+    .str();
+}
+
+template <typename T>
+static void printIndexedList(raw_ostream &os, ArrayRef<T> values) {
+  os << "[";
+  llvm::interleaveComma(llvm::enumerate(values), os, [&](auto indexedValue) {
+    os << indexedValue.index() << ":" << indexedValue.value();
+  });
+  os << "]";
+}
+
+struct PeftMaterializationReportSummary {
+  size_t scalarGraphCompute = 0;
+  size_t graphComputeBatchOps = 0;
+  size_t scalarGraphComputeInstances = 0;
+  size_t graphComputeBatchInstances = 0;
+  size_t peftClasses = 0;
+  size_t singleCpuClasses = 0;
+  size_t multiCpuClasses = 0;
+  size_t scheduledCompute = 0;
+  size_t scheduledComputeBatch = 0;
+  size_t deferredCommunication = 0;
+  size_t deferredCommunicationScalarMetadata = 0;
+  size_t deferredCommunicationScheduledLaneMetadata = 0;
+  size_t deferredCommunicationMultiSourcePayloads = 0;
+};
+
+static PeftMaterializationReportSummary buildPeftMaterializationReportSummary(
+  func::FuncOp funcOp,
+  const MergeScheduleResult &schedule,
+  const llvm::MapVector<size_t, PeftClassPlan> &peftClassPlans,
+  ArrayRef<ScheduledMaterializationRecord> materializedSchedules) {
+  PeftMaterializationReportSummary summary;
+  for (Operation &op : funcOp.getOps()) {
+    if (isa<SpatGraphCompute>(op))
+      summary.scalarGraphCompute++;
+    else if (isa<SpatGraphComputeBatch>(op)) {
+      summary.graphComputeBatchOps++;
+    }
+  }
+  for (const ComputeInstance &instance : schedule.dominanceOrderCompute)
+    (isa<SpatCompute>(instance.op) ? summary.scalarGraphComputeInstances : summary.graphComputeBatchInstances)++;
+  summary.peftClasses = peftClassPlans.size();
+  for (const auto &entry : peftClassPlans)
+    (entry.second.cpus.size() == 1 ? summary.singleCpuClasses : summary.multiCpuClasses)++;
+  for (const ScheduledMaterializationRecord &record : materializedSchedules) {
+    if (isa<SpatScheduledCompute>(record.scheduledOp))
+      summary.scheduledCompute++;
+    else
+      summary.scheduledComputeBatch++;
+  }
+  funcOp.walk([&](SpatDeferredCommunicationOp transfer) {
+    summary.deferredCommunication++;
+    if (auto batchedAttr = transfer->getAttrOfType<BoolAttr>("batched"); batchedAttr && batchedAttr.getValue())
+      summary.deferredCommunicationScheduledLaneMetadata++;
+    else
+      summary.deferredCommunicationScalarMetadata++;
+    if (transfer.getSources().size() > 1)
+      summary.deferredCommunicationMultiSourcePayloads++;
+  });
+  return summary;
+}
+
+} // namespace
+
+LogicalResult verifyPeftMaterializationReportSummary(func::FuncOp funcOp,
+                                                            const MergeScheduleResult &schedule,
+                                                            const llvm::MapVector<size_t, PeftClassPlan> &peftClassPlans,
+                                                            ArrayRef<ScheduledMaterializationRecord> materializedSchedules) {
+  PeftMaterializationReportSummary summary =
+    buildPeftMaterializationReportSummary(funcOp, schedule, peftClassPlans, materializedSchedules);
+  pim::CappedDiagnosticReporter diagnostics;
+  if (summary.peftClasses != peftClassPlans.size())
+    diagnostics.report(funcOp.getOperation(), [&](Operation *illegalOp) {
+      illegalOp->emitOpError() << "phase-check report PEFT total " << summary.peftClasses
+                               << " does not match classes.size() " << peftClassPlans.size();
+    });
+  if (summary.scalarGraphComputeInstances + summary.graphComputeBatchInstances != schedule.dominanceOrderCompute.size())
+    diagnostics.report(funcOp.getOperation(), [&](Operation *illegalOp) {
+      illegalOp->emitOpError() << "phase-check report total compute instances "
+                               << (summary.scalarGraphComputeInstances + summary.graphComputeBatchInstances)
+                               << " does not match schedule size " << schedule.dominanceOrderCompute.size();
+    });
+  if (summary.scheduledCompute + summary.scheduledComputeBatch != materializedSchedules.size())
+    diagnostics.report(funcOp.getOperation(), [&](Operation *illegalOp) {
+      illegalOp->emitOpError() << "phase-check report scheduled total "
+                               << (summary.scheduledCompute + summary.scheduledComputeBatch)
+                               << " does not match materialized scheduled ops " << materializedSchedules.size();
+    });
+  diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial report verification failed");
+  return success(!diagnostics.hasFailure());
+}
+
+namespace {
+
+static std::string formatStepResultRange(size_t resultOffset, size_t resultCount) {
+  if (resultCount == 1)
+    return llvm::formatv("result[{0}]", resultOffset).str();
+  return llvm::formatv("result[{0}:{1}]", resultOffset, resultOffset + resultCount).str();
+}
+
+static void printMultiSourceDeferredInputs(raw_ostream &os, Block &block) {
+  unsigned deferredInputIndex = 0;
+  for (Operation &op : block.getOperations()) {
+    auto transfer = dyn_cast<SpatDeferredCommunicationOp>(&op);
+    if (!transfer)
+      continue;
+    auto multiSourcePayload = transfer->getAttrOfType<BoolAttr>("multi_source_payload");
+    auto sourceOperandForScheduledLane =
+      transfer->getAttrOfType<DenseI64ArrayAttr>("source_operand_for_scheduled_lane");
+    if (multiSourcePayload && multiSourcePayload.getValue() && sourceOperandForScheduledLane) {
+      SmallVector<size_t> sourceOperandIndexes;
+      for (int64_t sourceOperandIndex : sourceOperandForScheduledLane.asArrayRef())
+        sourceOperandIndexes.push_back(static_cast<size_t>(sourceOperandIndex));
+      os << "    deferred input " << deferredInputIndex << ": multi-source uniqueSources="
+         << transfer.getSources().size() << " sourceOperandForScheduledLane=";
+      printIndexedList(os, ArrayRef<size_t>(sourceOperandIndexes));
+      os << "\n";
+    }
+    deferredInputIndex++;
+  }
+}
+
+static void dumpPeftMaterializationReport(ModuleOp moduleOp,
+                                          func::FuncOp funcOp,
+                                          const MergeScheduleResult &schedule,
+                                          const llvm::MapVector<size_t, PeftClassPlan> &peftClassPlans,
+                                          ArrayRef<ScheduledMaterializationRecord> materializedSchedules,
+                                          ScheduledComputePrintContext &printContext) {
+  std::fstream file = openDialectDumpFileWithExtension("spatial2_scheduled_no_comm", "/reports", "txt");
+  if (!file.is_open())
+    return;
+
+  llvm::raw_os_ostream os(file);
+  AsmState &asmState = printContext.asmState;
+  PeftMaterializationReportSummary summary =
+    buildPeftMaterializationReportSummary(funcOp, schedule, peftClassPlans, materializedSchedules);
+
+  os << "Summary\n";
+  os << "=======\n";
+  os << "Graph computes:\n";
+  os << "  total: " << (summary.scalarGraphCompute + summary.graphComputeBatchOps) << "\n";
+  os << "  scalar graph_compute: " << summary.scalarGraphCompute << "\n";
+  os << "  graph_compute_batch: " << summary.graphComputeBatchOps << "\n";
+  os << "Compute instances:\n";
+  os << "  total: " << (summary.scalarGraphComputeInstances + summary.graphComputeBatchInstances) << "\n";
+  os << "  scalar graph_compute instances: " << summary.scalarGraphComputeInstances << "\n";
+  os << "  graph_compute_batch instances: " << summary.graphComputeBatchInstances << "\n";
+  os << "PEFT classes:\n";
+  os << "  total: " << summary.peftClasses << "\n";
+  os << "  single-cpu: " << summary.singleCpuClasses << "\n";
+  os << "  multi-cpu: " << summary.multiCpuClasses << "\n";
+  os << "Scheduled ops:\n";
+  os << "  total: " << (summary.scheduledCompute + summary.scheduledComputeBatch) << "\n";
+  os << "  scheduled_compute: " << summary.scheduledCompute << "\n";
+  os << "  scheduled_compute_batch: " << summary.scheduledComputeBatch << "\n";
+  os << "Deferred communications:\n";
+  os << "  total: " << summary.deferredCommunication << "\n";
+  os << "  scalar metadata: " << summary.deferredCommunicationScalarMetadata << "\n";
+  os << "  scheduled-lane metadata: " << summary.deferredCommunicationScheduledLaneMetadata << "\n";
+  os << "  multi-source payloads: " << summary.deferredCommunicationMultiSourcePayloads << "\n\n";
+
+  os << "PEFT Classes\n";
+  os << "============\n";
+  for (const auto &entry : peftClassPlans) {
+    const PeftClassPlan &peftClassPlan = entry.second;
+    os << "C" << peftClassPlan.canonicalClassId << " "
+       << (peftClassPlan.cpus.size() == 1 ? "single-cpu" : "multi-cpu") << " PEFT class\n";
+    if (peftClassPlan.cpus.size() == 1) {
+      size_t cpu = peftClassPlan.cpus.front();
+      os << "  cpu: " << cpu << "\n";
+      os << "  steps: " << peftClassPlan.instancesByCpu.lookup(cpu).size() << "\n";
+      for (auto [stepIndex, instance] : llvm::enumerate(peftClassPlan.instancesByCpu.lookup(cpu)))
+        os << "  step " << stepIndex << ": " << formatComputeInstanceForReport(instance, asmState) << "\n";
+    } else {
+      os << "  scheduled lanes: " << peftClassPlan.cpus.size() << "\n";
+      os << "  steps: " << peftClassPlan.instancesByCpu.lookup(peftClassPlan.cpus.front()).size() << "\n";
+      os << "  cpus by scheduled lane:\n";
+      os << "    ";
+      printIndexedList(os, ArrayRef<size_t>(peftClassPlan.cpus));
+      os << "\n";
+      os << "  step sources:\n";
+      for (auto [stepIndex, stepTuple] : llvm::enumerate(buildComputeStepTuples(peftClassPlan)))
+        os << "    step " << stepIndex << ": "
+           << formatGraphComputeBlockKey(getGraphComputeBlockKey(stepTuple.instances.front()), asmState) << "\n";
+    }
+    os << "\n";
+  }
+
+  os << "Materialized Scheduled Ops\n";
+  os << "=========================\n";
+  for (const ScheduledMaterializationRecord &record : materializedSchedules) {
+    os << "C" << record.canonicalPeftClassId << " -> " << formatOperationLabel(record.scheduledOp, asmState) << " "
+       << record.scheduledOp->getName().getStringRef() << "\n";
+    os << "  kind: "
+       << (isa<SpatScheduledCompute>(record.scheduledOp) ? "single-cpu scheduled_compute"
+                                                         : "multi-cpu scheduled_compute_batch")
+       << "\n";
+    if (isa<SpatScheduledCompute>(record.scheduledOp))
+      os << "  cpu: " << record.cpus.front() << "\n";
+    else
+      os << "  scheduled lanes: " << record.cpus.size() << "\n";
+    os << "  results: " << record.scheduledOp->getNumResults() << "\n";
+    os << "  steps: "
+       << (isa<SpatScheduledCompute>(record.scheduledOp)
+             ? peftClassPlans.lookup(record.canonicalPeftClassId).instancesByCpu.lookup(record.cpus.front()).size()
+             : record.stepPlans.size())
+       << "\n";
+    if (isa<SpatScheduledComputeBatch>(record.scheduledOp)) {
+      os << "  cpus by scheduled lane:\n";
+      os << "    ";
+      printIndexedList(os, ArrayRef<size_t>(record.cpus));
+      os << "\n\n";
+    }
+    if (isa<SpatScheduledCompute>(record.scheduledOp)) {
+      const PeftClassPlan &peftClassPlan = peftClassPlans.lookup(record.canonicalPeftClassId);
+      size_t cpu = peftClassPlan.cpus.front();
+      size_t resultOffset = 0;
+      for (auto [stepIndex, instance] : llvm::enumerate(peftClassPlan.instancesByCpu.lookup(cpu))) {
+        size_t resultCount = getComputeInstanceResultValueCount(instance);
+        os << "  step " << stepIndex << " " << formatStepResultRange(resultOffset, resultCount) << " "
+           << formatComputeInstanceForReport(instance, asmState) << "\n";
+        resultOffset += resultCount;
+      }
+    } else {
+      auto scheduledBatch = cast<SpatScheduledComputeBatch>(record.scheduledOp);
+      for (auto [stepIndex, stepPlan] : llvm::enumerate(record.stepPlans)) {
+        const ComputeInstance &representative = stepPlan.stepTuple.instances.front();
+        SmallVector<uint32_t> sourceLaneStarts = collectSourceLaneStarts(stepPlan.stepTuple);
+        os << "  step " << stepIndex << " " << formatStepResultRange(stepPlan.resultOffset, stepPlan.resultCount) << " "
+           << formatGraphComputeBlockKey(getGraphComputeBlockKey(representative), asmState)
+           << " lanesPerScheduledLane=" << representative.laneCount << " sourceLaneSelector="
+           << (usesAffineSourceLaneMapping(stepPlan.stepTuple) ? "affine" : "table") << "\n";
+        os << "    source lanes by scheduled lane:\n";
+        os << "      ";
+        printIndexedList(os, ArrayRef<uint32_t>(sourceLaneStarts));
+        os << "\n";
+        Block &stepBlock = *std::next(scheduledBatch.getBody().begin(), stepIndex);
+        printMultiSourceDeferredInputs(os, stepBlock);
+      }
+    }
+    os << "\n";
+  }
+}
+
+
+} // namespace
+
+void dumpScheduledComputeReportAndModule(ModuleOp moduleOp,
+                                         func::FuncOp funcOp,
+                                         const MergeScheduleResult &schedule,
+                                         const llvm::MapVector<size_t, PeftClassPlan> &peftClassPlans,
+                                         ArrayRef<ScheduledMaterializationRecord> materializedSchedules) {
+  OpPrintingFlags flags;
+  flags.elideLargeElementsAttrs().enableDebugInfo(false, false).assumeVerified();
+  ScheduledComputePrintContext printContext(moduleOp, flags);
+  dumpPeftMaterializationReport(moduleOp, funcOp, schedule, peftClassPlans, materializedSchedules, printContext);
+
+  std::fstream file = openDialectDumpFileWithExtension("spatial2_scheduled_no_comm", "/dialects", "mlir");
+  if (!file.is_open())
+    return;
+  llvm::raw_os_ostream os(file);
+  moduleOp.getOperation()->print(os, printContext.asmState);
+  os.flush();
+}
+
+} // namespace spatial
+} // namespace onnx_mlir

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

@@ -0,0 +1,21 @@
+#pragma once
+
+#include "ScheduledComputeMaterialization.hpp"
+
+namespace onnx_mlir {
+namespace spatial {
+
+LogicalResult verifyPeftMaterializationReportSummary(
+  func::FuncOp funcOp,
+  const MergeScheduleResult &schedule,
+  const llvm::MapVector<size_t, PeftClassPlan> &peftClassPlans,
+  ArrayRef<ScheduledMaterializationRecord> materializedSchedules);
+
+void dumpScheduledComputeReportAndModule(ModuleOp moduleOp,
+                                         func::FuncOp funcOp,
+                                         const MergeScheduleResult &schedule,
+                                         const llvm::MapVector<size_t, PeftClassPlan> &peftClassPlans,
+                                         ArrayRef<ScheduledMaterializationRecord> materializedSchedules);
+
+} // namespace spatial
+} // namespace onnx_mlir

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

@@ -0,0 +1,316 @@
+#include "ScheduledComputeVerification.hpp"
+
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
+
+namespace onnx_mlir {
+namespace spatial {
+
+using namespace mlir;
+
+LogicalResult verifyMaterializedScheduleMapping(
+  func::FuncOp funcOp,
+  const MergeScheduleResult &schedule,
+  const llvm::MapVector<size_t, PeftClassPlan> &peftClassPlans,
+  const DenseMap<GraphComputeBlockKey, Block *> &graphComputeToBlockMap,
+  ArrayRef<ScheduledMaterializationRecord> materializedSchedules) {
+  pim::CappedDiagnosticReporter diagnostics;
+  size_t expectedClassCount = countPeftEquivalenceClasses(schedule);
+  if (expectedClassCount != materializedSchedules.size()) {
+    diagnostics.report(funcOp.getOperation(), [&](Operation *illegalOp) {
+      illegalOp->emitOpError() << "phase-check expected " << expectedClassCount
+                               << " PEFT equivalence classes but materialized " << materializedSchedules.size()
+                               << " scheduled computes";
+    });
+  }
+
+  llvm::SmallDenseSet<size_t, 16> seenClasses;
+  for (const ScheduledMaterializationRecord &record : materializedSchedules) {
+    if (!seenClasses.insert(record.canonicalPeftClassId).second) {
+      diagnostics.report(record.scheduledOp, [&](Operation *illegalOp) {
+        illegalOp->emitOpError("phase-check multiple scheduled ops own the same PEFT equivalence class");
+      });
+    }
+    if (!peftClassPlans.count(record.canonicalPeftClassId)) {
+      diagnostics.report(record.scheduledOp, [&](Operation *illegalOp) {
+        illegalOp->emitOpError("phase-check scheduled op refers to a missing PEFT materialization class");
+      });
+    }
+  }
+
+  for (const auto &entry : peftClassPlans) {
+    if (!seenClasses.count(entry.first)) {
+      diagnostics.report(funcOp.getOperation(), [&](Operation *illegalOp) {
+        illegalOp->emitOpError() << "phase-check PEFT equivalence class " << entry.first
+                                 << " was not materialized by any scheduled op";
+      });
+    }
+  }
+
+  for (GraphComputeBlockKey key : collectExpectedGraphComputeBlockKeys(funcOp)) {
+    if (graphComputeToBlockMap.count(key))
+      continue;
+    diagnostics.report(key.op, [&](Operation *illegalOp) {
+      illegalOp->emitOpError() << "phase-check graph compute is missing a scheduled MLIR block mapping for lanes ["
+                               << key.laneStart << ":" << (key.laneStart + key.laneCount) << "]";
+    });
+  }
+
+  for (const auto &entry : graphComputeToBlockMap) {
+    Block *block = entry.second;
+    if (!block || !isa<SpatScheduledCompute, SpatScheduledComputeBatch>(block->getParentOp())) {
+      diagnostics.report(entry.first.op, [&](Operation *illegalOp) {
+        illegalOp->emitOpError("phase-check graph compute block mapping does not target a scheduled compute block");
+      });
+    }
+  }
+
+  if (graphComputeToBlockMap.size() != collectExpectedGraphComputeBlockKeys(funcOp).size()) {
+    diagnostics.report(funcOp.getOperation(), [&](Operation *illegalOp) {
+      illegalOp->emitOpError() << "phase-check expected "
+                               << collectExpectedGraphComputeBlockKeys(funcOp).size()
+                               << " graph compute block mappings but saw " << graphComputeToBlockMap.size();
+    });
+  }
+
+  diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial materialization verification failed");
+  return success(!diagnostics.hasFailure());
+}
+
+LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp) {
+  pim::CappedDiagnosticReporter diagnostics;
+  funcOp.walk([&](SpatDeferredCommunicationOp transfer) {
+    for (Value source : transfer.getSources()) {
+      if (isa_and_nonnull<SpatScheduledCompute, SpatScheduledComputeBatch>(source.getDefiningOp())) {
+        diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) {
+          illegalOp->emitOpError("phase-check deferred communication source operand must remain an original graph SSA producer value");
+        });
+      }
+    }
+
+    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;
+          }
+        }
+      }
+    }
+  });
+
+  diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial deferred communication verification failed");
+  return success(!diagnostics.hasFailure());
+}
+
+LogicalResult verifyMultiCpuStepResultRouting(SpatScheduledComputeBatch scheduled,
+                                                     ArrayRef<ScheduledStepPlan> stepPlans) {
+  pim::CappedDiagnosticReporter diagnostics;
+  unsigned resultArgBase = getScheduledBatchResultArgBase(scheduled);
+  if (scheduled.getBody().getBlocks().size() != stepPlans.size()) {
+    diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) {
+      illegalOp->emitOpError() << "scheduled batch step routing expected " << stepPlans.size()
+                               << " blocks but saw " << scheduled.getBody().getBlocks().size();
+    });
+    diagnostics.emitSuppressedSummary(scheduled.getOperation(),
+                                      "scheduled batch step routing verification failed");
+    return failure();
+  }
+
+  SmallVector<unsigned> globalResultWrites(scheduled.getNumResults(), 0);
+  size_t stepIndex = 0;
+  for (Block &block : scheduled.getBody().getBlocks()) {
+    const ScheduledStepPlan &stepPlan = stepPlans[stepIndex++];
+    SmallVector<bool> localWrites(stepPlan.resultCount, false);
+    auto inParallel = dyn_cast<SpatInParallelOp>(block.getTerminator());
+    if (!inParallel) {
+      diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) {
+        illegalOp->emitOpError() << "scheduled batch step " << stepPlan.stepIndex
+                                 << " is missing spat.in_parallel";
+      });
+      continue;
+    }
+
+    for (Operation &op : inParallel.getRegion().front()) {
+      auto insert = dyn_cast<tensor::ParallelInsertSliceOp>(&op);
+      if (!insert)
+        continue;
+      auto dest = dyn_cast<BlockArgument>(insert.getDest());
+      if (!dest || dest.getOwner() != &block) {
+        diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) {
+          illegalOp->emitOpError() << "scheduled batch step " << stepPlan.stepIndex
+                                   << " writes to a non-block result destination";
+        });
+        continue;
+      }
+      unsigned resultIndex = dest.getArgNumber() - resultArgBase;
+      if (dest.getArgNumber() < resultArgBase || resultIndex >= scheduled.getNumResults()) {
+        diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) {
+          illegalOp->emitOpError() << "scheduled batch step " << stepPlan.stepIndex
+                                   << " writes to invalid result block argument " << dest.getArgNumber();
+        });
+        continue;
+      }
+      if (resultIndex < stepPlan.resultOffset
+          || resultIndex >= stepPlan.resultOffset + stepPlan.resultCount) {
+        diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) {
+          illegalOp->emitOpError() << "scheduled batch step " << stepPlan.stepIndex
+                                   << " expected result range [" << stepPlan.resultOffset << ":"
+                                   << (stepPlan.resultOffset + stepPlan.resultCount)
+                                   << ") but wrote result " << resultIndex;
+        });
+        continue;
+      }
+      localWrites[resultIndex - stepPlan.resultOffset] = true;
+      globalResultWrites[resultIndex]++;
+    }
+
+    for (size_t index = 0; index < localWrites.size(); ++index)
+      if (!localWrites[index])
+        diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) {
+          illegalOp->emitOpError() << "scheduled batch step " << stepPlan.stepIndex
+                                   << " did not write expected result " << (stepPlan.resultOffset + index);
+        });
+  }
+
+  for (size_t resultIndex = 0; resultIndex < globalResultWrites.size(); ++resultIndex)
+    if (globalResultWrites[resultIndex] != 1)
+      diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) {
+        illegalOp->emitOpError() << "scheduled batch result " << resultIndex << " expected one producing step but saw "
+                                 << globalResultWrites[resultIndex];
+      });
+
+  diagnostics.emitSuppressedSummary(scheduled.getOperation(), "scheduled batch step routing verification failed");
+  return success(!diagnostics.hasFailure());
+}
+
+LogicalResult verifyMultiCpuLocalFragmentOffsets(SpatScheduledComputeBatch scheduled) {
+  pim::CappedDiagnosticReporter diagnostics;
+  unsigned resultArgBase = getScheduledBatchResultArgBase(scheduled);
+  for (auto enumeratedBlock : llvm::enumerate(scheduled.getBody().getBlocks())) {
+    size_t stepIndex = enumeratedBlock.index();
+    Block &block = enumeratedBlock.value();
+    Value scheduledLane = block.getArgument(0);
+    auto inParallel = dyn_cast<SpatInParallelOp>(block.getTerminator());
+    if (!inParallel) {
+      diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) {
+        illegalOp->emitOpError() << "phase-check scheduled batch step " << stepIndex
+                                 << " is missing spat.in_parallel";
+      });
+      continue;
+    }
+    auto isFinalScheduledOutputInsert = [&](Operation *op) {
+      auto insert = dyn_cast<tensor::ParallelInsertSliceOp>(op);
+      if (!insert || op->getParentOp() != inParallel.getOperation())
+        return false;
+      auto dest = dyn_cast<BlockArgument>(insert.getDest());
+      return dest && dest.getOwner() == &block && dest.getArgNumber() >= resultArgBase;
+    };
+
+    block.walk([&](Operation *op) {
+      if (op == block.getTerminator())
+        return;
+      if (isFinalScheduledOutputInsert(op)) {
+        if (scheduled.getLaneCount() > 1) {
+          auto insert = cast<tensor::ParallelInsertSliceOp>(op);
+          bool dependsOnScheduledLane = false;
+          for (OpFoldResult offset : insert.getMixedOffsets()) {
+            if (auto value = dyn_cast<Value>(offset); value && valueTransitivelyDependsOn(value, scheduledLane)) {
+              dependsOnScheduledLane = true;
+              break;
+            }
+          }
+          if (!dependsOnScheduledLane)
+            diagnostics.report(insert.getOperation(), [&](Operation *illegalOp) {
+              illegalOp->emitOpError(
+                "phase-check scheduled batch final output insert must be indexed by scheduled lane");
+            });
+        }
+        return;
+      }
+
+      auto insertSlice = dyn_cast<tensor::InsertSliceOp>(op);
+      if (!insertSlice)
+        return;
+      auto dest = dyn_cast<BlockArgument>(insertSlice.getDest());
+      if (dest && dest.getOwner() == &block && dest.getArgNumber() >= resultArgBase)
+        return;
+
+      auto destType = dyn_cast<RankedTensorType>(insertSlice.getDestType());
+      if (!destType || !destType.hasStaticShape() || destType.getRank() == 0)
+        return;
+
+      for (OpFoldResult offset : insertSlice.getMixedOffsets()) {
+        auto value = dyn_cast<Value>(offset);
+        if (!value)
+          continue;
+        if (!valueTransitivelyDependsOn(value, scheduledLane))
+          continue;
+        diagnostics.report(insertSlice.getOperation(), [&](Operation *illegalOp) {
+          illegalOp->emitOpError()
+            << "phase-check scheduled batch local fragment insert offset must use the source-instance inner lane, not the scheduled lane"
+            << " step " << stepIndex;
+        });
+        break;
+      }
+    });
+  }
+
+  diagnostics.emitSuppressedSummary(scheduled.getOperation(),
+                                    "scheduled batch local fragment offset verification failed");
+  return success(!diagnostics.hasFailure());
+}
+
+
+LogicalResult verifyScheduledMaterializationRecords(ArrayRef<ScheduledMaterializationRecord> materializedSchedules) {
+  for (const ScheduledMaterializationRecord &record : materializedSchedules) {
+    auto scheduled = dyn_cast<SpatScheduledComputeBatch>(record.scheduledOp);
+    if (!scheduled)
+      continue;
+    if (failed(verifyMultiCpuStepResultRouting(scheduled, record.stepPlans)))
+      return failure();
+    if (failed(verifyMultiCpuLocalFragmentOffsets(scheduled)))
+      return failure();
+  }
+  return success();
+}
+
+} // namespace spatial
+} // namespace onnx_mlir

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

@@ -0,0 +1,22 @@
+#pragma once
+
+#include "ScheduledComputeMaterialization.hpp"
+
+namespace onnx_mlir {
+namespace spatial {
+
+LogicalResult verifyMaterializedScheduleMapping(
+  func::FuncOp funcOp,
+  const MergeScheduleResult &schedule,
+  const llvm::MapVector<size_t, PeftClassPlan> &peftClassPlans,
+  const DenseMap<GraphComputeBlockKey, Block *> &graphComputeToBlockMap,
+  ArrayRef<ScheduledMaterializationRecord> materializedSchedules);
+
+LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp);
+LogicalResult verifyMultiCpuStepResultRouting(SpatScheduledComputeBatch scheduled,
+                                              ArrayRef<ScheduledStepPlan> stepPlans);
+LogicalResult verifyMultiCpuLocalFragmentOffsets(SpatScheduledComputeBatch scheduled);
+LogicalResult verifyScheduledMaterializationRecords(ArrayRef<ScheduledMaterializationRecord> materializedSchedules);
+
+} // namespace spatial
+} // namespace onnx_mlir

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

@@ -1,832 +0,0 @@
-#include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/IR/AsmState.h"
-#include "mlir/IR/IRMapping.h"
-#include "mlir/Pass/Pass.h"
-
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/DenseSet.h"
-#include "llvm/ADT/MapVector.h"
-#include "llvm/ADT/SmallString.h"
-#include "llvm/ADT/TypeSwitch.h"
-#include "llvm/Support/FormatVariadic.h"
-#include "llvm/Support/raw_os_ostream.h"
-#include "llvm/Support/raw_ostream.h"
-
-#include <fstream>
-#include <optional>
-#include <string>
-#include <vector>
-
-#include "ComputeInstanceUtils.hpp"
-#include "MergeSchedulingAnalysis.hpp"
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
-#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
-
-using namespace mlir;
-
-namespace onnx_mlir {
-namespace spatial {
-namespace {
-
-struct ProducerValueKey {
-  ComputeInstance instance;
-  size_t resultIndex = 0;
-
-  bool operator==(const ProducerValueKey& other) const {
-    return instance == other.instance && resultIndex == other.resultIndex;
-  }
-};
-
-struct GraphComputeBlockKey {
-  Operation* op = nullptr;
-  uint32_t laneStart = 0;
-  uint32_t laneCount = 1;
-
-  bool operator==(const GraphComputeBlockKey& other) const {
-    return op == other.op && laneStart == other.laneStart && laneCount == other.laneCount;
-  }
-};
-
-struct ScheduledMaterializationRecord {
-  Operation* scheduledOp = nullptr;
-  SmallVector<size_t> canonicalCpuClasses;
-  SmallVector<size_t> cpus;
-  SmallVector<GraphComputeBlockKey> computeKeys;
-  SmallVector<Block*> blocks;
-};
-
-} // namespace
-} // namespace spatial
-} // namespace onnx_mlir
-
-namespace llvm {
-template <>
-struct DenseMapInfo<onnx_mlir::spatial::ProducerValueKey> {
-  static onnx_mlir::spatial::ProducerValueKey getEmptyKey() {
-    return {DenseMapInfo<onnx_mlir::spatial::ComputeInstance>::getEmptyKey(), std::numeric_limits<size_t>::max()};
-  }
-  static onnx_mlir::spatial::ProducerValueKey getTombstoneKey() {
-    return {DenseMapInfo<onnx_mlir::spatial::ComputeInstance>::getTombstoneKey(), std::numeric_limits<size_t>::max()};
-  }
-  static unsigned getHashValue(const onnx_mlir::spatial::ProducerValueKey& key) {
-    return hash_combine(DenseMapInfo<onnx_mlir::spatial::ComputeInstance>::getHashValue(key.instance), key.resultIndex);
-  }
-  static bool isEqual(const onnx_mlir::spatial::ProducerValueKey& lhs,
-                      const onnx_mlir::spatial::ProducerValueKey& rhs) {
-    return lhs == rhs;
-  }
-};
-
-template <>
-struct DenseMapInfo<onnx_mlir::spatial::GraphComputeBlockKey> {
-  static onnx_mlir::spatial::GraphComputeBlockKey getEmptyKey() {
-    return {DenseMapInfo<mlir::Operation*>::getEmptyKey(), UINT32_MAX, UINT32_MAX};
-  }
-  static onnx_mlir::spatial::GraphComputeBlockKey getTombstoneKey() {
-    return {DenseMapInfo<mlir::Operation*>::getTombstoneKey(), UINT32_MAX, UINT32_MAX - 1};
-  }
-  static unsigned getHashValue(const onnx_mlir::spatial::GraphComputeBlockKey& key) {
-    return hash_combine(key.op, key.laneStart, key.laneCount);
-  }
-  static bool isEqual(const onnx_mlir::spatial::GraphComputeBlockKey& lhs,
-                      const onnx_mlir::spatial::GraphComputeBlockKey& rhs) {
-    return lhs == rhs;
-  }
-};
-} // namespace llvm
-
-namespace onnx_mlir {
-namespace spatial {
-namespace {
-
-struct ScalarClass {
-  size_t cpu = 0;
-  SmallVector<size_t> cpus;
-  SmallVector<ComputeInstance> instances;
-  SmallVector<Value> weights;
-  SmallVector<Value> inputs;
-  SmallVector<Type> resultTypes;
-};
-
-struct CompactCpuIndexMap {
-  DenseMap<size_t, size_t> rawToCompact;
-
-  size_t lookup(size_t rawCpu) const {
-    auto it = rawToCompact.find(rawCpu);
-    assert(it != rawToCompact.end() && "missing compact cpu index");
-    return it->second;
-  }
-};
-
-static ComputeInstance getWholeBatchInstance(SpatComputeBatch batch) {
-  return {batch.getOperation(), 0, static_cast<uint32_t>(batch.getLaneCount())};
-}
-
-static GraphComputeBlockKey getGraphComputeBlockKey(const ComputeInstance& instance) {
-  return {instance.op, instance.laneStart, instance.laneCount};
-}
-
-static CompactCpuIndexMap buildCompactCpuIndexMap(const MergeScheduleResult& schedule) {
-  SmallVector<size_t> rawCpus;
-  auto append = [&](size_t rawCpu) {
-    if (!llvm::is_contained(rawCpus, rawCpu))
-      rawCpus.push_back(rawCpu);
-  };
-  for (const auto& entry : schedule.computeToCpuMap)
-    append(entry.second);
-  for (const auto& entry : schedule.equivalentClass) {
-    append(entry.first);
-    for (size_t rawCpu : entry.second)
-      append(rawCpu);
-  }
-  for (const auto& entry : schedule.cpuToLastComputeMap)
-    append(entry.first);
-  llvm::sort(rawCpus);
-
-  CompactCpuIndexMap compactMap;
-  for (auto [index, rawCpu] : llvm::enumerate(rawCpus))
-    compactMap.rawToCompact[rawCpu] = index;
-  return compactMap;
-}
-
-static SmallVector<size_t> getEquivalentCpus(size_t cpu, const MergeScheduleResult& schedule) {
-  llvm::SmallDenseSet<size_t, 8> seen;
-  SmallVector<size_t> cpus;
-  auto append = [&](size_t value) {
-    if (seen.insert(value).second)
-      cpus.push_back(value);
-  };
-  append(cpu);
-  if (auto it = schedule.equivalentClass.find(cpu); it != schedule.equivalentClass.end())
-    for (size_t equivalentCpu : it->second)
-      append(equivalentCpu);
-  llvm::sort(cpus);
-  return cpus;
-}
-
-static size_t getCanonicalCpuClass(size_t cpu, const MergeScheduleResult& schedule) {
-  SmallVector<size_t> cpus = getEquivalentCpus(cpu, schedule);
-  return cpus.empty() ? cpu : cpus.front();
-}
-
-static Value getProducerOutput(const ProducerValueRef& ref) {
-  auto outputs = getComputeInstanceOutputValues(ref.instance);
-  assert(ref.resultIndex < outputs.size() && "producer result index out of range");
-  return outputs[ref.resultIndex];
-}
-
-static SmallVector<size_t> getCanonicalCpuClassesForBatch(SpatComputeBatch batch, const MergeScheduleResult& schedule) {
-  llvm::SmallDenseSet<size_t, 4> seen;
-  SmallVector<size_t> classes;
-  size_t chunkCount = getBatchChunkTargetCount(batch.getLaneCount());
-  for (size_t chunkIndex = 0; chunkIndex < chunkCount; ++chunkIndex) {
-    ComputeInstance chunk = getBatchChunkForIndex(batch, chunkIndex);
-    auto cpuIt = schedule.computeToCpuMap.find(chunk);
-    if (cpuIt == schedule.computeToCpuMap.end())
-      continue;
-    size_t cpuClass = getCanonicalCpuClass(cpuIt->second, schedule);
-    if (seen.insert(cpuClass).second)
-      classes.push_back(cpuClass);
-  }
-  llvm::sort(classes);
-  return classes;
-}
-
-static SmallVector<size_t> remapCpuList(ArrayRef<size_t> rawCpus, const CompactCpuIndexMap& compactCpuMap) {
-  SmallVector<size_t> compactCpus;
-  compactCpus.reserve(rawCpus.size());
-  for (size_t rawCpu : rawCpus)
-    compactCpus.push_back(compactCpuMap.lookup(rawCpu));
-  llvm::sort(compactCpus);
-  return compactCpus;
-}
-
-static std::string getInstanceName(const ComputeInstance& instance) {
-  return llvm::formatv("{0}[lanes={1}:{2}]",
-                       instance.op->getName().getStringRef(),
-                       instance.laneStart,
-                       instance.laneStart + instance.laneCount)
-    .str();
-}
-
-static void appendUnique(SmallVectorImpl<Value>& values, Value value) {
-  if (!llvm::is_contained(values, value))
-    values.push_back(value);
-}
-
-static Value getBlockOperand(Block& block, ValueRange operands, Value value, unsigned firstArgument = 0) {
-  auto it = llvm::find(operands, value);
-  assert(it != operands.end() && "missing scheduled operand");
-  return block.getArgument(firstArgument + std::distance(operands.begin(), it));
-}
-
-static void appendComputeBlockArguments(SmallVectorImpl<Type>& argTypes,
-                                        SmallVectorImpl<Location>& argLocs,
-                                        ValueRange weights,
-                                        ValueRange inputs,
-                                        ArrayRef<ProducerValueKey> carriedKeys,
-                                        Location loc) {
-  for (Value weight : weights)
-    argTypes.push_back(weight.getType());
-  for (Value input : inputs)
-    argTypes.push_back(input.getType());
-  for (ProducerValueKey key : carriedKeys) {
-    auto outputs = getComputeInstanceOutputValues(key.instance);
-    assert(key.resultIndex < outputs.size() && "missing carried result");
-    argTypes.push_back(outputs[key.resultIndex].getType());
-  }
-  argLocs.append(argTypes.size(), loc);
-}
-
-static Block* createScheduledComputeBlock(PatternRewriter& rewriter,
-                                          SpatScheduledCompute scheduled,
-                                          ArrayRef<ProducerValueKey> carriedKeys,
-                                          Location loc) {
-  SmallVector<Type> argTypes;
-  SmallVector<Location> argLocs;
-  appendComputeBlockArguments(argTypes, argLocs, scheduled.getWeights(), scheduled.getInputs(), carriedKeys, loc);
-  return rewriter.createBlock(&scheduled.getBody(), scheduled.getBody().end(), TypeRange(argTypes), argLocs);
-}
-
-static void appendBlockYieldBaseAndCarriedOperands(Block& block,
-                                                   unsigned baseArgCount,
-                                                   size_t carriedCount,
-                                                   SmallVectorImpl<Value>& operands) {
-  for (unsigned index = 0; index < baseArgCount; ++index)
-    operands.push_back(block.getArgument(index));
-  for (size_t index = 0; index < carriedCount; ++index)
-    operands.push_back(block.getArgument(baseArgCount + index));
-}
-
-static void createBlockYield(PatternRewriter& rewriter, Location loc, ValueRange outputs, Block* next = nullptr) {
-  OperationState state(loc, SpatBlockYieldOp::getOperationName());
-  state.addOperands(outputs);
-  if (next)
-    state.addSuccessors(next);
-  rewriter.create(state);
-}
-
-static std::string formatValueLabel(Value value, AsmState& asmState) {
-  std::string storage;
-  llvm::raw_string_ostream os(storage);
-  value.printAsOperand(os, asmState);
-  return storage;
-}
-
-static std::string formatBlockLabel(Block* block, AsmState& asmState) {
-  std::string storage;
-  llvm::raw_string_ostream os(storage);
-  block->printAsOperand(os, asmState);
-  return storage;
-}
-
-static std::string formatOperationLabel(Operation* op, AsmState& asmState) {
-  if (op->getNumResults() == 0)
-    return op->getName().getStringRef().str();
-  std::string storage;
-  llvm::raw_string_ostream os(storage);
-  llvm::interleaveComma(op->getResults(), os, [&](Value result) { os << formatValueLabel(result, asmState); });
-  return os.str();
-}
-
-static std::string formatGraphComputeBlockKey(const GraphComputeBlockKey& key, AsmState& asmState) {
-  return llvm::formatv("{0} {1} lanes [{2}:{3}]",
-                       formatOperationLabel(key.op, asmState),
-                       key.op->getName().getStringRef(),
-                       key.laneStart,
-                       key.laneStart + key.laneCount)
-    .str();
-}
-
-static size_t countPeftEquivalenceClasses(const MergeScheduleResult& schedule) {
-  llvm::SmallDenseSet<size_t, 16> classes;
-  for (const ComputeInstance& instance : schedule.dominanceOrderCompute) {
-    auto cpuIt = schedule.computeToCpuMap.find(instance);
-    if (cpuIt == schedule.computeToCpuMap.end())
-      continue;
-    classes.insert(getCanonicalCpuClass(cpuIt->second, schedule));
-  }
-  return classes.size();
-}
-
-static SmallVector<GraphComputeBlockKey> collectExpectedGraphComputeBlockKeys(func::FuncOp funcOp) {
-  SmallVector<GraphComputeBlockKey> keys;
-  for (Operation& op : funcOp.getOps()) {
-    if (auto compute = dyn_cast<SpatGraphCompute>(&op))
-      keys.push_back(getGraphComputeBlockKey({compute.getOperation(), 0, 1}));
-    else if (auto batch = dyn_cast<SpatGraphComputeBatch>(&op))
-      keys.push_back(getGraphComputeBlockKey(getWholeBatchInstance(batch)));
-  }
-  return keys;
-}
-
-static LogicalResult verifyMaterializedScheduleMapping(
-  func::FuncOp funcOp,
-  const MergeScheduleResult& schedule,
-  const DenseMap<GraphComputeBlockKey, Block*>& graphComputeToBlockMap,
-  ArrayRef<ScheduledMaterializationRecord> materializedSchedules) {
-  pim::CappedDiagnosticReporter diagnostics;
-  size_t expectedClassCount = countPeftEquivalenceClasses(schedule);
-  if (expectedClassCount != materializedSchedules.size()) {
-    diagnostics.report(funcOp.getOperation(), [&](Operation* illegalOp) {
-      illegalOp->emitOpError() << "phase-check expected " << expectedClassCount
-                               << " PEFT equivalence classes but materialized " << materializedSchedules.size()
-                               << " scheduled computes";
-    });
-  }
-
-  for (const ScheduledMaterializationRecord& record : materializedSchedules) {
-    if (record.canonicalCpuClasses.empty()) {
-      diagnostics.report(record.scheduledOp, [&](Operation* illegalOp) {
-        illegalOp->emitOpError("phase-check scheduled compute is missing an owning PEFT equivalence class");
-      });
-      continue;
-    }
-    if (record.canonicalCpuClasses.size() > 1) {
-      diagnostics.report(record.scheduledOp, [&](Operation* illegalOp) {
-        std::string classes;
-        llvm::raw_string_ostream os(classes);
-        llvm::interleaveComma(record.canonicalCpuClasses, os, [&](size_t cpuClass) { os << cpuClass; });
-        illegalOp->emitOpError("phase-check scheduled compute spans multiple PEFT equivalence classes")
-          << " (classes " << os.str() << ")";
-      });
-    }
-  }
-
-  for (GraphComputeBlockKey key : collectExpectedGraphComputeBlockKeys(funcOp)) {
-    if (graphComputeToBlockMap.count(key))
-      continue;
-    diagnostics.report(key.op, [&](Operation* illegalOp) {
-      illegalOp->emitOpError() << "phase-check graph compute is missing a scheduled MLIR block mapping for lanes ["
-                               << key.laneStart << ":" << (key.laneStart + key.laneCount) << "]";
-    });
-  }
-
-  for (const auto& entry : graphComputeToBlockMap) {
-    Block* block = entry.second;
-    if (!block || !isa<SpatScheduledCompute, SpatScheduledComputeBatch>(block->getParentOp())) {
-      diagnostics.report(entry.first.op, [&](Operation* illegalOp) {
-        illegalOp->emitOpError("phase-check graph compute block mapping does not target a scheduled compute block");
-      });
-    }
-  }
-
-  diagnostics.emitSuppressedSummary(funcOp, "scheduled Spatial materialization verification failed");
-  return success(!diagnostics.hasFailure());
-}
-
-static void dumpPeftMaterializationReport(func::FuncOp funcOp,
-                                          const MergeScheduleResult& schedule,
-                                          const CompactCpuIndexMap& compactCpuMap,
-                                          ArrayRef<ScheduledMaterializationRecord> materializedSchedules) {
-  std::fstream file = openDialectDumpFileWithExtension("spatial1_schedule_materialization", "/reports", "txt");
-  if (!file.is_open())
-    return;
-
-  llvm::raw_os_ostream os(file);
-  AsmState asmState(funcOp);
-
-  llvm::MapVector<size_t, SmallVector<size_t>> classToCpus;
-  DenseMap<size_t, SmallVector<ComputeInstance>> cpuToInstances;
-  for (const ComputeInstance& instance : schedule.dominanceOrderCompute) {
-    auto cpuIt = schedule.computeToCpuMap.find(instance);
-    if (cpuIt == schedule.computeToCpuMap.end())
-      continue;
-    size_t cpu = compactCpuMap.lookup(cpuIt->second);
-    size_t cpuClass = compactCpuMap.lookup(getCanonicalCpuClass(cpuIt->second, schedule));
-    if (!llvm::is_contained(classToCpus[cpuClass], cpu))
-      classToCpus[cpuClass].push_back(cpu);
-    cpuToInstances[cpu].push_back(instance);
-  }
-
-  os << "PEFT Schedule\n";
-  os << "=============\n";
-  os << "equivalence classes: " << classToCpus.size() << "\n\n";
-  for (auto& entry : classToCpus) {
-    size_t cpuClass = entry.first;
-    llvm::sort(entry.second);
-    os << "class " << cpuClass << " cpus [";
-    llvm::interleaveComma(entry.second, os, [&](size_t cpu) { os << cpu; });
-    os << "]\n";
-    for (size_t cpu : entry.second) {
-      os << "  cpu " << cpu << ":\n";
-      auto& instances = cpuToInstances[cpu];
-      llvm::sort(instances, [&](const ComputeInstance& lhs, const ComputeInstance& rhs) {
-        return schedule.computeToCpuSlotMap.lookup(lhs) < schedule.computeToCpuSlotMap.lookup(rhs);
-      });
-      for (const ComputeInstance& instance : instances) {
-        os << "    slot " << schedule.computeToCpuSlotMap.lookup(instance) << ": "
-           << formatGraphComputeBlockKey(getGraphComputeBlockKey(instance), asmState) << "\n";
-      }
-    }
-    os << "\n";
-  }
-
-  os << "Materialized Scheduled Computes\n";
-  os << "===============================\n";
-  os << "scheduled ops: " << materializedSchedules.size() << "\n\n";
-  for (const ScheduledMaterializationRecord& record : materializedSchedules) {
-    os << formatOperationLabel(record.scheduledOp, asmState) << " " << record.scheduledOp->getName().getStringRef();
-    os << " classes [";
-    llvm::interleaveComma(record.canonicalCpuClasses, os, [&](size_t cpuClass) { os << cpuClass; });
-    os << "] cpus [";
-    llvm::interleaveComma(record.cpus, os, [&](size_t cpu) { os << cpu; });
-    os << "]\n";
-    for (auto [key, block] : llvm::zip(record.computeKeys, record.blocks))
-      os << "  " << formatGraphComputeBlockKey(key, asmState) << " -> " << formatBlockLabel(block, asmState) << "\n";
-    os << "\n";
-  }
-}
-
-static LogicalResult createDeferredCommunication(PatternRewriter& rewriter,
-                                                 Location loc,
-                                                 Value consumerInput,
-                                                 Value source,
-                                                 const ProducerValueRef& producer,
-                                                 const ComputeInstance& consumer,
-                                                 const MergeScheduleResult& schedule,
-                                                 const CompactCpuIndexMap& compactCpuMap,
-                                                 uint64_t exchangeId,
-                                                 Value& result) {
-  auto sourceCpu = compactCpuMap.lookup(schedule.computeToCpuMap.lookup(producer.instance));
-  auto targetCpu = compactCpuMap.lookup(schedule.computeToCpuMap.lookup(consumer));
-  auto transfer = SpatDeferredCommunicationOp::create(
-    rewriter,
-    loc,
-    consumerInput.getType(),
-    ValueRange {source},
-    rewriter.getStringAttr(("x" + std::to_string(exchangeId)).c_str()),
-    rewriter.getStringAttr(getInstanceName(producer.instance)),
-    rewriter.getStringAttr(getInstanceName(consumer)),
-    rewriter.getI64IntegerAttr(static_cast<int64_t>(sourceCpu)),
-    rewriter.getI64IntegerAttr(static_cast<int64_t>(targetCpu)),
-    rewriter.getI64IntegerAttr(static_cast<int64_t>(sourceCpu)),
-    rewriter.getI64IntegerAttr(static_cast<int64_t>(targetCpu)),
-    rewriter.getI64IntegerAttr(producer.instance.laneStart),
-    rewriter.getI64IntegerAttr(consumer.laneStart),
-    rewriter.getStringAttr(consumerInput.getDefiningOp<tensor::ExtractSliceOp>() ? "projected" : "ordinary"),
-    rewriter.getI64IntegerAttr(static_cast<int64_t>(producer.resultIndex)),
-    rewriter.getDenseI64ArrayAttr({}),
-    rewriter.getI64IntegerAttr(-1));
-
-  Block* block = rewriter.createBlock(&transfer.getBody(), transfer.getBody().end(), source.getType(), loc);
-  rewriter.setInsertionPointToStart(block);
-  Value yielded = block->getArgument(0);
-  if (yielded.getType() != consumerInput.getType()) {
-    auto slice = consumerInput.getDefiningOp<tensor::ExtractSliceOp>();
-    if (!slice || slice.getSource().getType() != source.getType())
-      return transfer.emitOpError("cannot derive deferred communication payload type from recorded source");
-    IRMapping mapper;
-    mapper.map(slice.getSource(), yielded);
-    auto cloned = cast<tensor::ExtractSliceOp>(rewriter.clone(*slice, mapper));
-    yielded = cloned.getResult();
-  }
-  SpatYieldOp::create(rewriter, loc, yielded);
-  result = transfer.getOutput();
-  rewriter.setInsertionPointAfter(transfer);
-  return success();
-}
-
-static LogicalResult mapComputeInputs(PatternRewriter& rewriter,
-                                      Block& block,
-                                      SpatCompute compute,
-                                      const ComputeInstance& instance,
-                                      const ScalarClass& klass,
-                                      const MergeScheduleResult& schedule,
-                                      const CompactCpuIndexMap& compactCpuMap,
-                                      ValueRange scheduledWeights,
-                                      ValueRange scheduledInputs,
-                                      IRMapping& mapper,
-                                      uint64_t& exchangeId) {
-  for (auto [index, weight] : llvm::enumerate(compute.getWeights()))
-    mapper.map(*compute.getWeightArgument(index), getBlockOperand(block, scheduledWeights, weight));
-
-  for (auto [index, input] : llvm::enumerate(compute.getInputs())) {
-    Value mapped;
-    std::optional<ProducerValueRef> producer = getProducerValueRef(input, &instance);
-    if (producer) {
-      Value source = getProducerOutput(*producer);
-      if (failed(createDeferredCommunication(
-            rewriter, input.getLoc(), input, source, *producer, instance, schedule, compactCpuMap, exchangeId++, mapped)))
-        return failure();
-    }
-    else {
-      mapped = getBlockOperand(block, scheduledInputs, input, scheduledWeights.size());
-    }
-    mapper.map(*compute.getInputArgument(index), mapped);
-  }
-  return success();
-}
-
-static LogicalResult mapBatchInputs(PatternRewriter& rewriter,
-                                    Block& block,
-                                    SpatComputeBatch batch,
-                                    const MergeScheduleResult& schedule,
-                                    const CompactCpuIndexMap& compactCpuMap,
-                                    ValueRange scheduledWeights,
-                                    ValueRange scheduledInputs,
-                                    IRMapping& mapper,
-                                    uint64_t& exchangeId) {
-  mapper.map(*batch.getLaneArgument(), block.getArgument(0));
-  unsigned firstWeightArg = 1;
-  unsigned firstInputArg = firstWeightArg + scheduledWeights.size();
-  for (auto [index, weight] : llvm::enumerate(batch.getWeights()))
-    mapper.map(*batch.getWeightArgument(index), block.getArgument(firstWeightArg + index));
-
-  ComputeInstance consumer {batch.getOperation(), 0, static_cast<uint32_t>(batch.getLaneCount())};
-  for (auto [index, input] : llvm::enumerate(batch.getInputs())) {
-    Value mapped;
-    std::optional<ProducerValueRef> producer = getProducerValueRef(input, &consumer);
-    if (producer) {
-      Value source = getProducerOutput(*producer);
-      if (failed(createDeferredCommunication(
-            rewriter, input.getLoc(), input, source, *producer, consumer, schedule, compactCpuMap, exchangeId++, mapped)))
-        return failure();
-    }
-    else {
-      auto it = llvm::find(scheduledInputs, input);
-      if (it == scheduledInputs.end())
-        return batch.emitOpError("host batch input is missing from scheduled batch operands");
-      mapped = block.getArgument(firstInputArg + std::distance(scheduledInputs.begin(), it));
-    }
-    mapper.map(*batch.getInputArgument(index), mapped);
-  }
-
-  unsigned firstOutputArg = firstInputArg + scheduledInputs.size();
-  for (unsigned index = 0; index < batch.getNumResults(); ++index)
-    mapper.map(*batch.getOutputArgument(index), block.getArgument(firstOutputArg + index));
-  return success();
-}
-
-static void cloneComputeBody(PatternRewriter& rewriter,
-                             Block& source,
-                             IRMapping& mapper,
-                             SmallVectorImpl<Value>& yieldedValues) {
-  for (Operation& op : source.without_terminator())
-    rewriter.clone(op, mapper);
-  auto yield = cast<SpatYieldOp>(source.getTerminator());
-  for (Value output : yield.getOutputs())
-    yieldedValues.push_back(mapper.lookup(output));
-}
-
-struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, OperationPass<ModuleOp>> {
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MergeComputeNodesPass)
-
-  StringRef getArgument() const override { return "pim-merge-compute-nodes"; }
-  StringRef getDescription() const override {
-    return "Materialize scheduled Spatial compute with deferred communication placeholders.";
-  }
-
-  void runOnOperation() override {
-    ModuleOp moduleOp = getOperation();
-    auto entryFunc = getPimEntryFunc(moduleOp);
-    if (failed(entryFunc)) {
-      moduleOp.emitError("failed to locate the PIM entry function during MergeComputeNodes");
-      signalPassFailure();
-      return;
-    }
-
-    func::FuncOp funcOp = *entryFunc;
-    MergeScheduleResult schedule = MergeSchedulingAnalysis(funcOp).getResult();
-    CompactCpuIndexMap compactCpuMap = buildCompactCpuIndexMap(schedule);
-    llvm::MapVector<size_t, ScalarClass> scalarClasses;
-    SmallVector<SpatComputeBatch> batchOps;
-    for (const ComputeInstance& instance : schedule.dominanceOrderCompute) {
-      if (isa<SpatCompute>(instance.op)) {
-        size_t cpu = schedule.computeToCpuMap.lookup(instance);
-        size_t cpuClass = getCanonicalCpuClass(cpu, schedule);
-        auto& klass = scalarClasses[cpuClass];
-        klass.cpu = cpuClass;
-        if (klass.cpus.empty())
-          klass.cpus = getEquivalentCpus(cpu, schedule);
-        klass.instances.push_back(instance);
-      }
-      else {
-        auto batch = cast<SpatComputeBatch>(instance.op);
-        if (!llvm::is_contained(batchOps, batch))
-          batchOps.push_back(batch);
-      }
-    }
-    for (auto& entry : scalarClasses) {
-      llvm::sort(entry.second.instances, [&](const ComputeInstance& lhs, const ComputeInstance& rhs) {
-        return schedule.computeToCpuSlotMap.lookup(lhs) < schedule.computeToCpuSlotMap.lookup(rhs);
-      });
-    }
-
-    PatternRewriter rewriter(moduleOp.getContext());
-    Operation* insertionPoint = funcOp.getBody().front().getTerminator();
-    DenseMap<ScalarClass*, SpatScheduledCompute> scalarOps;
-    DenseMap<Operation*, SpatScheduledComputeBatch> batchReplacements;
-    DenseMap<Operation*, size_t> scheduledOpToRecordIndex;
-    DenseMap<GraphComputeBlockKey, Block*> graphComputeToBlockMap;
-    std::vector<ScheduledMaterializationRecord> materializedSchedules;
-
-    for (auto& entry : scalarClasses) {
-      ScalarClass& klass = entry.second;
-      for (const ComputeInstance& instance : klass.instances) {
-        auto compute = cast<SpatCompute>(instance.op);
-        llvm::append_range(klass.resultTypes, compute.getResultTypes());
-        for (Value weight : compute.getWeights())
-          appendUnique(klass.weights, weight);
-        for (Value input : compute.getInputs()) {
-          std::optional<ProducerValueRef> producer = getProducerValueRef(input, &instance);
-          if (!producer)
-            appendUnique(klass.inputs, input);
-        }
-      }
-      rewriter.setInsertionPoint(insertionPoint);
-      auto scheduled = SpatScheduledCompute::create(
-        rewriter, klass.instances.front().op->getLoc(), TypeRange(klass.resultTypes), klass.weights, klass.inputs);
-      scheduled->setAttr(
-        kCoreIdAttrName, rewriter.getI32IntegerAttr(static_cast<int32_t>(compactCpuMap.lookup(klass.cpu))));
-      SmallVector<int64_t> equivalentCpus;
-      for (size_t cpu : klass.cpus)
-        equivalentCpus.push_back(static_cast<int64_t>(compactCpuMap.lookup(cpu)));
-      scheduled->setAttr("scheduled.equivalent_cpus", rewriter.getDenseI64ArrayAttr(equivalentCpus));
-      SmallVector<Attribute> sources;
-      SmallVector<int64_t> lanes;
-      for (const ComputeInstance& instance : klass.instances) {
-        sources.push_back(rewriter.getStringAttr(getInstanceName(instance)));
-        lanes.push_back(instance.laneStart);
-        lanes.push_back(instance.laneCount);
-      }
-      scheduled->setAttr("scheduled.sources", rewriter.getArrayAttr(sources));
-      scheduled->setAttr("scheduled.lane_ranges", rewriter.getDenseI64ArrayAttr(lanes));
-      scalarOps[&klass] = scheduled;
-      ScheduledMaterializationRecord record;
-      record.scheduledOp = scheduled.getOperation();
-      record.canonicalCpuClasses.push_back(compactCpuMap.lookup(klass.cpu));
-      record.cpus = remapCpuList(klass.cpus, compactCpuMap);
-      scheduledOpToRecordIndex[scheduled.getOperation()] = materializedSchedules.size();
-      materializedSchedules.push_back(std::move(record));
-    }
-
-    for (SpatComputeBatch batch : batchOps) {
-      SmallVector<Value> inputs;
-      ComputeInstance wholeBatch = getWholeBatchInstance(batch);
-      for (Value input : batch.getInputs())
-        if (!getProducerValueRef(input, &wholeBatch))
-          appendUnique(inputs, input);
-
-      rewriter.setInsertionPoint(insertionPoint);
-      auto scheduled = SpatScheduledComputeBatch::create(rewriter,
-                                                         batch.getLoc(),
-                                                         batch.getResultTypes(),
-                                                         rewriter.getI32IntegerAttr(batch.getLaneCount()),
-                                                         batch.getWeights(),
-                                                         inputs);
-      SmallVector<int32_t> coreIds;
-      uint32_t laneCount = batch.getLaneCount();
-      coreIds.reserve(laneCount);
-      for (uint32_t lane = 0; lane < laneCount; ++lane) {
-        ComputeInstance chunk = getBatchChunkForLane(batch, lane);
-        coreIds.push_back(static_cast<int32_t>(compactCpuMap.lookup(schedule.computeToCpuMap.lookup(chunk))));
-      }
-      scheduled->setAttr(kCoreIdsAttrName, rewriter.getDenseI32ArrayAttr(coreIds));
-      scheduled->setAttr("scheduled.sources", rewriter.getArrayAttr({rewriter.getStringAttr(getInstanceName(wholeBatch))}));
-      batchReplacements[batch.getOperation()] = scheduled;
-      ScheduledMaterializationRecord record;
-      record.scheduledOp = scheduled.getOperation();
-      record.canonicalCpuClasses = remapCpuList(getCanonicalCpuClassesForBatch(batch, schedule), compactCpuMap);
-      for (int32_t coreId : coreIds)
-        if (!llvm::is_contained(record.cpus, static_cast<size_t>(coreId)))
-          record.cpus.push_back(static_cast<size_t>(coreId));
-      scheduledOpToRecordIndex[scheduled.getOperation()] = materializedSchedules.size();
-      materializedSchedules.push_back(std::move(record));
-    }
-
-    uint64_t exchangeId = 0;
-    for (auto& entry : scalarClasses) {
-      ScalarClass& klass = entry.second;
-      SpatScheduledCompute scheduled = scalarOps.lookup(&klass);
-      SmallVector<ProducerValueKey> carriedKeys;
-      Block* block = nullptr;
-      for (auto [ordinal, instance] : llvm::enumerate(klass.instances)) {
-        auto compute = cast<SpatCompute>(instance.op);
-        if (!block)
-          block = createScheduledComputeBlock(rewriter, scheduled, carriedKeys, compute.getLoc());
-        GraphComputeBlockKey key = getGraphComputeBlockKey(instance);
-        graphComputeToBlockMap[key] = block;
-        auto& record = materializedSchedules[scheduledOpToRecordIndex.lookup(scheduled.getOperation())];
-        record.computeKeys.push_back(key);
-        record.blocks.push_back(block);
-        rewriter.setInsertionPointToStart(block);
-        IRMapping mapper;
-        if (failed(mapComputeInputs(rewriter,
-                                    *block,
-                                    compute,
-                                    instance,
-                                    klass,
-                                    schedule,
-                                    compactCpuMap,
-                                    scheduled.getWeights(),
-                                    scheduled.getInputs(),
-                                    mapper,
-                                    exchangeId))) {
-          signalPassFailure();
-          return;
-        }
-        SmallVector<Value> yieldedValues;
-        cloneComputeBody(rewriter, compute.getBody().front(), mapper, yieldedValues);
-        SmallVector<ProducerValueKey> currentKeys;
-        for (auto [index, value] : llvm::enumerate(yieldedValues)) {
-          ProducerValueKey key {instance, index};
-          currentKeys.push_back(key);
-        }
-
-        unsigned baseArgCount = scheduled.getWeights().size() + scheduled.getInputs().size();
-        SmallVector<Value> blockYieldOperands;
-        bool hasNextBlock = ordinal + 1 < klass.instances.size();
-        if (hasNextBlock) {
-          SmallVector<ProducerValueKey> nextCarriedKeys(carriedKeys);
-          llvm::append_range(nextCarriedKeys, currentKeys);
-          Block* nextBlock = createScheduledComputeBlock(rewriter, scheduled, nextCarriedKeys, compute.getLoc());
-          appendBlockYieldBaseAndCarriedOperands(*block, baseArgCount, carriedKeys.size(), blockYieldOperands);
-          llvm::append_range(blockYieldOperands, yieldedValues);
-          rewriter.setInsertionPointToEnd(block);
-          createBlockYield(rewriter, compute.getLoc(), blockYieldOperands, nextBlock);
-          carriedKeys = std::move(nextCarriedKeys);
-          block = nextBlock;
-        }
-        else {
-          for (size_t index = 0; index < carriedKeys.size(); ++index)
-            blockYieldOperands.push_back(block->getArgument(baseArgCount + index));
-          llvm::append_range(blockYieldOperands, yieldedValues);
-          rewriter.setInsertionPointToEnd(block);
-          createBlockYield(rewriter, compute.getLoc(), blockYieldOperands);
-        }
-      }
-    }
-
-    for (SpatComputeBatch batch : batchOps) {
-      SpatScheduledComputeBatch scheduled = batchReplacements.lookup(batch.getOperation());
-      SmallVector<Type> blockArgTypes {rewriter.getIndexType()};
-      SmallVector<Location> blockArgLocs {batch.getLoc()};
-      for (Value weight : scheduled.getWeights()) {
-        blockArgTypes.push_back(weight.getType());
-        blockArgLocs.push_back(weight.getLoc());
-      }
-      for (Value input : scheduled.getInputs()) {
-        blockArgTypes.push_back(input.getType());
-        blockArgLocs.push_back(input.getLoc());
-      }
-      for (Type resultType : scheduled.getResultTypes()) {
-        blockArgTypes.push_back(resultType);
-        blockArgLocs.push_back(batch.getLoc());
-      }
-      Block* block = rewriter.createBlock(
-        &scheduled.getBody(), scheduled.getBody().end(), TypeRange(blockArgTypes), blockArgLocs);
-      GraphComputeBlockKey key = getGraphComputeBlockKey(getWholeBatchInstance(batch));
-      graphComputeToBlockMap[key] = block;
-      auto& record = materializedSchedules[scheduledOpToRecordIndex.lookup(scheduled.getOperation())];
-      record.computeKeys.push_back(key);
-      record.blocks.push_back(block);
-      rewriter.setInsertionPointToStart(block);
-      IRMapping mapper;
-      if (failed(mapBatchInputs(rewriter,
-                                *block,
-                                batch,
-                                schedule,
-                                compactCpuMap,
-                                scheduled.getWeights(),
-                                scheduled.getInputs(),
-                                mapper,
-                                exchangeId))) {
-        signalPassFailure();
-        return;
-      }
-      for (Operation& op : batch.getBody().front())
-        rewriter.clone(op, mapper);
-    }
-
-    dumpPeftMaterializationReport(funcOp, schedule, compactCpuMap, materializedSchedules);
-    if (failed(verifyMaterializedScheduleMapping(funcOp, schedule, graphComputeToBlockMap, materializedSchedules))) {
-      moduleOp.emitError("scheduled Spatial materialization mapping verification failed");
-      signalPassFailure();
-      return;
-    }
-    if (failed(verifyScheduledSpatialInvariants(funcOp))) {
-      moduleOp.emitError("scheduled Spatial phase 1 verification failed");
-      signalPassFailure();
-      return;
-    }
-    dumpModule(moduleOp, "spatial1_scheduled_no_comm", /*assumeVerified=*/true);
-    moduleOp.emitError("MergeComputeNodes stopped after spatial1_scheduled_no_comm; "
-                       "Phase 2 communication realization is not implemented");
-    signalPassFailure();
-  }
-};
-
-} // namespace
-
-} // namespace spatial
-
-std::unique_ptr<Pass> createMergeComputeNodesPass() { return std::make_unique<spatial::MergeComputeNodesPass>(); }
-
-} // namespace onnx_mlir