blazingly faster

.gitignore

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

README.md

@@ -99,15 +99,13 @@ Pass these to `onnx-mlir` when compiling for PIM:
 - `--core-count=<N>` - required positive core count for PIM compilation.
 - `--crossbar-size=<N>` - crossbar width/height. Default in code is `2`.
 - `--crossbar-count=<N>` - crossbars per core. Default in code is `256`.
-- `--pim-merge-scheduler=peft` - merge scheduler. `peft` is the only accepted
-  value in the current code.
 - `--pim-only-codegen` - assume input is already bufferized PIM IR and only run
   the codegen tail.
 - `--pim-emit-json` - also emit `core_*.json` instruction files alongside
   `core_*.pim`.
-- `--pim-export-spatial-dataflow=<none|spatial1|spatial2|spatial3|all>` - control Spatial
-  dataflow CSV reports. The default `all` emits graph, scheduled, and realized
-  snapshots under `reports/`.
+- `--pim-export-spatial-dataflow=<none|spatial1|spatial2|spatial3|spatial4|all>` - control Spatial
+  dataflow CSV reports for the graph, trivially merged graph, scheduled, and
+  realized snapshots under `reports/`.
 - `--use-experimental-conv-impl` - use the alternate convolution lowering.
 - `--ignore-concat-error` - soft-fail a ConcatOp corner case.
 
@@ -170,8 +168,8 @@ Each validation run writes artifacts in the model workspace, for example under
 - `simulation/out.bin` - raw simulator output used for comparison.
 
 The compiler currently dumps dialect snapshots such as `spatial0.mlir`,
-`spatial1_graph.mlir`, `spatial2_scheduled_no_comm.mlir`,
-`spatial3_scheduled.mlir`, `pim0.mlir`, `pim1_buff.mlir`,
+`spatial1_graph.mlir`, `spatial2_trivial_merged.mlir`,
+`spatial3_scheduled_no_comm.mlir`, `spatial4_scheduled.mlir`, `pim0.mlir`, `pim1_buff.mlir`,
 `pim2_coalesced.mlir`, and `pim3_folded.mlir` when an output directory is
 available.
 

requirements.txt

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

src/PIM/Common/CMakeLists.txt

@@ -10,6 +10,7 @@ add_pim_library(OMPimCommon
   IR/ShapeUtils.cpp
   IR/ShapingUtils.cpp
   IR/StaticIntSequence.cpp
+  IR/StaticIntGrid.cpp
   IR/SubviewUtils.cpp
   IR/TensorSliceUtils.cpp
   IR/WeightUtils.cpp

src/PIM/Common/IR/AddressAnalysis.cpp

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

src/PIM/Common/IR/CoreBlockUtils.cpp

@@ -1,3 +1,4 @@
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
@@ -10,7 +11,8 @@
 namespace onnx_mlir {
 
 bool isCoreStaticAddressOp(mlir::Operation* op) {
-  if (mlir::isa<mlir::arith::ConstantOp,
+  if (mlir::isa<mlir::affine::AffineApplyOp,
+                mlir::arith::ConstantOp,
                 mlir::arith::AddIOp,
                 mlir::arith::SubIOp,
                 mlir::arith::MulIOp,

src/PIM/Common/IR/StaticIntGrid.cpp

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

src/PIM/Common/IR/StaticIntGrid.hpp

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

src/PIM/Common/IR/StaticIntSequence.hpp

@@ -36,6 +36,12 @@ public:
   StaticIntSequence slice(size_t begin, size_t count) const;
   StaticIntSequence remap(llvm::ArrayRef<unsigned> indices) const;
   StaticIntSequenceKind getKind() const { return kind; }
+  std::optional<int64_t> getAffineStep() const {
+    if (kind == StaticIntSequenceKind::Uniform)
+      return 0;
+    return kind == StaticIntSequenceKind::Affine
+         ? std::optional<int64_t>(step) : std::nullopt;
+  }
 
   bool operator==(const StaticIntSequence& other) const;
   llvm::hash_code hash() const;

src/PIM/Common/IR/TensorSliceUtils.cpp

@@ -22,7 +22,7 @@ Value extractAxisSlice(
     .getResult();
 }
 
-Value extractStaticSliceOrIdentity(RewriterBase& rewriter,
+Value extractStaticSliceOrIdentity(OpBuilder& rewriter,
                                    Location loc,
                                    Value source,
                                    RankedTensorType resultType,
@@ -52,7 +52,8 @@ Value extractStaticSliceOrIdentity(RewriterBase& rewriter,
   if (isIdentitySlice)
     return source;
 
-  return tensor::ExtractSliceOp::create(rewriter, loc, resultType, source, offsets, sizes, strides).getResult();
+  return rewriter.createOrFold<tensor::ExtractSliceOp>(
+    loc, resultType, source, offsets, sizes, strides);
 }
 
 Value insertStaticSlice(
@@ -68,7 +69,7 @@ Value insertStaticSlice(
     .getResult();
 }
 
-Value extractMixedSliceOrIdentity(RewriterBase &rewriter,
+Value extractMixedSliceOrIdentity(OpBuilder &rewriter,
                                   Location loc,
                                   Value source,
                                   RankedTensorType resultType,

src/PIM/Common/IR/TensorSliceUtils.hpp

@@ -17,7 +17,7 @@ struct MixedSliceGeometry {
 mlir::Value extractAxisSlice(
   mlir::PatternRewriter& rewriter, mlir::Location loc, mlir::Value source, int64_t axis, int64_t offset, int64_t size);
 
-mlir::Value extractStaticSliceOrIdentity(mlir::RewriterBase& rewriter,
+mlir::Value extractStaticSliceOrIdentity(mlir::OpBuilder& rewriter,
                                          mlir::Location loc,
                                          mlir::Value source,
                                          mlir::RankedTensorType resultType,
@@ -31,7 +31,7 @@ mlir::Value insertStaticSlice(mlir::PatternRewriter& rewriter,
                               mlir::Value dest,
                               llvm::ArrayRef<mlir::OpFoldResult> offsets);
 
-mlir::Value extractMixedSliceOrIdentity(mlir::RewriterBase &rewriter,
+mlir::Value extractMixedSliceOrIdentity(mlir::OpBuilder &rewriter,
                                         mlir::Location loc,
                                         mlir::Value source,
                                         mlir::RankedTensorType resultType,

src/PIM/Compiler/CMakeLists.txt

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

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -15,13 +15,6 @@ llvm::cl::opt<PimEmissionTargetType> pimEmissionTarget(
   llvm::cl::init(EmitPimCodegen),
   llvm::cl::cat(OnnxMlirOptions));
 
-llvm::cl::opt<PimMergeSchedulerType>
-  pimMergeScheduler("pim-merge-scheduler",
-                    llvm::cl::desc("Scheduler used by the Spatial merge-compute-nodes pass"),
-                    llvm::cl::values(clEnumValN(MergeSchedulerPeft, "peft", "Use PEFT scheduling")),
-                    llvm::cl::init(MergeSchedulerPeft),
-                    llvm::cl::cat(OnnxMlirOptions));
-
 llvm::cl::opt<PimMemoryReportLevel> pimMemoryReport(
   "pim-memory-report",
   llvm::cl::desc("Emit a human-readable PIM memory planning report"),
@@ -64,9 +57,11 @@ llvm::cl::opt<PimSpatialDataflowExportType> pimExportSpatialDataflow(
   llvm::cl::values(
     clEnumValN(SpatialDataflowExportSpatial1, "spatial1", "Emit spatial1 graph dataflow CSV reports")),
   llvm::cl::values(
-    clEnumValN(SpatialDataflowExportSpatial2, "spatial2", "Emit spatial2 scheduled dataflow CSV reports")),
+    clEnumValN(SpatialDataflowExportSpatial2, "spatial2", "Emit spatial2 trivially merged graph dataflow CSV reports")),
   llvm::cl::values(
-    clEnumValN(SpatialDataflowExportSpatial3, "spatial3", "Emit spatial3 realized dataflow CSV reports")),
+    clEnumValN(SpatialDataflowExportSpatial3, "spatial3", "Emit spatial3 scheduled dataflow CSV reports")),
+  llvm::cl::values(
+    clEnumValN(SpatialDataflowExportSpatial4, "spatial4", "Emit spatial4 realized dataflow CSV reports")),
   llvm::cl::values(clEnumValN(SpatialDataflowExportAll, "all", "Emit all Spatial dataflow CSV reports")),
   llvm::cl::init(SpatialDataflowExportNone),
   llvm::cl::cat(OnnxMlirOptions));

src/PIM/Compiler/PimCompilerOptions.hpp

@@ -20,10 +20,6 @@ typedef enum {
   EmitPimCodegen = 3
 } PimEmissionTargetType;
 
-typedef enum {
-  MergeSchedulerPeft = 0,
-} PimMergeSchedulerType;
-
 typedef enum {
   PimMemoryReportNone = 0,
   PimMemoryReportSummary = 1,
@@ -47,12 +43,12 @@ typedef enum {
   SpatialDataflowExportSpatial1 = 1,
   SpatialDataflowExportSpatial2 = 2,
   SpatialDataflowExportSpatial3 = 3,
-  SpatialDataflowExportAll = 4,
+  SpatialDataflowExportSpatial4 = 4,
+  SpatialDataflowExportAll = 5,
 } PimSpatialDataflowExportType;
 
 extern llvm::cl::OptionCategory OnnxMlirOptions;
 extern llvm::cl::opt<PimEmissionTargetType> pimEmissionTarget;
-extern llvm::cl::opt<PimMergeSchedulerType> pimMergeScheduler;
 extern llvm::cl::opt<PimMemoryReportLevel> pimMemoryReport;
 extern llvm::cl::opt<PimConvLoweringType> pimConvLowering;
 extern llvm::cl::opt<PimSpatialDataflowExportType> pimExportSpatialDataflow;

src/PIM/Compiler/PimCompilerUtils.cpp

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

src/PIM/Compiler/PimMemoryLiveness.cpp

@@ -154,7 +154,10 @@ static OperationOrdering buildOperationOrdering(Operation* coreLikeOp) {
 }
 
 static bool isSupportedAliasOp(Operation* op) {
-  return isa<memref::SubViewOp, memref::CastOp, memref::CollapseShapeOp, memref::ExpandShapeOp>(op);
+  return isa<memref::SubViewOp,
+             memref::CastOp,
+             memref::CollapseShapeOp,
+             memref::ExpandShapeOp>(op);
 }
 
 static bool isRuntimeMemoryTouchOp(Operation* op) {

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

@@ -13,6 +13,7 @@
 #include <utility>
 
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
@@ -374,9 +375,6 @@ extractGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
   auto physicalType = mlir::dyn_cast<mlir::RankedTensorType>(physicalBatch.getType());
   if (!physicalType || physicalType.getRank() != fragmentType.getRank() + 1)
     return mlir::failure();
-  mlir::SmallVector<int64_t> selectedShape {1};
-  llvm::append_range(selectedShape, fragmentType.getShape());
-  auto selectedType = mlir::RankedTensorType::get(selectedShape, fragmentType.getElementType(), fragmentType.getEncoding());
   mlir::SmallVector<mlir::OpFoldResult> offsets {slot};
   mlir::SmallVector<mlir::OpFoldResult> sizes {rewriter.getIndexAttr(1)};
   mlir::SmallVector<mlir::OpFoldResult> strides {rewriter.getIndexAttr(1)};
@@ -385,11 +383,8 @@ extractGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
     sizes.push_back(rewriter.getIndexAttr(dim));
     strides.push_back(rewriter.getIndexAttr(1));
   }
-  mlir::Value selected = mlir::tensor::ExtractSliceOp::create(rewriter, loc, selectedType, physicalBatch, offsets, sizes, strides);
-  mlir::SmallVector<mlir::ReassociationIndices> reassociation {{0, 1}};
-  for (int64_t dim = 2; dim <= fragmentType.getRank(); ++dim)
-    reassociation.push_back({dim});
-  return mlir::tensor::CollapseShapeOp::create(rewriter, loc, fragmentType, selected, reassociation).getResult();
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, physicalBatch, fragmentType, {offsets, sizes, strides});
 }
 
 template <typename BodyFn>

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

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

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -79,6 +79,57 @@ materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location
   return createRowStripAssemblyBlueprint(rowStripValue.storage, rowStripValue.logicalType, rewriter, loc);
 }
 
+static FailureOr<Value> lowerDenseBatchBiasAdd(Value input, Value bias, RankedTensorType resultType,
+                                               PatternRewriter& rewriter, Location loc) {
+  auto producer = input.getDefiningOp<spatial::SpatGraphComputeBatch>();
+  auto inputType = dyn_cast<RankedTensorType>(input.getType());
+  auto biasType = dyn_cast<RankedTensorType>(bias.getType());
+  if (!producer || !inputType || !biasType || !inputType.hasStaticShape() || !biasType.hasStaticShape()
+      || !resultType.hasStaticShape() || inputType.getDimSize(0) != producer.getLaneCount()
+      || biasType.getDimSize(0) != producer.getLaneCount() || resultType.getDimSize(0) != producer.getLaneCount())
+    return failure();
+  auto inputFragmentType = spatial::getGraphBatchFragmentType(inputType, producer.getLaneCount());
+  auto outputFragmentType = spatial::getGraphBatchFragmentType(resultType, producer.getLaneCount());
+  if (failed(inputFragmentType) || failed(outputFragmentType) || inputFragmentType->getRank() != biasType.getRank()
+      || inputFragmentType->getDimSize(0) != 1 || inputFragmentType->getShape().drop_front() != biasType.getShape().drop_front()
+      || inputFragmentType->getRank() != outputFragmentType->getRank() + 1)
+    return failure();
+  for (auto [inputDim, outputDim] : llvm::zip(inputFragmentType->getShape().drop_front(), outputFragmentType->getShape()))
+    if (outputDim > inputDim)
+      return failure();
+
+  auto batch = createSpatComputeBatch(rewriter, loc, TypeRange {resultType}, producer.getLaneCount(), {}, ValueRange {input, bias},
+    [&](detail::SpatComputeBatchBodyArgs args) -> LogicalResult {
+      FailureOr<Value> fragment = extractGraphBatchPhysicalFragment(rewriter, loc, args.inputs[0], args.lane, *inputFragmentType);
+      if (failed(fragment))
+        return failure();
+      MixedSliceGeometry biasSlice;
+      for (int64_t dim : inputFragmentType->getShape()) {
+        biasSlice.offsets.push_back(biasSlice.offsets.empty() ? OpFoldResult(args.lane) : rewriter.getIndexAttr(0));
+        biasSlice.sizes.push_back(rewriter.getIndexAttr(dim));
+        biasSlice.strides.push_back(rewriter.getIndexAttr(1));
+      }
+      Value biasFragment = extractMixedSliceOrIdentity(rewriter, loc, args.inputs[1], *inputFragmentType, biasSlice);
+      if (!biasFragment)
+        return failure();
+      Value added = spatial::SpatVAddOp::create(rewriter, loc, *inputFragmentType, *fragment, biasFragment);
+      MixedSliceGeometry outputSlice;
+      outputSlice.offsets.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(0));
+      outputSlice.sizes.push_back(rewriter.getIndexAttr(1));
+      outputSlice.strides.assign(inputFragmentType->getRank(), rewriter.getIndexAttr(1));
+      for (int64_t dim : outputFragmentType->getShape())
+        outputSlice.sizes.push_back(rewriter.getIndexAttr(dim));
+      Value output = extractMixedSliceOrIdentity(rewriter, loc, added, *outputFragmentType, outputSlice);
+      if (!output)
+        return failure();
+      publishGraphBatchPhysicalFragment(rewriter, loc, output, args.outputs.front(), args.lane);
+      return success();
+    });
+  if (failed(batch))
+    return failure();
+  return batch->getResult(0);
+}
+
 struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LowerSpatialPlansPass)
 
@@ -236,6 +287,13 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
           signalPassFailure();
           return;
         }
+        if (planOp.getInput().getDefiningOp<spatial::SpatGraphComputeBatch>()) {
+          FailureOr<Value> lowered = lowerDenseBatchBiasAdd(planOp.getInput(), *denseBias, resultType, rewriter, planOp.getLoc());
+          if (succeeded(lowered)) {
+            rewriter.replaceOp(planOp, *lowered);
+            continue;
+          }
+        }
         auto computeOp = createSpatCompute<2>(rewriter,
                                               planOp.getLoc(),
                                               planOp.getOutput().getType(),

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

@@ -20,6 +20,7 @@
 
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp"
 #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
@@ -1355,20 +1356,11 @@ static Value createWeightTile(Value packedWeights,
   SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
                                    rewriter.getIndexAttr(tiling.tileInputRows),
                                    rewriter.getIndexAttr(tiling.tileOutputChannels)};
-  auto sliceType =
-    RankedTensorType::get({1, tiling.tileInputRows, tiling.tileOutputChannels}, packedWeightType.getElementType());
-  Value slice = tensor::ExtractSliceOp::create(
-    rewriter, loc, sliceType, packedWeights, offsets, sizes, getUnitStrides(rewriter, 3));
   auto collapsedType =
     RankedTensorType::get({tiling.tileInputRows, tiling.tileOutputChannels}, packedWeightType.getElementType());
-  return tensor::CollapseShapeOp::create(rewriter,
-                                         loc,
-                                         collapsedType,
-                                         slice,
-                                         SmallVector<ReassociationIndices> {
-                                           {0, 1},
-                                           {2}
-  });
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, packedWeights, collapsedType,
+    {offsets, sizes, getUnitStrides(rewriter, 3)});
 }
 
 static Value createBiasTile(
@@ -1676,8 +1668,6 @@ struct ConvGemmPlan {
   int64_t effectiveMaxParallelPixels;
   int64_t packedNumRows;
 
-  RankedTensorType im2colType;
-  RankedTensorType im2colRowType;
   RankedTensorType gemmInputRowsType;
   RankedTensorType wFlatType;
   RankedTensorType wTransType;
@@ -1718,52 +1708,6 @@ static PreparedConvInput prepareInputForIm2Col(const ConvLoweringState& state,
   return {paddedInputOp.getResult(0), paddedType};
 }
 
-static Value createPaddedRows(Value rows,
-                              RankedTensorType rowsType,
-                              int64_t paddedRows,
-                              PatternRewriter& rewriter,
-                              Location loc) {
-  if (rowsType.getDimSize(0) == paddedRows)
-    return rows;
-
-  auto paddedType =
-    RankedTensorType::get({paddedRows, rowsType.getDimSize(1)}, rowsType.getElementType(), rowsType.getEncoding());
-  return createZeroPaddedTensor(
-    rows, paddedType, {0, 0}, {paddedRows - rowsType.getDimSize(0), 0}, rewriter, loc);
-}
-
-static Value packRowsForParallelGemm(
-  Value rows, RankedTensorType rowsType, int64_t packFactor, PatternRewriter& rewriter, Location loc) {
-  if (packFactor == 1)
-    return rows;
-
-  const int64_t paddedNumRows = ceilIntegerDivide(rowsType.getDimSize(0), packFactor) * packFactor;
-  const int64_t packedNumRows = paddedNumRows / packFactor;
-  const int64_t rowWidth = rowsType.getDimSize(1);
-  auto groupedType =
-    RankedTensorType::get({packedNumRows, packFactor, rowWidth}, rowsType.getElementType(), rowsType.getEncoding());
-  auto packedType =
-    RankedTensorType::get({packedNumRows, packFactor * rowWidth}, rowsType.getElementType(), rowsType.getEncoding());
-
-  Value padded = createPaddedRows(rows, rowsType, paddedNumRows, rewriter, loc);
-  Value grouped = tensor::ExpandShapeOp::create(rewriter,
-                                                loc,
-                                                groupedType,
-                                                padded,
-                                                SmallVector<ReassociationIndices> {
-                                                  {0, 1},
-                                                  {2}
-  });
-  return tensor::CollapseShapeOp::create(rewriter,
-                                         loc,
-                                         packedType,
-                                         grouped,
-                                         SmallVector<ReassociationIndices> {
-                                           {0},
-                                           {1, 2}
-  });
-}
-
 static Value unpackRowsFromParallelGemm(Value packedRows,
                                         RankedTensorType packedRowsType,
                                         int64_t unpackedRows,
@@ -2166,8 +2110,6 @@ buildConvGemmPlan(const ConvLoweringState& state,
 
   auto elemType = state.xType.getElementType();
   auto outElemType = state.outType.getElementType();
-  plan.im2colType = RankedTensorType::get({plan.chunkNumPatches, plan.patchSize}, elemType);
-  plan.im2colRowType = RankedTensorType::get({1, plan.patchSize}, elemType);
   plan.gemmInputRowsType =
     RankedTensorType::get({plan.packedNumRows, plan.effectiveMaxParallelPixels * plan.patchSize}, elemType);
   plan.wFlatType = RankedTensorType::get({state.numChannelsOut, plan.patchSize}, state.wType.getElementType());
@@ -2185,44 +2127,99 @@ static Value createIm2colRows(const ConvLoweringState& state,
                               const ConvGemmPlan& plan,
                               PatternRewriter& rewriter,
                               Location loc) {
-  constexpr size_t numInputs = 1;
-  auto im2colComputeOp =
-    createSpatCompute<numInputs>(rewriter, loc, TypeRange {plan.gemmInputRowsType}, {}, preparedInput.value, [&](Value xArg) {
-      auto elemType = preparedInput.type.getElementType();
-      // Keep the standard im2col view of convolution, flipped so filters sit in
-      // B / crossbar columns:
-      //   A (im2col):  [numPatches, patchSize]  -- one row per output spatial position
-      //   B (weights): [patchSize, cOut]
-      //   Gemm output: [numPatches, cOut]
-      Value im2colInit = tensor::EmptyOp::create(rewriter, loc, plan.im2colType.getShape(), elemType);
+  if (plan.gemmInputRowsType.getDimSize(1) > crossbarSize.getValue()) {
+    assert(plan.effectiveMaxParallelPixels == 1 && "multi-crossbar im2col rows cannot pack pixels");
+    auto compute = createSpatCompute<1>(
+      rewriter, loc, TypeRange {plan.gemmInputRowsType}, {}, preparedInput.value, [&](Value input) {
+        auto elemType = preparedInput.type.getElementType();
+        Value empty = tensor::EmptyOp::create(rewriter, loc, plan.gemmInputRowsType.getShape(), elemType);
+        Operation *anchor = rewriter.getInsertionBlock()->getParentOp();
+        Value c0 = getOrCreateIndexConstant(rewriter, anchor, 0);
+        Value c1 = getOrCreateIndexConstant(rewriter, anchor, 1);
+        Value upper = getOrCreateIndexConstant(rewriter, anchor, plan.chunkNumPatches);
+        auto patchType = RankedTensorType::get(
+          {1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType);
+        auto rowType = RankedTensorType::get({plan.patchSize}, elemType);
+        auto loop = buildNormalizedScfFor(
+          rewriter, loc, c0, upper, c1, ValueRange {empty},
+          [&](OpBuilder &, Location nestedLoc, Value patchIndex, ValueRange iterArgs,
+              SmallVectorImpl<Value> &yielded) {
+            Value batchIndex = affineAddFloorDivConst(
+              rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchor);
+            Value batchPatchIndex = affineAddModConst(
+              rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchor);
+            Value outHeight = affineFloorDivConst(
+              rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchor);
+            Value outWidth = affineModConst(
+              rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchor);
+            Value patch = createConvInputPatch(
+              input, patchType, batchIndex, c0,
+              affineMulConst(rewriter, nestedLoc, outHeight, state.strideHeight, anchor),
+              affineMulConst(rewriter, nestedLoc, outWidth, state.strideWidth, anchor),
+              state.dilationHeight, state.dilationWidth, rewriter, nestedLoc);
+            Value row = tensor::CollapseShapeOp::create(
+              rewriter, nestedLoc, rowType, patch,
+              SmallVector<ReassociationIndices> {{0, 1, 2, 3}});
+            Value next = tensor::InsertSliceOp::create(
+              rewriter, nestedLoc, row, iterArgs.front(),
+              SmallVector<OpFoldResult> {patchIndex, rewriter.getIndexAttr(0)},
+              SmallVector<OpFoldResult> {rewriter.getIndexAttr(1), rewriter.getIndexAttr(plan.patchSize)},
+              getUnitStrides(rewriter, 2));
+            yielded.push_back(next);
+            return success();
+          });
+        if (failed(loop))
+          return failure();
+        spatial::SpatYieldOp::create(rewriter, loc, loop->results.front());
+        return success();
+      });
+    assert(succeeded(compute) && "Conv im2col compute construction must succeed");
+    return compute->getResult(0);
+  }
+
+  auto elemType = preparedInput.type.getElementType();
+  auto packedRowType = RankedTensorType::get(
+    {plan.effectiveMaxParallelPixels * plan.patchSize}, elemType, plan.gemmInputRowsType.getEncoding());
+  auto zeroAttr = DenseElementsAttr::get(packedRowType, rewriter.getZeroAttr(elemType));
+  Value zeroRow = getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), zeroAttr, packedRowType);
+  auto im2colComputeOp = createSpatComputeBatch(
+    rewriter,
+    loc,
+    TypeRange {plan.gemmInputRowsType},
+    plan.packedNumRows,
+    {},
+    ValueRange {preparedInput.value, zeroRow},
+    [&](detail::SpatComputeBatchBodyArgs args) {
       Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
       Value c0 = getOrCreateIndexConstant(rewriter, anchorOp, 0);
       Value c1 = getOrCreateIndexConstant(rewriter, anchorOp, 1);
+      Value cPack = getOrCreateIndexConstant(rewriter, anchorOp, plan.effectiveMaxParallelPixels);
       Value cNumPatches = getOrCreateIndexConstant(rewriter, anchorOp, plan.chunkNumPatches);
+      Value laneStart = affineMulConst(rewriter, loc, args.lane, plan.effectiveMaxParallelPixels, anchorOp);
+      Value remaining = arith::SubIOp::create(rewriter, loc, cNumPatches, laneStart);
+      Value isPartial = arith::CmpIOp::create(rewriter, loc, arith::CmpIPredicate::ult, remaining, cPack);
+      Value lanePatches = arith::SelectOp::create(rewriter, loc, isPartial, remaining, cPack);
+      auto patchType = RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType);
+      auto patchRowType = RankedTensorType::get({plan.patchSize}, elemType);
 
-      auto im2colLoop = buildNormalizedScfFor(
+      auto rowLoop = buildNormalizedScfFor(
         rewriter,
         loc,
         c0,
-        cNumPatches,
+        lanePatches,
         c1,
-        ValueRange {im2colInit},
-        [&](OpBuilder&, Location nestedLoc, Value patchIndex, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
-          Value im2colAcc = iterArgs.front();
+        ValueRange {args.inputs[1]},
+        [&](OpBuilder&, Location nestedLoc, Value copyIndex, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+          Value patchIndex = arith::AddIOp::create(rewriter, nestedLoc, laneStart, copyIndex);
           Value batchIndex =
             affineAddFloorDivConst(rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchorOp);
           Value batchPatchIndex =
             affineAddModConst(rewriter, nestedLoc, patchIndex, plan.chunkStart, plan.numPatchesPerBatch, anchorOp);
           Value outHeightIndex = affineFloorDivConst(rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchorOp);
           Value outWidthIndex = affineModConst(rewriter, nestedLoc, batchPatchIndex, state.outWidth, anchorOp);
-          Value inputHeightOffset =
-            affineMulConst(rewriter, nestedLoc, outHeightIndex, state.strideHeight, anchorOp);
-          Value inputWidthOffset =
-            affineMulConst(rewriter, nestedLoc, outWidthIndex, state.strideWidth, anchorOp);
-
-          auto patchType =
-            RankedTensorType::get({1, state.numChannelsIn, state.wHeight, state.wWidth}, elemType);
-          Value patch = createConvInputPatch(xArg,
+          Value inputHeightOffset = affineMulConst(rewriter, nestedLoc, outHeightIndex, state.strideHeight, anchorOp);
+          Value inputWidthOffset = affineMulConst(rewriter, nestedLoc, outWidthIndex, state.strideWidth, anchorOp);
+          Value patch = createConvInputPatch(args.inputs.front(),
                                              patchType,
                                              batchIndex,
                                              c0,
@@ -2232,34 +2229,27 @@ static Value createIm2colRows(const ConvLoweringState& state,
                                              state.dilationWidth,
                                              rewriter,
                                              nestedLoc);
-          Value row = tensor::CollapseShapeOp::create(rewriter,
-                                                      nestedLoc,
-                                                      plan.im2colRowType,
-                                                      patch,
-                                                      SmallVector<ReassociationIndices> {
-                                                        {0},
-                                                        {1, 2, 3}
+          Value patchRow = tensor::CollapseShapeOp::create(rewriter,
+                                                           nestedLoc,
+                                                           patchRowType,
+                                                           patch,
+                                                           SmallVector<ReassociationIndices> {
+                                                             {0, 1, 2, 3}
           });
-
-          SmallVector<OpFoldResult> rowOffsets {patchIndex, rewriter.getIndexAttr(0)};
-          SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(1), rewriter.getIndexAttr(plan.patchSize)};
-          Value next = tensor::InsertSliceOp::create(
-            rewriter, nestedLoc, row, im2colAcc, rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
+          Value rowOffset = affineMulConst(rewriter, nestedLoc, copyIndex, plan.patchSize, anchorOp);
+          Value next = tensor::InsertSliceOp::create(rewriter,
+                                                     nestedLoc,
+                                                     patchRow,
+                                                     iterArgs.front(),
+                                                     SmallVector<OpFoldResult> {rowOffset},
+                                                     SmallVector<OpFoldResult> {rewriter.getIndexAttr(plan.patchSize)},
+                                                     getUnitStrides(rewriter, 1));
           yielded.push_back(next);
           return success();
         });
-      if (failed(im2colLoop))
+      if (failed(rowLoop))
         return failure();
-
-      Value gemmInputRows = im2colLoop->results.front();
-      // Pack N old im2col rows into one longer row so one GEMM can cover N
-      // pixels in parallel. The corresponding packed weight matrix contains N
-      // block-diagonal copies of W^T, and the packed output must be unpacked
-      // back to one row per spatial patch.
-      if (plan.effectiveMaxParallelPixels != 1)
-        gemmInputRows = packRowsForParallelGemm(gemmInputRows, plan.im2colType, plan.effectiveMaxParallelPixels, rewriter, loc);
-
-      spatial::SpatYieldOp::create(rewriter, loc, gemmInputRows);
+      publishGraphBatchPhysicalFragment(rewriter, loc, rowLoop->results.front(), args.outputs.front(), args.lane);
       return success();
     });
 

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

@@ -17,6 +17,7 @@
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/Diagnostics.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
@@ -501,9 +502,9 @@ static Value extractReductionPiece(Value partialPiecesArg,
   SmallVector<OpFoldResult> pieceSizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
   SmallVector<OpFoldResult> pieceOffsets {
     createPartialGroupOffset(hSlice, kSlice, numKSlices, numOutRows, rewriter, loc), rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-  auto selectedType = RankedTensorType::get({numOutRows, 1, static_cast<int64_t>(crossbarSize.getValue())}, pieceType.getElementType());
-  Value selected = tensor::ExtractSliceOp::create(rewriter, loc, selectedType, partialPiecesArg, pieceOffsets, pieceSizes, unitStrides);
-  return tensor::CollapseShapeOp::create(rewriter, loc, pieceType, selected, SmallVector<ReassociationIndices> {{0, 1}, {2}});
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, partialPiecesArg, pieceType,
+    {pieceOffsets, pieceSizes, unitStrides});
 }
 
 static Value reducePartialPiecesForHSlice(Value partialPiecesArg,

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

@@ -9,6 +9,7 @@
 
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
@@ -299,22 +300,14 @@ static Value extractBatchedATile(Value a,
                                  RankedTensorType aTileType,
                                  PatternRewriter& rewriter,
                                  Location loc) {
-  auto aSliceType = RankedTensorType::get({1, 1, aTileType.getDimSize(1)}, aTileType.getElementType());
   Value sourceBatchIndex =
     mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
   SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), row, kOffset};
   SmallVector<OpFoldResult> sizes {
     rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(aTileType.getDimSize(1))};
-  auto slice =
-    tensor::ExtractSliceOp::create(rewriter, loc, aSliceType, a, offsets, sizes, getUnitStrides(rewriter, 3));
-  return tensor::CollapseShapeOp::create(rewriter,
-                                         loc,
-                                         aTileType,
-                                         slice,
-                                         SmallVector<ReassociationIndices> {
-                                           {0, 1},
-                                           {2}
-  });
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, a, aTileType,
+    {offsets, sizes, getUnitStrides(rewriter, 3)});
 }
 
 static Value extractBatchedBTile(Value b,
@@ -326,24 +319,15 @@ static Value extractBatchedBTile(Value b,
                                  RankedTensorType bTileType,
                                  PatternRewriter& rewriter,
                                  Location loc) {
-  auto bSliceType =
-    RankedTensorType::get({1, bTileType.getDimSize(0), bTileType.getDimSize(1)}, bTileType.getElementType());
   Value sourceBatchIndex =
     mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
   SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), kOffset, hOffset};
   SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(1),
                                    rewriter.getIndexAttr(bTileType.getDimSize(0)),
                                    rewriter.getIndexAttr(bTileType.getDimSize(1))};
-  auto slice =
-    tensor::ExtractSliceOp::create(rewriter, loc, bSliceType, b, offsets, sizes, getUnitStrides(rewriter, 3));
-  return tensor::CollapseShapeOp::create(rewriter,
-                                         loc,
-                                         bTileType,
-                                         slice,
-                                         SmallVector<ReassociationIndices> {
-                                           {0, 1},
-                                           {2}
-  });
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, b, bTileType,
+    {offsets, sizes, getUnitStrides(rewriter, 3)});
 }
 
 static Value getBatchLaneIndex(
@@ -448,22 +432,14 @@ static Value extractDynamicBatchedRowVector(Value matrix,
                                             RankedTensorType vectorType,
                                             PatternRewriter& rewriter,
                                             Location loc) {
-  auto rowSliceType = RankedTensorType::get({1, 1, vectorType.getDimSize(1)}, vectorType.getElementType());
   Value sourceBatchIndex =
     mapOutputBatchIndexToSourceBatchIndex(outputBatchIndex, sourceBatchShape, outputBatchShape, rewriter, loc);
   SmallVector<OpFoldResult> offsets {OpFoldResult(sourceBatchIndex), row, rewriter.getIndexAttr(0)};
   SmallVector<OpFoldResult> sizes {
     rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(vectorType.getDimSize(1))};
-  auto rowSlice =
-    tensor::ExtractSliceOp::create(rewriter, loc, rowSliceType, matrix, offsets, sizes, getUnitStrides(rewriter, 3));
-  return tensor::CollapseShapeOp::create(rewriter,
-                                         loc,
-                                         vectorType,
-                                         rowSlice,
-                                         SmallVector<ReassociationIndices> {
-                                           {0, 1},
-                                           {2}
-  });
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, matrix, vectorType,
+    {offsets, sizes, getUnitStrides(rewriter, 3)});
 }
 
 static FailureOr<spatial::SpatComputeBatch> createBatchedVvdmulBatch(Value a,
@@ -519,7 +495,6 @@ static FailureOr<Value> createBatchedDynamicOutputCompute(Value scalarPieces,
   const int64_t numOutRows = outType.getDimSize(1);
   const int64_t numOutCols = outType.getDimSize(2);
   auto scalarType = RankedTensorType::get({1, 1}, outType.getElementType());
-  auto outputScalarType = RankedTensorType::get({1, 1, 1}, outType.getElementType());
 
   auto computeOp = createSpatCompute<1>(
     rewriter, loc, TypeRange {outType}, {}, ValueRange {scalarPieces}, [&](Value pieces) -> LogicalResult {
@@ -545,20 +520,12 @@ static FailureOr<Value> createBatchedDynamicOutputCompute(Value scalarPieces,
           FailureOr<Value> scalar = extractGraphBatchPhysicalFragment(rewriter, nestedLoc, pieces, lane, scalarType);
           if (failed(scalar))
             return failure();
-          Value expanded = tensor::ExpandShapeOp::create(rewriter,
-                                                         nestedLoc,
-                                                         outputScalarType,
-                                                         *scalar,
-                                                         SmallVector<ReassociationIndices> {
-                                                           {0},
-                                                           {1, 2}
-          });
           SmallVector<OpFoldResult> outputOffsets {batch, row, column};
           SmallVector<OpFoldResult> outputSizes = {
             rewriter.getIndexAttr(1), rewriter.getIndexAttr(1), rewriter.getIndexAttr(1)};
           Value next =
             tensor::InsertSliceOp::create(
-              rewriter, nestedLoc, expanded, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
+              rewriter, nestedLoc, *scalar, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
               .getResult();
           yielded.push_back(next);
           return success();
@@ -591,9 +558,9 @@ static Value extractBatchedReductionPiece(Value partialPiecesArg,
   Value pieceOffset = arith::AddIOp::create(rewriter, loc, batchAndHSlice, kOffset);
   SmallVector<OpFoldResult> offsets {pieceOffset, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
   SmallVector<OpFoldResult> sizes {rewriter.getIndexAttr(numOutRows), rewriter.getIndexAttr(1), rewriter.getIndexAttr(crossbarSize.getValue())};
-  auto selectedType = RankedTensorType::get({numOutRows, 1, static_cast<int64_t>(crossbarSize.getValue())}, pieceType.getElementType());
-  Value selected = tensor::ExtractSliceOp::create(rewriter, loc, selectedType, partialPiecesArg, offsets, sizes, getUnitStrides(rewriter, 3));
-  return tensor::CollapseShapeOp::create(rewriter, loc, pieceType, selected, SmallVector<ReassociationIndices> {{0, 1}, {2}});
+  return extractMixedSliceOrIdentity(
+    rewriter, loc, partialPiecesArg, pieceType,
+    {offsets, sizes, getUnitStrides(rewriter, 3)});
 }
 
 static Value reduceBatchedPartialPiecesForHSlice(Value partialPiecesArg,
@@ -640,8 +607,6 @@ static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
       const int64_t numOutHSlices = ceilIntegerDivide(outType.getDimSize(2), crossbarSize.getValue());
       auto pieceType = RankedTensorType::get({numOutRows, static_cast<int64_t>(crossbarSize.getValue())},
                                              partialPiecesType.getElementType());
-      auto outputSliceType = RankedTensorType::get({1, numOutRows, static_cast<int64_t>(crossbarSize.getValue())},
-                                                   partialPiecesType.getElementType());
 
       Value outputInit =
         tensor::EmptyOp::create(rewriter, loc, paddedOutType.getShape(), paddedOutType.getElementType()).getResult();
@@ -671,14 +636,6 @@ static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
               Value outputAcc = hIterArgs.front();
               Value reduced = reduceBatchedPartialPiecesForHSlice(
                 partialPiecesArg, batch, hSlice, pieceType, numKSlices, numOutHSlices, numOutRows, rewriter, hLoc);
-              Value expandedReduced = tensor::ExpandShapeOp::create(rewriter,
-                                                                    hLoc,
-                                                                    outputSliceType,
-                                                                    reduced,
-                                                                    SmallVector<ReassociationIndices> {
-                                                                      {0, 1},
-                                                                      {2}
-              });
               Value hOffset = affineMulConst(
                 rewriter, hLoc, hSlice, crossbarSize.getValue(), rewriter.getInsertionBlock()->getParentOp());
               SmallVector<OpFoldResult> outputOffsets {batch, rewriter.getIndexAttr(0), hOffset};
@@ -687,7 +644,7 @@ static FailureOr<Value> createBatchedReductionCompute(Value partialPieces,
                                                      rewriter.getIndexAttr(crossbarSize.getValue())};
               Value next =
                 tensor::InsertSliceOp::create(
-                  rewriter, hLoc, expandedReduced, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
+                  rewriter, hLoc, reduced, outputAcc, outputOffsets, outputSizes, getUnitStrides(rewriter, 3))
                   .getResult();
               hYielded.push_back(next);
               return success();

src/PIM/Conversion/SpatialToPim/Patterns/ChannelLowering.cpp

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

src/PIM/Conversion/SpatialToPim/SpatialToPimPass.cpp

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

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

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

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

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

src/PIM/Dialect/Spatial/CMakeLists.txt

@@ -20,11 +20,13 @@ add_pim_library(SpatialOps
   Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp
   Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp
   Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp
+  Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp
   Transforms/MergeComputeNodes/ScheduledComputeReport.cpp
   Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp
   Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp
   Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp
   Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp
+  Transforms/TrivialGraphComputeMergePass.cpp
 
   EXCLUDE_FROM_OM_LIBS
 

src/PIM/Dialect/Spatial/Spatial.td

@@ -49,6 +49,7 @@ class SpatComputeLikeBase<string mnemonic> : SpatOp<mnemonic,
 }
 
 def SpatGraphCompute : SpatComputeLikeBase<"graph_compute"> {
+  let hasCanonicalizer = 1;
   let extraClassDeclaration = [{
     std::optional<::mlir::BlockArgument> getWeightArgument(unsigned idx);
     std::optional<::mlir::BlockArgument> getInputArgument(unsigned idx);
@@ -186,7 +187,8 @@ def SpatDeferredCommunicationOp : SpatOp<"deferred_communication", [SingleBlock]
   let summary = "Temporary scheduled payload derivation placeholder";
 
   let arguments = (ins
-    Variadic<SpatTensor>:$sources
+    Variadic<SpatTensor>:$sources,
+    OptionalAttr<I64Attr>:$specialization_count
   );
 
   let results = (outs
@@ -199,6 +201,22 @@ def SpatDeferredCommunicationOp : SpatOp<"deferred_communication", [SingleBlock]
   let hasCustomAssemblyFormat = 1;
 }
 
+def SpatDeferredSourceSelectOp : SpatOp<"deferred_source_select", []> {
+  let summary = "Select a deferred tensor source with a statically analyzable index";
+
+  let arguments = (ins
+    Index:$selector,
+    Variadic<SpatTensor>:$sources
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
 def SpatExtractRowsOp : SpatOp<"extract_rows", []> {
   let summary = "Extract every row of a rank-2 tensor as separate rank-2 row tensors";
 

src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp

@@ -457,39 +457,80 @@ void SpatDeferredCommunicationOp::print(OpAsmPrinter& printer) {
   printCompressedValueSequence(printer, getSources());
   printer.printOptionalAttrDict((*this)->getAttrs());
   printer << " : ";
-  printer.printFunctionalType(getSources().getTypes(), getOperation()->getResultTypes());
+  printCompressedTypeList(
+    printer, getSources().getTypes(), ListDelimiter::Paren);
+  printer << " -> ";
+  printCompressedTypeSequence(printer, getOperation()->getResultTypes());
   printer << " ";
   printer.printRegion(getBody(), /*printEntryBlockArgs=*/false);
 }
 
 ParseResult SpatDeferredCommunicationOp::parse(OpAsmParser& parser, OperationState& result) {
   SmallVector<OpAsmParser::UnresolvedOperand> sources;
-  Type functionTypeStorage;
+  SmallVector<Type> sourceTypes, outputTypes;
 
   if (parseCompressedOperandSequence(parser, sources) || parser.parseOptionalAttrDict(result.attributes)
-      || parser.parseColon() || parser.parseType(functionTypeStorage))
+      || parser.parseColon()
+      || parseCompressedRepeatedList(
+        parser, ListDelimiter::Paren, sourceTypes,
+        [&](Type& type) { return parser.parseType(type); })
+      || parser.parseArrow()
+      || parseCompressedTypeSequence(
+        parser, outputTypes, /*allowEmpty=*/false))
     return failure();
 
-  auto functionType = dyn_cast<FunctionType>(functionTypeStorage);
-  if (!functionType)
-    return parser.emitError(parser.getCurrentLocation(), "expected deferred communication function type");
-  if (sources.size() != functionType.getNumInputs())
+  if (sources.size() != sourceTypes.size())
     return parser.emitError(parser.getCurrentLocation(), "number of sources and source types must match");
 
-  if (parser.resolveOperands(sources, functionType.getInputs(), parser.getCurrentLocation(), result.operands))
+  if (parser.resolveOperands(sources, sourceTypes, parser.getCurrentLocation(), result.operands))
     return failure();
-  result.addTypes(functionType.getResults());
+  result.addTypes(outputTypes);
 
   Region* body = result.addRegion();
   SmallVector<OpAsmParser::Argument> bodyArgs;
-  for (Type type : functionType.getInputs()) {
+  for (Type type : sourceTypes) {
     OpAsmParser::Argument argument;
     argument.type = type;
     bodyArgs.push_back(argument);
   }
+  if (auto count = dyn_cast_or_null<IntegerAttr>(
+        result.attributes.get("specialization_count"));
+      count && count.getInt() > 1) {
+    OpAsmParser::Argument argument;
+    argument.type = parser.getBuilder().getIndexType();
+    bodyArgs.push_back(argument);
+  }
   return parser.parseRegion(*body, bodyArgs);
 }
 
+void SpatDeferredSourceSelectOp::print(OpAsmPrinter& printer) {
+  printer << " " << getSelector() << " of ";
+  printCompressedValueSequence(printer, getSources());
+  printer.printOptionalAttrDict((*this)->getAttrs());
+  printer << " : " << getOutput().getType();
+}
+
+ParseResult SpatDeferredSourceSelectOp::parse(
+    OpAsmParser& parser, OperationState& result) {
+  OpAsmParser::UnresolvedOperand selector;
+  SmallVector<OpAsmParser::UnresolvedOperand> sources;
+  Type outputType;
+  if (parser.parseOperand(selector) || parser.parseKeyword("of")
+      || parseCompressedOperandSequence(parser, sources)
+      || parser.parseOptionalAttrDict(result.attributes)
+      || parser.parseColon() || parser.parseType(outputType))
+    return failure();
+  if (parser.resolveOperand(selector, parser.getBuilder().getIndexType(),
+                            result.operands))
+    return failure();
+  SmallVector<Type> sourceTypes(sources.size(), outputType);
+  if (parser.resolveOperands(sources, sourceTypes,
+                             parser.getCurrentLocation(), result.operands))
+    return failure();
+  result.addTypes(outputType);
+  return success();
+}
+
 void SpatExtractRowsOp::print(OpAsmPrinter& printer) {
   printer << " ";
   printer.printOperand(getInput());

src/PIM/Dialect/Spatial/SpatialOpsCanonicalization.cpp

@@ -42,6 +42,33 @@ LogicalResult SpatGraphCompute::fold(FoldAdaptor adaptor, ::llvm::SmallVectorImp
   return foldComputeLike(*this, results);
 }
 
+struct RemoveUnusedGraphComputeInputsPattern : OpRewritePattern<SpatGraphCompute> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(SpatGraphCompute compute, PatternRewriter& rewriter) const override {
+    SmallVector<unsigned> unusedInputs;
+    for (unsigned index = 0; index < compute.getInputs().size(); ++index) {
+      auto argument = compute.getInputArgument(index);
+      if (argument && argument->use_empty())
+        unusedInputs.push_back(index);
+    }
+    if (unusedInputs.empty())
+      return failure();
+
+    rewriter.modifyOpInPlace(compute, [&] {
+      for (unsigned index : llvm::reverse(unusedInputs)) {
+        compute.getBody().front().eraseArgument(compute.getWeights().size() + index);
+        compute.getInputsMutable().erase(index);
+      }
+    });
+    return success();
+  }
+};
+
+void SpatGraphCompute::getCanonicalizationPatterns(RewritePatternSet& results, MLIRContext* context) {
+  results.add<RemoveUnusedGraphComputeInputsPattern>(context);
+}
+
 LogicalResult SpatScheduledCompute::fold(FoldAdaptor adaptor, ::llvm::SmallVectorImpl<::mlir::OpFoldResult>& results) {
   return foldComputeLike(*this, results);
 }

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -225,33 +225,6 @@ static LogicalResult verifyOnlyConstantExternalValues(Operation* ownerOp, Region
   return success(!hasFailure);
 }
 
-static LogicalResult verifyYieldTypes(Operation* op, Region& region, TypeRange resultTypes, StringRef kind) {
-  if (region.empty())
-    return op->emitOpError() << kind << " requires a body block";
-  Block& block = region.front();
-  auto yield = dyn_cast_or_null<SpatYieldOp>(block.getTerminator());
-  if (!yield)
-    return op->emitOpError() << kind << " body must terminate with spat.yield";
-  if (yield.getOutputs().size() != resultTypes.size())
-    return op->emitOpError() << kind << " yield operand count must match result count";
-  for (auto [yieldType, resultType] : llvm::zip(yield.getOutputs().getTypes(), resultTypes))
-    if (yieldType != resultType)
-      return op->emitOpError() << kind << " yield operand types must match result types";
-  return success();
-}
-
-static LogicalResult verifyRegionArguments(Operation* op, Region& region, ValueRange operands, StringRef kind) {
-  if (region.empty())
-    return op->emitOpError() << kind << " requires a body block";
-  Block& block = region.front();
-  if (block.getNumArguments() != operands.size())
-    return op->emitOpError() << kind << " body argument count must match operand count";
-  for (auto [arg, operand] : llvm::zip(block.getArguments(), operands))
-    if (arg.getType() != operand.getType())
-      return op->emitOpError() << kind << " body argument types must match operand types";
-  return success();
-}
-
 template <typename ComputeBatchOpTy>
 static LogicalResult verifyBatchBody(ComputeBatchOpTy batchOp, Block& block, bool verifyLaneSliceOffsets = true) {
   if (batchOp.getNumResults() == 0) {
@@ -741,11 +714,12 @@ LogicalResult verifyComputeLikeOp(ComputeOpTy compute, StringRef opName) {
   bool isScheduled = isa<SpatScheduledCompute>(compute.getOperation());
   if (compute.getBody().empty())
     return compute.emitOpError("compute body must have at least one block");
+  if (isScheduled && !compute.getBody().hasOneBlock())
+    return compute.emitOpError("scheduled compute must have exactly one block");
 
   SmallVector<Type> yieldedTypes;
-  for (Block& block : compute.getBody()) {
-    if ((!isScheduled && block.getNumArguments() != expectedArgCount)
-        || (isScheduled && block.getNumArguments() < expectedArgCount))
+  for (Block &block : compute.getBody()) {
+    if (block.getNumArguments() != expectedArgCount)
       return compute.emitOpError("compute body must have weight and input block arguments");
 
     for (auto [weightIndex, weight] : llvm::enumerate(compute.getWeights()))
@@ -757,17 +731,13 @@ LogicalResult verifyComputeLikeOp(ComputeOpTy compute, StringRef opName) {
 
     Operation* terminator = block.getTerminator();
     if (auto yieldOp = dyn_cast_or_null<SpatYieldOp>(terminator)) {
-      auto realized = compute->template getAttrOfType<BoolAttr>("scheduled.realized");
-      if (isScheduled && (!realized || !realized.getValue() || !compute.getBody().hasOneBlock()))
-        return compute.emitOpError("scheduled compute blocks must terminate with spat.block_yield");
       llvm::append_range(yieldedTypes, yieldOp->getOperandTypes());
       continue;
     }
     auto blockYield = dyn_cast_or_null<SpatBlockYieldOp>(terminator);
-    if (!blockYield || !isScheduled)
+    if (!blockYield || isScheduled)
       return compute.emitOpError("ComputeOp must have a single yield operation");
-    if (blockYield->getNumSuccessors() == 0)
-      llvm::append_range(yieldedTypes, blockYield->getOperandTypes());
+    llvm::append_range(yieldedTypes, blockYield->getOperandTypes());
   }
 
   auto resultTypes = compute.getResultTypes();
@@ -832,6 +802,11 @@ LogicalResult SpatBlockYieldOp::verify() {
 LogicalResult SpatDeferredCommunicationOp::verify() {
   if (getSources().empty())
     return emitOpError("requires at least one source");
+  auto specialization = (*this)->getAttrOfType<IntegerAttr>(
+    "specialization_count");
+  int64_t specializationCount = specialization ? specialization.getInt() : 1;
+  if (specializationCount <= 0)
+    return emitOpError("specialization_count must be positive");
   static constexpr StringLiteral staleAttributes[] = {
     "exchangeId", "logicalProducer", "logicalConsumer", "sourceClass", "targetClass", "sourceCore",
     "targetCore", "sourceLane", "targetLane", "transferKind", "resultIndex", "projectedTransfer",
@@ -842,9 +817,56 @@ LogicalResult SpatDeferredCommunicationOp::verify() {
     if (getOperation()->hasAttr(name))
       return emitOpError() << "does not accept stale routing attribute '" << name
                            << "'; source selection and shaping belong in the body and routing is derived in Phase 2";
-  if (failed(verifyRegionArguments(getOperation(), getBody(), getSources(), "spat.deferred_communication")))
-    return failure();
-  return verifyYieldTypes(getOperation(), getBody(), getOperation()->getResultTypes(), "spat.deferred_communication");
+  if (getBody().empty())
+    return emitOpError("spat.deferred_communication requires a body block");
+  Block &body = getBody().front();
+  unsigned expectedArguments = getSources().size()
+                             + (specializationCount > 1 ? 1 : 0);
+  if (body.getNumArguments() != expectedArguments)
+    return emitOpError("body argument count must match sources plus the grouped specialization argument");
+  for (auto [argument, source] : llvm::zip(
+         body.getArguments().take_front(getSources().size()), getSources()))
+    if (argument.getType() != source.getType())
+      return emitOpError("body source argument types must match source operand types");
+  if (specializationCount > 1
+      && !body.getArguments().back().getType().isIndex())
+    return emitOpError("grouped specialization argument must have index type");
+  auto yield = dyn_cast_or_null<SpatYieldOp>(body.getTerminator());
+  if (!yield || yield.getOutputs().size() != 1)
+    return emitOpError("body must yield exactly one fragment");
+  Type fragmentType = yield.getOutputs().front().getType();
+  Type outputType = getOutput().getType();
+  if (specializationCount == 1)
+    return fragmentType == outputType
+         ? success()
+         : emitOpError("ordinary deferred yield type must match its output type");
+  auto fragmentTensor = dyn_cast<RankedTensorType>(fragmentType);
+  auto outputTensor = dyn_cast<RankedTensorType>(outputType);
+  if (!fragmentTensor || !outputTensor || !fragmentTensor.hasStaticShape()
+      || !outputTensor.hasStaticShape())
+    return emitOpError("grouped specialization requires static ranked tensor types");
+  if (outputTensor.getRank() != fragmentTensor.getRank() + 1
+      || outputTensor.getDimSize(0) != specializationCount
+      || outputTensor.getShape().drop_front() != fragmentTensor.getShape()
+      || outputTensor.getElementType() != fragmentTensor.getElementType())
+    return emitOpError("grouped output must have shape specialization_count x fragment shape");
+  return success();
+}
+
+LogicalResult SpatDeferredSourceSelectOp::verify() {
+  if (getSources().empty())
+    return emitOpError("requires at least one source");
+  if (!getSelector().getType().isIndex())
+    return emitOpError("requires an index selector");
+  if (llvm::any_of(getSources(), [&](Value source) {
+        return source.getType() != getOutput().getType();
+      }))
+    return emitOpError("source and output types must match");
+  if (!getOperation()->getParentOfType<SpatDeferredCommunicationOp>()
+      && !getOperation()->getParentOfType<SpatScheduledCompute>()
+      && !getOperation()->getParentOfType<SpatScheduledComputeBatch>())
+    return emitOpError("must be nested in deferred or scheduled computation");
+  return success();
 }
 
 template <typename ComputeBatchOpTy>

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

@@ -1,709 +1,313 @@
-#include "mlir/Dialect/SCF/IR/SCF.h"
-
-#include "llvm/ADT/MapVector.h"
-
 #include "DeferredBoundaryPlanning.hpp"
 #include "DeferredCommunicationScheduling.hpp"
-#include "DeferredProjectionAnalysis.hpp"
 #include "DeferredTransferPlanning.hpp"
 
 namespace onnx_mlir::spatial {
 using namespace mlir;
 namespace {
 
-enum class BoundaryEventKind { Send, Receive };
+static BoundaryProgram &getBoundary(
+    SmallVectorImpl<BoundaryProgram> &boundaries,
+    DenseMap<BoundaryKey, unsigned> &indices, BoundaryKey key) {
+  auto [it, inserted] = indices.try_emplace(key, boundaries.size());
+  if (inserted)
+    boundaries.push_back({key, {}});
+  return boundaries[it->second];
+}
 
-struct BoundaryEvent {
-  BoundaryEventKind kind = BoundaryEventKind::Send;
-  ScheduledTransferSlice slice;
-  LaneSet activeLanes;
-  TransferEmissionSignature emission;
-};
+static LaneSet getReceiveLanes(const ScheduledTransferSlice &slice) {
+  LaneInterval family = slice.family->targetLanes.intervals().front();
+  unsigned begin = family.begin + slice.familyOffset;
+  return LaneSet::range(begin, begin + slice.transferCount);
+}
 
-struct SequenceNode {
-  unsigned previous = 0;
-  unsigned instruction = 0;
-};
+static void appendSend(BoundaryProgram &boundary,
+                       const ScheduledTransferSlice &slice) {
+  unsigned lane = slice.family->requirement->producer->scheduledLane;
+  LaneSet lanes = LaneSet::range(lane, lane + 1);
+  if (!boundary.instructions.empty())
+    if (auto *run = std::get_if<EmitSendRun>(&boundary.instructions.back());
+        run && haveSameTransferEmissionSignature(
+                 *run->slices.back().family, *slice.family)) {
+      run->slices.push_back(slice);
+      run->lanes = run->lanes.unite(lanes);
+      return;
+    }
+  boundary.instructions.push_back(EmitSendRun {{slice}, lanes});
+}
 
-struct IntervalClass {
-  LaneSet lanes;
-  unsigned sequence = 0;
-};
+static LogicalResult addCoverage(
+    RequirementFamily &requirement, const LaneSet &lanes,
+    DenseMap<RequirementFamily *, LaneSet> &coverage) {
+  LaneSet &covered = coverage[&requirement];
+  if (!covered.intersect(lanes).empty())
+    return requirement.exchange->deferred.emitOpError(
+      "deferred availability covers a target lane more than once");
+  covered = covered.unite(lanes);
+  return success();
+}
 
-enum class SemanticKind {
-  Send,
-  Availability,
-  Result
-};
+static bool canGroupLocalAvailability(RequirementFamily &lhs,
+                                      RequirementFamily &rhs) {
+  if (!(lhs.coordinate == rhs.coordinate)
+      || lhs.publicationFragmentType != rhs.publicationFragmentType
+      || lhs.producer->payload != rhs.producer->payload
+      || lhs.producerProjection.has_value()
+           != rhs.producerProjection.has_value())
+    return false;
+  if (!lhs.producerProjection)
+    return true;
+  auto ranks = [](const DeferredStaticSliceGeometry &geometry) {
+    return std::tuple(geometry.offsets.size(), geometry.sizes.size(),
+                      geometry.strides.size());
+  };
+  return ranks(*lhs.producerProjection) == ranks(*rhs.producerProjection);
+}
 
-struct SemanticKey {
-  SemanticKind kind = SemanticKind::Send;
-  TransferEmissionSignature emission;
-  DeferredExchangePlan* exchange = nullptr;
-  RequirementCoordinate coordinate;
-  Type fragmentType;
-
-  bool operator==(const SemanticKey& other) const {
-    if (kind != other.kind)
-      return false;
-    if (kind == SemanticKind::Send)
-      return emission == other.emission;
-    if (kind == SemanticKind::Result)
-      return exchange == other.exchange;
-    return exchange == other.exchange && coordinate == other.coordinate
-        && fragmentType == other.fragmentType;
-  }
-};
-
-struct PendingToken {
-  LaneSet lanes;
-  unsigned semantic = 0;
-  std::variant<unsigned, LocalAvailabilityFamily*, DeferredExchangePlan*> value;
-};
-
-struct SequenceCursor {
-  LaneSet lanes;
+struct CollectionTarget {
+  const FragmentCollectionPlan *collection = nullptr;
   unsigned position = 0;
 };
 
-struct CanonicalAction {
-  SemanticKey key;
-  SmallVector<ScheduledTransferSlice> slices;
-  SmallVector<LocalAvailabilityFamily*> locals;
-  LaneSet instructionLanes;
-  LaneSet receiveLanes;
-  LaneSet localLanes;
-  DeferredExchangePlan* result = nullptr;
-};
-
-struct BoundaryWork {
-  BoundaryProgram program;
-  SmallVector<BoundaryEvent> events;
-};
-
-static unsigned getBoundaryIndex(SmallVectorImpl<BoundaryWork>& boundaries,
-                                 DenseMap<BoundaryKey, unsigned>& indices,
-                                 BoundaryKey key) {
-  if (auto it = indices.find(key); it != indices.end())
-    return it->second;
-  unsigned index = boundaries.size();
-  indices[key] = index;
-  BoundaryWork work;
-  work.program.key = key;
-  boundaries.push_back(std::move(work));
-  return index;
-}
-
-static size_t hashSemanticKey(const SemanticKey& key) {
-  if (key.kind == SemanticKind::Send)
-    return llvm::hash_combine(key.kind,
-                              key.emission.scheduled,
-                              key.emission.payload.getAsOpaquePointer(),
-                              key.emission.fragmentType.getAsOpaquePointer(),
-                              key.emission.hasGraphLane,
-                              key.emission.hasProducerProjection,
-                              key.emission.sourceIsBatch);
-  if (key.kind == SemanticKind::Result)
-    return llvm::hash_combine(key.kind, key.exchange);
-  return llvm::hash_combine(key.kind, key.exchange,
-                            key.coordinate.leafIndex,
-                            key.coordinate.selectedPosition,
-                            key.fragmentType.getAsOpaquePointer());
-}
-
-static unsigned internSemantic(const SemanticKey& key,
-                               SmallVectorImpl<SemanticKey>& keys,
-                               DenseMap<size_t, SmallVector<unsigned>>& byHash) {
-  size_t hash = hashSemanticKey(key);
-  for (unsigned candidate : byHash.lookup(hash))
-    if (keys[candidate] == key)
-      return candidate;
-  unsigned id = keys.size();
-  keys.push_back(key);
-  byHash[hash].push_back(id);
-  return id;
-}
-
-static SemanticKey getEventKey(const BoundaryEvent& event) {
-  if (event.kind == BoundaryEventKind::Send) {
-    SemanticKey key;
-    key.kind = SemanticKind::Send;
-    key.emission = event.emission;
-    return key;
-  }
-  RequirementFamily& requirement = *event.slice.family->requirement;
-  SemanticKey key;
-  key.kind = SemanticKind::Availability;
-  key.exchange = requirement.exchange;
-  key.coordinate = requirement.coordinate;
-  key.fragmentType = requirement.publicationFragmentType;
-  return key;
-}
-
-static SemanticKey getLocalKey(LocalAvailabilityFamily& local) {
-  RequirementFamily& requirement = *local.requirement;
-  SemanticKey key;
-  key.kind = SemanticKind::Availability;
-  key.exchange = requirement.exchange;
-  key.coordinate = requirement.coordinate;
-  key.fragmentType = requirement.publicationFragmentType;
-  return key;
-}
-
-static SmallVector<ScheduledTransferSlice> intersectReceiveSlice(const ScheduledTransferSlice& slice,
-                                                                 const LaneSet& lanes) {
-  LaneInterval family = slice.family->targetLanes.intervals().front();
-  unsigned sliceBegin = family.begin + slice.familyOffset;
-  LaneSet active = LaneSet::range(sliceBegin, sliceBegin + slice.transferCount).intersect(lanes);
-  SmallVector<ScheduledTransferSlice> result;
-  for (LaneInterval selected : active.intervals()) {
-    ScheduledTransferSlice part = slice;
-    part.familyOffset += selected.begin - sliceBegin;
-    part.transferCount = selected.end - selected.begin;
-    result.push_back(part);
-  }
-  return result;
-}
-
-static void mergeSequenceCursors(SmallVectorImpl<SequenceCursor>& states) {
-  DenseMap<unsigned, unsigned> byPosition;
-  SmallVector<SequenceCursor> merged;
-  for (SequenceCursor& state : states) {
-    auto [it, inserted] = byPosition.try_emplace(state.position, merged.size());
-    if (inserted)
-      merged.push_back(std::move(state));
-    else
-      merged[it->second].lanes = merged[it->second].lanes.unite(state.lanes);
-  }
-  states = std::move(merged);
-}
-
-static LogicalResult appendReplayToken(const PendingToken& token,
-                                       const LaneSet& lanes,
-                                       CanonicalAction& action,
-                                       ArrayRef<BoundaryEvent> events) {
-  if (auto eventId = std::get_if<unsigned>(&token.value)) {
-    const BoundaryEvent& event = events[*eventId];
-    if (event.kind == BoundaryEventKind::Send) {
-      LaneSet active = event.activeLanes.intersect(lanes);
-      if (!active.empty())
-        action.slices.push_back(event.slice);
-      action.instructionLanes =
-        action.instructionLanes.unite(active);
-    }
-    else {
-      llvm::append_range(action.slices, intersectReceiveSlice(event.slice, lanes));
-      action.receiveLanes = action.receiveLanes.unite(lanes);
-    }
-    return success();
-  }
-  if (auto local = std::get_if<LocalAvailabilityFamily*>(&token.value)) {
-    if (!llvm::is_contained(action.locals, *local))
-      action.locals.push_back(*local);
-    action.localLanes = action.localLanes.unite(lanes);
-    return success();
-  }
-  action.result = *std::get_if<DeferredExchangePlan*>(&token.value);
-  return success();
-}
-
-static FailureOr<SmallVector<CanonicalAction, 0>> collectCanonicalActions(
-  ArrayRef<unsigned> sequence, ArrayRef<PendingToken> tokens,
-  ArrayRef<SemanticKey> semantics, ArrayRef<BoundaryEvent> events,
-  const LaneSet& lanes) {
-  SmallVector<CanonicalAction, 0> actions;
-  for (unsigned semantic : sequence) {
-    CanonicalAction action;
-    action.key = semantics[semantic];
-    actions.push_back(std::move(action));
-  }
-  SmallVector<SequenceCursor> states {
-    {lanes, 0}
-  };
-  for (const PendingToken& token : tokens) {
-    SmallVector<SequenceCursor> next;
-    for (const SequenceCursor& state : states) {
-      LaneSet intersection = state.lanes.intersect(token.lanes);
-      LaneSet difference = state.lanes.subtract(token.lanes);
-      if (!difference.empty())
-        next.push_back({difference, state.position});
-      if (intersection.empty())
-        continue;
-      if (state.position == sequence.size()) {
-        next.push_back({intersection, state.position});
-        continue;
-      }
-      if (state.position >= sequence.size() || sequence[state.position] != token.semantic
-          || failed(appendReplayToken(token, intersection, actions[state.position], events)))
-        return failure();
-      next.push_back({intersection, state.position + 1});
-    }
-    mergeSequenceCursors(next);
-    states = std::move(next);
-  }
-  if (llvm::any_of(states, [&](const SequenceCursor& state) { return state.position != sequence.size(); }))
-    return failure();
-  return actions;
-}
-
-static bool matchesAssembly(ArrayRef<CanonicalAction> actions, size_t begin, SmallVectorImpl<unsigned>& entryOrder) {
-  if (begin >= actions.size())
+static bool sameCollectionEmissionContract(
+    const CollectionTarget &lhs, const CollectionTarget &rhs) {
+  if (lhs.collection != rhs.collection)
     return false;
-  const CanonicalAction& first = actions[begin];
-  DeferredExchangePlan* exchange = first.key.exchange;
-  auto& assembly = exchange->program.insertAssembly;
-  if (first.key.kind != SemanticKind::Availability || !first.locals.empty()
-      || !assembly || assembly->entries.empty()
-      || begin + assembly->entries.size() > actions.size())
-    return false;
-  SmallVector<bool> matched(assembly->entries.size());
-  for (size_t offset = 0; offset < assembly->entries.size(); ++offset) {
-    const CanonicalAction& action = actions[begin + offset];
-    if (action.key.kind != SemanticKind::Availability || !action.locals.empty()
-        || action.key.exchange != exchange || action.slices.empty())
-      return false;
-    std::optional<unsigned> entry;
-    for (auto [entryIndex, candidate] : llvm::enumerate(assembly->entries))
-      if (!matched[entryIndex] && action.key.coordinate == candidate.coordinate) {
-        entry = entryIndex;
-        break;
-      }
-    if (!entry)
-      return false;
-    matched[*entry] = true;
-    entryOrder.push_back(*entry);
-  }
-  if (assembly->entries.size() > 1
-      && llvm::any_of(ArrayRef(assembly->entries).drop_front(), [&](const DeferredInsertAssemblyEntryTemplate& entry) {
-           return entry.sourceTransform != assembly->entries.front().sourceTransform
-               || entry.sourceType != assembly->entries.front().sourceType;
-         }))
-    return false;
-  return true;
+  if (lhs.collection->key.kind != FragmentCollectionKind::InsertAssembly)
+    return true;
+  const auto &entries =
+    lhs.collection->key.exchange->program.insertAssembly->entries;
+  const auto &left = entries[lhs.position];
+  const auto &right = entries[rhs.position];
+  return left.sourceTransform == right.sourceTransform
+      && left.sourceType == right.sourceType;
 }
 
-static bool matchesProjectionAssembly(ArrayRef<CanonicalAction> actions,
-                                      size_t begin,
-                                      const LaneSet& lanes,
-                                      unsigned& leafIndex,
-                                      SmallVectorImpl<unsigned>& positions) {
-  if (begin >= actions.size() || lanes.empty())
-    return false;
-  const CanonicalAction& first = actions[begin];
-  DeferredExchangePlan* exchange = first.key.exchange;
-  if (first.key.kind != SemanticKind::Availability || !first.locals.empty()
-      || !exchange || exchange->program.insertAssembly
-      || first.key.coordinate.leafIndex >= exchange->program.leaves.size())
-    return false;
-  leafIndex = first.key.coordinate.leafIndex;
-  const DeferredProjectionLeafTemplate& leaf = exchange->program.leaves[leafIndex];
-  if (leaf.form != DeferredLeafForm::DirectSource)
-    return false;
-  unsigned representative = lanes.intervals().front().begin;
-  unsigned positionCount = 0;
-  unsigned requirementCount = 0;
-  Type fragmentType;
-  for (RequirementFamily& requirement : exchange->requirements) {
-    if (requirement.coordinate.leafIndex != leafIndex || !requirement.targetLanes.contains(representative))
-      continue;
-    ++requirementCount;
-    positionCount = std::max(positionCount, requirement.coordinate.selectedPosition + 1);
-    if (fragmentType && fragmentType != requirement.publicationFragmentType)
-      return false;
-    fragmentType = requirement.publicationFragmentType;
+static std::optional<EmitLocalCollectionRun> buildLocalConcat(
+    const FragmentCollectionPlan &collection,
+    const DenseMap<RequirementFamily *, LocalAvailabilityFamily *> &locals,
+    unsigned targetLaneCount) {
+  RankedTensorType type = collection.collectionType;
+  if (collection.key.kind != FragmentCollectionKind::Leaf
+      || targetLaneCount != 1 || type.getRank() == 0
+      || collection.positionCount == 0
+      || type.getDimSize(0) != collection.positionCount)
+    return std::nullopt;
+  SmallVector<RequirementFamily *> requirements(collection.positionCount);
+  for (const auto &entry : collection.requirements) {
+    if (entry.position >= requirements.size() || requirements[entry.position])
+      return std::nullopt;
+    requirements[entry.position] = entry.family;
   }
-  auto fragment = dyn_cast<RankedTensorType>(fragmentType);
-  if (positionCount < 2 || requirementCount != positionCount || !fragment
-      || leaf.reconstructedType.getRank() != fragment.getRank() + 1
-      || leaf.reconstructedType.getDimSize(0) != positionCount
-      || leaf.reconstructedType.getShape().drop_front() != fragment.getShape())
-    return false;
-  SmallVector<bool> seen(positionCount);
-  for (size_t offset = 0; begin + offset < actions.size(); ++offset) {
-    const CanonicalAction& action = actions[begin + offset];
-    if (action.key.kind != SemanticKind::Availability || !action.locals.empty()
-        || action.key.exchange != exchange
-        || action.key.coordinate.leafIndex != leafIndex || action.key.fragmentType != fragmentType)
-      break;
-    unsigned position = action.key.coordinate.selectedPosition;
-    if (action.slices.empty() || position >= positionCount || seen[position])
-      return false;
-    seen[position] = true;
-    positions.push_back(position);
-  }
-  if (positions.size() < 2)
-    return false;
-  for (RequirementFamily& requirement : exchange->requirements) {
-    if (requirement.coordinate.leafIndex != leafIndex || !requirement.targetLanes.contains(representative)
-        || seen[requirement.coordinate.selectedPosition])
-      continue;
-    bool local = llvm::any_of(exchange->local, [&](const LocalAvailabilityFamily& availability) {
-      return availability.requirement == &requirement && availability.targetLanes.contains(representative);
-    });
-    if (!local)
-      return false;
-  }
-  return true;
-}
-
-static size_t getReceiveBundleLength(ArrayRef<CanonicalAction> actions,
-                                     size_t begin,
-                                     const LaneSet& lanes) {
-  if (begin >= actions.size())
-    return 0;
-  const CanonicalAction& first = actions[begin];
-  if (first.key.kind != SemanticKind::Availability)
-    return 0;
-  Type fragmentType = first.key.fragmentType;
-  auto rankedFragment = dyn_cast<RankedTensorType>(fragmentType);
-  if (!rankedFragment || !rankedFragment.hasStaticShape())
-    return 0;
-  Value firstOutput = first.key.exchange->deferred.getOutput();
-  if (!firstOutput.hasOneUse())
-    return 0;
-  Operation* firstUser = *firstOutput.getUsers().begin();
-  auto selection = firstUser->getParentOfType<scf::IndexSwitchOp>();
-  if (!selection)
-    return 0;
-  size_t end = begin;
-  while (end < actions.size()) {
-    const CanonicalAction& action = actions[end];
-    Value output = action.key.exchange
-                     ? action.key.exchange->deferred.getOutput()
-                     : Value();
-    Operation* user = output && output.hasOneUse()
-                        ? *output.getUsers().begin()
-                        : nullptr;
-    if (action.key.kind != SemanticKind::Availability || !action.locals.empty()
-        || action.slices.empty() || !(action.receiveLanes == lanes)
-        || action.key.fragmentType != fragmentType
-        || !user || user->getParentOfType<scf::IndexSwitchOp>() != selection)
-      break;
-    ++end;
-  }
-  return end - begin >= 2 ? end - begin : 0;
-}
-
-static BoundaryInstructionList materializeInstructions(ArrayRef<CanonicalAction> actions,
-                                                  const LaneSet& lanes) {
-  BoundaryInstructionList result;
-  auto& instructions = result.instructions;
-  for (size_t index = 0; index < actions.size();) {
-    const CanonicalAction& action = actions[index];
-    if (action.key.kind == SemanticKind::Send) {
-      EmitSendRun run;
-      run.lanes = action.instructionLanes;
-      do {
-        llvm::append_range(run.slices, actions[index].slices);
-        run.lanes = run.lanes.unite(actions[index].instructionLanes);
-        ++index;
-      }
-      while (index < actions.size() && actions[index].key.kind == SemanticKind::Send
-             && actions[index].key.emission == action.key.emission);
-      if (!run.slices.empty())
-        instructions.push_back(std::move(run));
-      continue;
+  LaneSet all = LaneSet::all(targetLaneCount);
+  EmitLocalCollectionRun run {&collection, 0, {}, all, true};
+  Value payload;
+  int64_t payloadBegin = 0;
+  int64_t payloadEnd = 0;
+  for (auto [position, requirement] : llvm::enumerate(requirements)) {
+    if (!requirement)
+      return std::nullopt;
+    LocalAvailabilityFamily *local = locals.lookup(requirement);
+    if (!local || !(requirement->targetLanes == all)
+        || !(local->targetLanes == all) || !requirement->graphLanes
+        || requirement->graphLanes->size() != 1
+        || requirement->graphLanes->valueAt(0) != static_cast<int64_t>(position)
+        || !requirement->producerLocalOffsets
+        || requirement->producerLocalOffsets->size() != 1)
+      return std::nullopt;
+    ProducedValue *producer = requirement->producer;
+    int64_t payloadOffset =
+      requirement->producerLocalOffsets->valueAt(0);
+    if (static_cast<int64_t>(position) == payloadEnd) {
+      if (!producer)
+        return std::nullopt;
+      auto payloadType = dyn_cast<RankedTensorType>(producer->payload.getType());
+      if (!payloadType || payloadOffset != 0
+          || payloadType.getRank() != type.getRank()
+          || payloadType.getElementType() != type.getElementType()
+          || payloadType.getShape().drop_front() != type.getShape().drop_front())
+        return std::nullopt;
+      payloadBegin = position;
+      payloadEnd = payloadBegin + payloadType.getDimSize(0);
+      if (payloadEnd > collection.positionCount)
+        return std::nullopt;
+      payload = producer->payload;
+      run.families.push_back(local);
     }
-    SmallVector<unsigned> assemblyEntries;
-    if (matchesAssembly(actions, index, assemblyEntries)) {
-      EmitReceiveAssemblyRun run;
-      run.lanes = lanes;
-      run.assemblyEntries = assemblyEntries;
-      run.entryOffsets.push_back(0);
-      for (size_t offset = 0; offset < assemblyEntries.size(); ++offset) {
-        llvm::append_range(run.slices, actions[index + offset].slices);
-        run.entryOffsets.push_back(run.slices.size());
-      }
-      instructions.push_back(std::move(run));
-      index += assemblyEntries.size();
-      continue;
-    }
-    unsigned projectionLeaf = 0;
-    SmallVector<unsigned> projectionPositions;
-    if (matchesProjectionAssembly(actions, index, lanes, projectionLeaf, projectionPositions)) {
-      EmitReceiveAssemblyRun run;
-      run.lanes = lanes;
-      run.projectionLeaf = projectionLeaf;
-      run.assemblyEntries = projectionPositions;
-      run.entryOffsets.push_back(0);
-      for (size_t offset = 0; offset < projectionPositions.size(); ++offset) {
-        llvm::append_range(run.slices, actions[index + offset].slices);
-        run.entryOffsets.push_back(run.slices.size());
-      }
-      instructions.push_back(std::move(run));
-      index += projectionPositions.size();
-      continue;
-    }
-    if (size_t bundleLength = getReceiveBundleLength(actions, index, lanes)) {
-      EmitReceiveBundle bundle;
-      for (size_t offset = 0; offset < bundleLength; ++offset) {
-        const CanonicalAction& entry = actions[index + offset];
-        EmitReceiveRun receive;
-        receive.slices = entry.slices;
-        receive.entryOffsets = {0, receive.slices.size()};
-        receive.lanes = entry.receiveLanes;
-        bundle.entries.push_back(std::move(receive));
-      }
-      instructions.push_back(std::move(bundle));
-      index += bundleLength;
-      continue;
-    }
-    if (action.key.kind == SemanticKind::Availability) {
-      ResolveAvailability availability;
-      availability.exchange = action.key.exchange;
-      availability.coordinate = action.key.coordinate;
-      availability.fragmentType = action.key.fragmentType;
-      if (!action.slices.empty()) {
-        EmitReceiveRun receive;
-        receive.slices = action.slices;
-        receive.entryOffsets = {0, receive.slices.size()};
-        receive.lanes = action.receiveLanes;
-        availability.alternatives.push_back(
-          {receive.lanes, AvailabilitySource(std::move(receive))});
-      }
-      if (!action.locals.empty()) {
-        MaterializeLocalFamily local {action.locals, action.localLanes};
-        availability.alternatives.push_back(
-          {local.lanes, AvailabilitySource(std::move(local))});
-      }
-      instructions.push_back(std::move(availability));
-      ++index;
-      continue;
-    }
-    instructions.push_back(ProduceDeferredResult {action.result, lanes});
-    ++index;
+    if (producer->payload != payload
+        || payloadOffset != static_cast<int64_t>(position) - payloadBegin)
+      return std::nullopt;
   }
-  return result;
+  return payloadEnd == collection.positionCount
+       ? std::optional<EmitLocalCollectionRun>(std::move(run)) : std::nullopt;
 }
 
-static void addTokenToClasses(const LaneSet& active,
-                              unsigned semantic,
-                              SmallVectorImpl<IntervalClass>& classes,
-                              SmallVectorImpl<SequenceNode>& nodes,
-                              DenseMap<std::pair<unsigned, unsigned>, unsigned>& interned) {
-  SmallVector<IntervalClass> next;
-  for (const IntervalClass& current : classes) {
-    LaneSet intersection = current.lanes.intersect(active);
-    LaneSet difference = current.lanes.subtract(active);
-    if (!difference.empty())
-      next.push_back({difference, current.sequence});
-    if (intersection.empty())
-      continue;
-    auto key = std::make_pair(current.sequence, semantic);
-    auto [it, inserted] = interned.try_emplace(key, nodes.size());
-    if (inserted)
-      nodes.push_back({current.sequence, semantic});
-    next.push_back({intersection, it->second});
-  }
-  classes = std::move(next);
-}
-
-struct SequenceClass {
-  LaneSet lanes;
-  SmallVector<unsigned> sequence;
-};
-
-static SmallVector<SequenceClass>
-buildSequenceClasses(unsigned laneCount, ArrayRef<PendingToken> tokens) {
-  SmallVector<IntervalClass> classes {
-    {LaneSet::all(laneCount), 0}
-  };
-  SmallVector<SequenceNode> nodes {
-    {0, 0}
-  };
-  DenseMap<std::pair<unsigned, unsigned>, unsigned> interned;
-  for (const PendingToken& token : tokens)
-    addTokenToClasses(token.lanes, token.semantic, classes, nodes, interned);
-
-  SmallVector<SequenceClass> result;
-  DenseMap<unsigned, unsigned> classBySequence;
-  SmallVector<unsigned> classSequences;
-  for (const IntervalClass& interval : classes) {
-    auto [it, inserted] = classBySequence.try_emplace(
-      interval.sequence, result.size());
-    if (inserted) {
-      result.push_back({interval.lanes, {}});
-      classSequences.push_back(interval.sequence);
+static void appendReceive(BoundaryProgram &boundary,
+                          const ScheduledTransferSlice &slice,
+                          CollectionTarget target) {
+  RequirementFamily *requirement = slice.family->requirement;
+  LaneSet lanes = getReceiveLanes(slice);
+  if (!boundary.instructions.empty())
+    if (auto *run = std::get_if<EmitReceiveAssemblyRun>(
+          &boundary.instructions.back())) {
+      RequirementFamily *previous = run->slices[
+        run->entryOffsets[run->entryOffsets.size() - 2]].family->requirement;
+      CollectionTarget previousTarget {run->collection, run->positions.back()};
+      bool sameEntry = previous == requirement;
+      if (sameEntry
+          || (sameCollectionEmissionContract(previousTarget, target)
+              && previous->publicationFragmentType
+                   == requirement->publicationFragmentType)) {
+        run->slices.push_back(slice);
+        if (sameEntry) {
+          run->entryOffsets.back() = run->slices.size();
+          run->entryLanes.back() = run->entryLanes.back().unite(lanes);
+        } else {
+          run->entryOffsets.push_back(run->slices.size());
+          run->positions.push_back(target.position);
+          run->entryLanes.push_back(lanes);
+        }
+        run->lanes = run->lanes.unite(lanes);
+        return;
+      }
     }
-    else
-      result[it->second].lanes = result[it->second].lanes.unite(interval.lanes);
-  }
-  for (auto [behavior, sequence] : llvm::zip_equal(result, classSequences)) {
-    for (unsigned node = sequence; node != 0; node = nodes[node].previous)
-      behavior.sequence.push_back(nodes[node].instruction);
-    std::reverse(behavior.sequence.begin(), behavior.sequence.end());
-  }
-  return result;
-}
-
-static SmallVector<DeferredExchangePlan*>
-getProducedExchanges(const BoundaryInstructionList& list) {
-  SmallVector<DeferredExchangePlan*> result;
-  for (const BoundaryInstruction& instruction : list.instructions)
-    if (auto produced = std::get_if<ProduceDeferredResult>(&instruction))
-      result.push_back(produced->exchange);
-  return result;
-}
-
-static LogicalResult buildCanonicalBoundary(
-  BoundaryProgram& boundary, ArrayRef<BoundaryEvent> events,
-  ArrayRef<PendingToken> tokens, ArrayRef<SemanticKey> semantics) {
-  unsigned laneCount = boundary.key.scheduled->cores.size();
-  SmallVector<SequenceClass> classes =
-    buildSequenceClasses(laneCount, tokens);
-  if (classes.empty())
-    return failure();
-  size_t prefix = classes.front().sequence.size();
-  for (const SequenceClass& behavior : ArrayRef(classes).drop_front()) {
-    prefix = std::min(prefix, behavior.sequence.size());
-    size_t index = 0;
-    while (index < prefix
-           && behavior.sequence[index] == classes.front().sequence[index])
-      ++index;
-    prefix = index;
-  }
-  auto prefixActions = collectCanonicalActions(
-    ArrayRef(classes.front().sequence).take_front(prefix), tokens, semantics,
-    events, LaneSet::all(laneCount));
-  if (failed(prefixActions))
-    return failure();
-  boundary.root = materializeInstructions(*prefixActions, LaneSet::all(laneCount));
-  if (classes.size() == 1)
-    return success();
-
-  auto dispatch = std::make_unique<LaneDispatch>();
-  SmallVector<int64_t> classIds(laneCount);
-  for (auto [classId, behavior] : llvm::enumerate(classes)) {
-    for (LaneInterval interval : behavior.lanes.intervals())
-      for (unsigned lane = interval.begin; lane < interval.end; ++lane)
-        classIds[lane] = classId;
-    auto actions = collectCanonicalActions(
-      behavior.sequence, tokens, semantics, events, behavior.lanes);
-    if (failed(actions))
-      return failure();
-    dispatch->branches.push_back(materializeInstructions(
-      ArrayRef(*actions).drop_front(prefix), behavior.lanes));
-    auto produced = getProducedExchanges(dispatch->branches.back());
-    if (classId == 0)
-      dispatch->producedExchanges = std::move(produced);
-    else if (produced != dispatch->producedExchanges)
-      return failure();
-  }
-  dispatch->branchByLane = StaticIntSequence::fromValues(classIds);
-  boundary.root.instructions.push_back(std::move(dispatch));
-  return success();
+  boundary.instructions.push_back(EmitReceiveAssemblyRun {
+    target.collection, {slice}, {0, 1}, {target.position}, {lanes}, lanes});
 }
 
 } // namespace
 
-FailureOr<DeferredBoundaryPlan> buildDeferredBoundaryPlan(DeferredTransferPlan& transfers,
-                                                          const ScheduledCommunicationPlan& schedule) {
+FailureOr<DeferredBoundaryPlan> buildDeferredBoundaryPlan(
+    DeferredTransferPlan &transfers,
+    const ScheduledCommunicationPlan &schedule) {
   DeferredBoundaryPlan result;
-  SmallVector<BoundaryWork> boundaries;
-  DenseMap<BoundaryKey, unsigned> boundaryIndices;
-  DenseMap<DeferredExchangePlan*, unsigned> resultBoundarySteps;
-  for (const ScheduledTransferSlice& slice : schedule.slices) {
-    ExternalTransferFamily& family = *slice.family;
-    auto [resultStep, inserted] =
-      resultBoundarySteps.try_emplace(family.requirement->exchange, slice.targetInsertionStep);
-    if (!inserted)
-      resultStep->second = std::max(resultStep->second, slice.targetInsertionStep);
-    unsigned sourceIndex =
-      getBoundaryIndex(boundaries, boundaryIndices,
-                       {family.sourceScheduled, slice.sourceInsertionStep});
-    unsigned targetIndex =
-      getBoundaryIndex(boundaries, boundaryIndices,
-                       {family.targetScheduled, slice.targetInsertionStep});
-    unsigned sourceLane = family.requirement->producer->scheduledLane;
-    boundaries[sourceIndex].events.push_back({BoundaryEventKind::Send,
-                                              slice,
-                                              LaneSet::range(sourceLane, sourceLane + 1),
-                                              getTransferEmissionSignature(family)});
-    LaneInterval familyLanes = family.targetLanes.intervals().front();
-    unsigned targetBegin = familyLanes.begin + slice.familyOffset;
-    boundaries[targetIndex].events.push_back({BoundaryEventKind::Receive,
-                                              slice,
-                                              LaneSet::range(targetBegin, targetBegin + slice.transferCount),
-                                              getTransferEmissionSignature(family)});
+  SmallVector<BoundaryProgram> boundaries;
+  DenseMap<BoundaryKey, unsigned> indices;
+  DenseMap<DeferredExchangePlan *, unsigned> resultSteps;
+  DenseMap<RequirementFamily *, LaneSet> coverage;
+
+  for (const std::unique_ptr<DeferredExchangePlan> &exchange :
+       transfers.exchanges) {
+    auto plan = buildDeferredResultPlan(*exchange);
+    if (failed(plan))
+      return exchange->deferred.emitOpError(
+        "cannot evaluate deferred result lane functions"), failure();
+    result.results.push_back(std::move(*plan));
+  }
+  DenseMap<RequirementFamily *, CollectionTarget> collections;
+  for (const DeferredResultPlan &plan : result.results)
+    for (const FragmentCollectionPlan &collection : plan.collections)
+      for (const FragmentCollectionPlan::Requirement &requirement :
+           collection.requirements)
+        if (!collections.try_emplace(
+              requirement.family,
+              CollectionTarget{&collection, requirement.position}).second)
+          return requirement.family->exchange->deferred.emitOpError(
+            "deferred requirement is owned by multiple fragment collections"),
+            failure();
+
+  for (const ScheduledTransferSlice &slice : schedule.slices) {
+    ExternalTransferFamily &family = *slice.family;
+    BoundaryProgram &source = getBoundary(
+      boundaries, indices,
+      {family.sourceScheduled, slice.sourceInsertionStep});
+    appendSend(source, slice);
+    BoundaryProgram &target = getBoundary(
+      boundaries, indices,
+      {family.targetScheduled, slice.targetInsertionStep});
+    CollectionTarget collection = collections.lookup(family.requirement);
+    if (!collection.collection)
+      return family.requirement->exchange->deferred.emitOpError(
+        "deferred requirement has no complete result-owned collection"),
+        failure();
+    appendReceive(target, slice, collection);
+    resultSteps[family.requirement->exchange] = std::max(
+      resultSteps.lookup(family.requirement->exchange),
+      slice.targetInsertionStep);
+    if (failed(addCoverage(*family.requirement, getReceiveLanes(slice),
+                           coverage)))
+      return failure();
   }
 
-  DenseMap<BoundaryKey, SmallVector<LocalAvailabilityFamily*>> locals;
-  DenseMap<BoundaryKey, SmallVector<DeferredExchangePlan*>> exchanges;
-  for (const std::unique_ptr<DeferredExchangePlan>& exchange : transfers.exchanges) {
-    auto resultStep = resultBoundarySteps.find(exchange.get());
-    unsigned step = resultStep == resultBoundarySteps.end() ? exchange->consumerStep : resultStep->second;
-    BoundaryKey key {exchange->target, step};
-    getBoundaryIndex(boundaries, boundaryIndices, key);
-    for (LocalAvailabilityFamily& local : exchange->local)
-      locals[key].push_back(&local);
-    exchanges[key].push_back(exchange.get());
-    auto resultPlan = buildDeferredResultPlan(*exchange);
-    if (failed(resultPlan))
-      return exchange->deferred.emitOpError("cannot evaluate deferred result lane functions"), failure();
-    result.results.push_back(std::move(*resultPlan));
-  }
+  for (const std::unique_ptr<DeferredExchangePlan> &exchange :
+       transfers.exchanges) {
+    unsigned resultStep = resultSteps.lookup(exchange.get());
+    for (LocalAvailabilityFamily &local : exchange->local)
+      resultStep = std::max(
+        resultStep, local.requirement->producer->step + 1);
+    if (exchange->requirements.empty())
+      resultStep = exchange->consumerStep;
+    if (resultStep > exchange->consumerStep)
+      return exchange->deferred.emitOpError(
+        "deferred result boundary is later than its consumer"), failure();
 
-  DenseMap<ScheduledInfo*, unsigned> scheduledOrder;
-  for (auto [index, scheduled] : llvm::enumerate(transfers.scheduled))
-    scheduledOrder[&scheduled] = index;
-  llvm::stable_sort(boundaries, [&](const BoundaryWork& lhs, const BoundaryWork& rhs) {
-    return std::tie(scheduledOrder[lhs.program.key.scheduled], lhs.program.key.insertionStep)
-         < std::tie(scheduledOrder[rhs.program.key.scheduled], rhs.program.key.insertionStep);
-  });
-  for (BoundaryWork& work : boundaries) {
-    BoundaryProgram& boundary = work.program;
-    SmallVector<PendingToken> tokens;
-    SmallVector<SemanticKey> semantics;
-    DenseMap<size_t, SmallVector<unsigned>> semanticsByHash;
-    DenseMap<unsigned, SmallVector<LocalAvailabilityFamily*>> localsBySemantic;
-    SmallPtrSet<LocalAvailabilityFamily*, 8> emittedLocals;
-    for (LocalAvailabilityFamily* local : locals[boundary.key]) {
-      unsigned semantic =
-        internSemantic(getLocalKey(*local), semantics, semanticsByHash);
-      localsBySemantic[semantic].push_back(local);
+    BoundaryProgram &boundary = getBoundary(
+      boundaries, indices, {exchange->target, resultStep});
+    DenseMap<RequirementFamily *, LocalAvailabilityFamily *> localByRequirement;
+    for (LocalAvailabilityFamily &local : exchange->local) {
+      auto [it, inserted] = localByRequirement.try_emplace(local.requirement, &local);
+      if (!inserted)
+        it->second = nullptr;
     }
-    SmallVector<unsigned> eventSemantics;
-    for (BoundaryEvent& event : work.events) {
-      unsigned semantic =
-        internSemantic(getEventKey(event), semantics, semanticsByHash);
-      eventSemantics.push_back(semantic);
-    }
-    llvm::SmallDenseSet<unsigned, 8> emittedAvailabilities;
-    for (auto [eventId, event] : llvm::enumerate(work.events)) {
-      unsigned semantic = eventSemantics[eventId];
-      if (event.kind == BoundaryEventKind::Send) {
-        tokens.push_back({LaneSet::all(
-                            boundary.key.scheduled->cores.size()),
-                          semantic, static_cast<unsigned>(eventId)});
+    DenseMap<const FragmentCollectionPlan *, std::optional<EmitLocalCollectionRun>> concatRuns;
+    llvm::SmallPtrSet<const FragmentCollectionPlan *, 4> emittedConcats;
+    SmallVector<EmitLocalCollectionRun> localUpdates;
+    for (LocalAvailabilityFamily &local : exchange->local) {
+      CollectionTarget target = collections.lookup(local.requirement);
+      if (!target.collection)
+        return exchange->deferred.emitOpError(
+          "local availability has no complete result-owned collection"),
+          failure();
+      auto [concat, inserted] = concatRuns.try_emplace(target.collection);
+      if (inserted)
+        concat->second = buildLocalConcat(*target.collection,
+                                          localByRequirement,
+                                          exchange->targetLaneCount);
+      if (concat->second) {
+        if (emittedConcats.insert(target.collection).second)
+          localUpdates.push_back(std::move(*concat->second));
+        if (failed(addCoverage(*local.requirement, local.targetLanes, coverage)))
+          return failure();
         continue;
       }
-      tokens.push_back(
-        {event.activeLanes, semantic, static_cast<unsigned>(eventId)});
-      if (emittedAvailabilities.insert(semantic).second) {
-        for (LocalAvailabilityFamily* local : localsBySemantic[semantic]) {
-          tokens.push_back({local->targetLanes, semantic, local});
-          emittedLocals.insert(local);
-        }
-        localsBySemantic.erase(semantic);
+      auto grouped = llvm::find_if(localUpdates, [&](EmitLocalCollectionRun &update) {
+        return update.lanes.intersect(local.targetLanes).empty()
+            && update.collection == target.collection
+            && update.collectionPosition == target.position
+            && canGroupLocalAvailability(*update.families.front()->requirement,
+                                         *local.requirement);
+      });
+      if (grouped == localUpdates.end()) {
+        localUpdates.push_back(EmitLocalCollectionRun {
+          target.collection, target.position, {&local}, local.targetLanes,
+          false});
+      } else {
+        grouped->families.push_back(&local);
+        grouped->lanes = grouped->lanes.unite(local.targetLanes);
       }
+      if (failed(addCoverage(*local.requirement, local.targetLanes, coverage)))
+        return failure();
     }
-    for (LocalAvailabilityFamily* local : locals[boundary.key])
-      if (!emittedLocals.contains(local)) {
-        unsigned semantic =
-          internSemantic(getLocalKey(*local), semantics, semanticsByHash);
-        tokens.push_back({local->targetLanes, semantic, local});
-      }
-    for (DeferredExchangePlan* exchange : exchanges[boundary.key]) {
-      SemanticKey key;
-      key.kind = SemanticKind::Result;
-      key.exchange = exchange;
-      unsigned semantic = internSemantic(key, semantics, semanticsByHash);
-      tokens.push_back({LaneSet::all(exchange->targetLaneCount), semantic, exchange});
-    }
-    if (failed(buildCanonicalBoundary(boundary, work.events, tokens, semantics)))
-      return boundary.key.scheduled->op->emitOpError("cannot construct canonical boundary program"), failure();
-    result.boundaries.push_back(std::move(boundary));
+    for (EmitLocalCollectionRun &update : localUpdates)
+      boundary.instructions.push_back(std::move(update));
+    for (RequirementFamily &requirement : exchange->requirements)
+      if (!(coverage.lookup(&requirement) == requirement.targetLanes))
+        return exchange->deferred.emitOpError(
+          "deferred availability does not cover every target lane exactly once"),
+          failure();
+    boundary.instructions.push_back(ProduceDeferredResult {exchange.get()});
+
   }
+
+  DenseMap<ScheduledInfo *, unsigned> scheduledOrder;
+  for (auto [index, scheduled] : llvm::enumerate(transfers.scheduled))
+    scheduledOrder[&scheduled] = index;
+  llvm::stable_sort(boundaries, [&](const BoundaryProgram &lhs,
+                                    const BoundaryProgram &rhs) {
+    return std::tie(scheduledOrder[lhs.key.first], lhs.key.second)
+         < std::tie(scheduledOrder[rhs.key.first], rhs.key.second);
+  });
+  result.boundaries = std::move(boundaries);
   return result;
 }
 

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

@@ -7,72 +7,37 @@ namespace onnx_mlir::spatial {
 
 struct DeferredTransferPlan;
 
-struct BoundaryKey {
-  ScheduledInfo* scheduled = nullptr;
-  unsigned insertionStep = 0;
-
-  bool operator==(const BoundaryKey& other) const {
-    return scheduled == other.scheduled
-        && insertionStep == other.insertionStep;
-  }
-};
+using BoundaryKey = std::pair<ScheduledInfo *, unsigned>;
 
 struct EmitSendRun {
   llvm::SmallVector<ScheduledTransferSlice> slices;
   LaneSet lanes;
 };
-struct EmitReceiveRun {
-  llvm::SmallVector<ScheduledTransferSlice> slices;
-  llvm::SmallVector<size_t> entryOffsets;
-  LaneSet lanes;
-};
-struct EmitReceiveBundle {
-  llvm::SmallVector<EmitReceiveRun> entries;
-};
-struct EmitReceiveAssemblyRun {
-  llvm::SmallVector<ScheduledTransferSlice> slices;
-  llvm::SmallVector<size_t> entryOffsets;
-  llvm::SmallVector<unsigned> assemblyEntries;
-  std::optional<unsigned> projectionLeaf;
-  LaneSet lanes;
-};
-struct MaterializeLocalFamily {
+struct EmitLocalCollectionRun {
+  const FragmentCollectionPlan* collection = nullptr;
+  unsigned collectionPosition = 0;
   llvm::SmallVector<LocalAvailabilityFamily*> families;
   LaneSet lanes;
+  bool concatenatePayloads = false;
 };
-using AvailabilitySource =
-  std::variant<EmitReceiveRun, MaterializeLocalFamily>;
-struct AvailabilityAlternative {
+struct EmitReceiveAssemblyRun {
+  const FragmentCollectionPlan* collection = nullptr;
+  llvm::SmallVector<ScheduledTransferSlice> slices;
+  llvm::SmallVector<size_t> entryOffsets;
+  llvm::SmallVector<unsigned> positions;
+  llvm::SmallVector<LaneSet> entryLanes;
   LaneSet lanes;
-  AvailabilitySource source;
-};
-struct ResolveAvailability {
-  DeferredExchangePlan* exchange = nullptr;
-  RequirementCoordinate coordinate;
-  mlir::Type fragmentType;
-  llvm::SmallVector<AvailabilityAlternative, 2> alternatives;
 };
 struct ProduceDeferredResult {
   DeferredExchangePlan* exchange = nullptr;
-  LaneSet lanes;
 };
 
-struct BoundaryInstructionList;
-struct LaneDispatch;
-using BoundaryInstruction = std::variant<EmitSendRun, ResolveAvailability,
-  EmitReceiveBundle, EmitReceiveAssemblyRun, ProduceDeferredResult,
-  std::unique_ptr<LaneDispatch>>;
-struct BoundaryInstructionList {
-  llvm::SmallVector<BoundaryInstruction, 0> instructions;
-};
-struct LaneDispatch {
-  StaticIntSequence branchByLane = StaticIntSequence::uniform(0, 1);
-  llvm::SmallVector<BoundaryInstructionList> branches;
-  llvm::SmallVector<DeferredExchangePlan*> producedExchanges;
-};
+using BoundaryInstruction =
+  std::variant<EmitSendRun, EmitLocalCollectionRun, EmitReceiveAssemblyRun,
+               ProduceDeferredResult>;
 struct BoundaryProgram {
   BoundaryKey key;
-  BoundaryInstructionList root;
+  llvm::SmallVector<BoundaryInstruction, 0> instructions;
 };
 
 struct DeferredBoundaryPlan {
@@ -84,22 +49,3 @@ mlir::FailureOr<DeferredBoundaryPlan> buildDeferredBoundaryPlan(DeferredTransfer
                                                                 const ScheduledCommunicationPlan& schedule);
 
 } // namespace onnx_mlir::spatial
-
-namespace llvm {
-template <>
-struct DenseMapInfo<onnx_mlir::spatial::BoundaryKey> {
-  static onnx_mlir::spatial::BoundaryKey getEmptyKey() {
-    return {DenseMapInfo<onnx_mlir::spatial::ScheduledInfo*>::getEmptyKey(), 0};
-  }
-  static onnx_mlir::spatial::BoundaryKey getTombstoneKey() {
-    return {DenseMapInfo<onnx_mlir::spatial::ScheduledInfo*>::getTombstoneKey(), 0};
-  }
-  static unsigned getHashValue(const onnx_mlir::spatial::BoundaryKey& key) {
-    return hash_combine(key.scheduled, key.insertionStep);
-  }
-  static bool isEqual(const onnx_mlir::spatial::BoundaryKey& lhs,
-                      const onnx_mlir::spatial::BoundaryKey& rhs) {
-    return lhs == rhs;
-  }
-};
-} // namespace llvm

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

@@ -9,6 +9,23 @@
 
 namespace onnx_mlir::spatial {
 
+struct FragmentCollectionKeyInfo {
+  static FragmentCollectionKey getEmptyKey() {
+    return {llvm::DenseMapInfo<DeferredExchangePlan*>::getEmptyKey()};
+  }
+  static FragmentCollectionKey getTombstoneKey() {
+    return {llvm::DenseMapInfo<DeferredExchangePlan*>::getTombstoneKey()};
+  }
+  static unsigned getHashValue(const FragmentCollectionKey& key) {
+    return llvm::hash_combine(
+      key.exchange, key.kind, key.leafIndex);
+  }
+  static bool isEqual(const FragmentCollectionKey& lhs,
+                      const FragmentCollectionKey& rhs) {
+    return lhs == rhs;
+  }
+};
+
 struct DeferredEmissionContext {
   DeferredEmissionContext(mlir::IRRewriter& rewriter,
                           ConstantPool& constants)
@@ -16,10 +33,8 @@ struct DeferredEmissionContext {
 
   mlir::IRRewriter& rewriter;
   ConstantPool& constants;
-  llvm::DenseMap<RequirementFamily*, mlir::Value> receives;
-  llvm::DenseMap<DeferredExchangePlan*, mlir::Value> assemblies;
-  llvm::DenseMap<std::pair<DeferredExchangePlan*, unsigned>, mlir::Value>
-    projectionAssemblies;
+  llvm::DenseMap<FragmentCollectionKey, mlir::Value,
+                 FragmentCollectionKeyInfo> fragmentCollections;
 };
 
 using DeferredReplacementMap =

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

@@ -5,8 +5,6 @@
 
 #include "llvm/ADT/DenseMap.h"
 
-#include <map>
-
 namespace onnx_mlir::spatial {
 using namespace mlir;
 namespace {
@@ -126,7 +124,6 @@ static LogicalResult simulatePlanned(
 }
 
 struct RealizedOperation {
-  Operation *op = nullptr;
   bool send = false;
   StaticIntSequence channels;
   StaticIntSequence parents;
@@ -135,27 +132,6 @@ struct RealizedOperation {
   StaticIntSequence targets;
 };
 
-static FailureOr<size_t> getBatchTransferCount(Operation *op) {
-  if (auto count = op->getAttrOfType<IntegerAttr>(
-        "raptor.batch_transfer_count")) {
-    if (count.getInt() > 0)
-      return count.getInt();
-    return op->emitOpError("has invalid compact transfer count"), failure();
-  }
-  if (op->hasAttr("raptor.batch_channel_ids_encoding"))
-    return op->emitOpError("is missing compact transfer count"), failure();
-  Attribute channels = op->getAttr("raptor.batch_channel_ids");
-  if (auto array = dyn_cast_or_null<DenseI64ArrayAttr>(channels))
-    return array.empty()
-      ? FailureOr<size_t>(failure())
-      : FailureOr<size_t>(array.size());
-  if (auto elements = dyn_cast_or_null<DenseIntElementsAttr>(channels);
-      elements && elements.getNumElements() > 0)
-    return elements.getNumElements();
-  return op->emitOpError("has invalid legacy compact transfer metadata"),
-         failure();
-}
-
 static FailureOr<RealizedOperation> parseRealizedOperation(Operation *op) {
   bool scalar = op->hasAttr("raptor.channel_id");
   bool batch = op->hasAttr("raptor.batch_channel_ids");
@@ -164,11 +140,14 @@ static FailureOr<RealizedOperation> parseRealizedOperation(Operation *op) {
       "must have exactly one scalar or compact metadata form");
     return failure();
   }
-  auto batchCount = scalar ? FailureOr<size_t>(1)
-                           : getBatchTransferCount(op);
-  if (failed(batchCount))
-    return failure();
-  size_t size = *batchCount;
+  size_t size = 1;
+  if (batch) {
+    auto count = op->getAttrOfType<IntegerAttr>(
+      "raptor.batch_transfer_count");
+    if (!count || count.getInt() <= 0)
+      return op->emitOpError("has invalid compact transfer count"), failure();
+    size = count.getInt();
+  }
   auto channels = getStaticIntSequenceAttr(
     op, scalar ? "raptor.channel_id" : "raptor.batch_channel_ids", size);
   auto parents = getStaticIntSequenceAttr(
@@ -191,37 +170,10 @@ static FailureOr<RealizedOperation> parseRealizedOperation(Operation *op) {
       return failure();
     }
   }
-  return RealizedOperation {op, isa<SpatChannelSendOp>(op),
+  return RealizedOperation {isa<SpatChannelSendOp>(op),
     *channels, *parents, *counts, *sources, *targets};
 }
 
-struct CoreTransferSequences {
-  DenseMap<int64_t, StaticIntSequenceChain> sends;
-  DenseMap<int64_t, StaticIntSequenceChain> receives;
-  DenseMap<int64_t, StaticIntSequenceChain> events;
-};
-
-struct ExpectedFamily {
-  ExternalTransferFamily *family = nullptr;
-  int64_t firstChannel = 0;
-  int64_t endChannel = 0;
-};
-
-static void appendByCore(DenseMap<int64_t, StaticIntSequenceChain> &result,
-                         const StaticIntSequence &channels,
-                         const StaticIntSequence &cores, size_t begin,
-                         size_t count) {
-  size_t end = begin + count;
-  cores.forEachEqualRun(
-    [&](int64_t core, size_t runBegin, size_t runCount) {
-      size_t selectedBegin = std::max(begin, runBegin);
-      size_t selectedEnd = std::min(end, runBegin + runCount);
-      if (selectedBegin < selectedEnd)
-        result[core].append(
-          channels, selectedBegin, selectedEnd - selectedBegin);
-    });
-}
-
 static void appendEventsByCore(
     DenseMap<int64_t, StaticIntSequenceChain> &result,
     const StaticIntSequence &channels, const StaticIntSequence &cores,
@@ -241,39 +193,6 @@ static void appendEventsByCore(
     });
 }
 
-static LogicalResult compareSequences(
-    func::FuncOp funcOp,
-    const DenseMap<int64_t, StaticIntSequenceChain> &expected,
-    const DenseMap<int64_t, StaticIntSequenceChain> &actual, StringRef kind) {
-  if (expected.size() != actual.size())
-    return funcOp.emitOpError()
-           << "realized " << kind << " stream set differs from plan";
-  for (const auto &[core, sequence] : expected) {
-    auto found = actual.find(core);
-    if (found == actual.end())
-      return funcOp.emitOpError() << "realized " << kind
-             << " stream is missing on core " << core;
-    StaticIntSequenceChainCursor expectedCursor(sequence);
-    StaticIntSequenceChainCursor actualCursor(found->second);
-    uint64_t ordinal = 0;
-    while (!expectedCursor.done() && !actualCursor.done()
-           && expectedCursor.value() == actualCursor.value()) {
-      expectedCursor.advance();
-      actualCursor.advance();
-      ++ordinal;
-    }
-    if (expectedCursor.done() && actualCursor.done())
-      continue;
-    return funcOp.emitOpError()
-           << "realized " << kind << " logical order differs on core "
-           << core << " at ordinal " << ordinal << ": expected channel "
-           << (expectedCursor.done() ? -1 : expectedCursor.value())
-           << ", actual channel "
-           << (actualCursor.done() ? -1 : actualCursor.value());
-  }
-  return success();
-}
-
 static LogicalResult compareEventSequences(
     func::FuncOp funcOp,
     const DenseMap<int64_t, StaticIntSequenceChain> &expected,
@@ -320,6 +239,37 @@ static LogicalResult compareEventSequences(
 LogicalResult verifyPlannedCommunicationDeadlockFree(
     Operation *anchor, ArrayRef<unsigned> stepCounts,
     const ScheduledCommunicationPlan &plan) {
+  SmallVector<std::pair<int64_t, int64_t>> familyChannels;
+  DenseMap<ExternalTransferFamily *, unsigned> familyIndex;
+  for (const ScheduledTransferSlice &slice : plan.slices) {
+    ExternalTransferFamily *family = slice.family;
+    if (!familyIndex.try_emplace(family, familyIndex.size()).second)
+      continue;
+    size_t count = family->channelIds.size();
+    if (count == 0)
+      return anchor->emitError(
+        "planned communication family has no channels");
+    int64_t first = family->channelIds.valueAt(0);
+    for (size_t index = 1; index < count; ++index)
+      if (family->channelIds.valueAt(index)
+          != first + static_cast<int64_t>(index))
+        return anchor->emitError(
+          "planned communication family has non-consecutive channels");
+    familyChannels.emplace_back(
+      first, first + static_cast<int64_t>(count));
+  }
+  llvm::sort(familyChannels);
+  int64_t nextChannel = 0;
+  for (auto [firstChannel, endChannel] : familyChannels) {
+    if (firstChannel != nextChannel)
+      return anchor->emitError(
+        "planned communication channels are not exactly contiguous");
+    nextChannel = endChannel;
+  }
+  if (static_cast<uint64_t>(nextChannel) != plan.logicalTransferCount)
+    return anchor->emitError(
+      "planned communication channel count is inconsistent");
+
   for (const ScheduledTransferSlice &slice : plan.slices) {
     ExternalTransferFamily &family = *slice.family;
     for (size_t offset = 0; offset < slice.transferCount; ++offset) {
@@ -341,55 +291,27 @@ LogicalResult verifyPlannedCommunicationDeadlockFree(
 
 LogicalResult verifyRealizedCommunicationDeadlockFree(
     func::FuncOp funcOp, const ScheduledCommunicationPlan &plan) {
-  SmallVector<ExpectedFamily> families;
+  SmallVector<ExternalTransferFamily *> familyByChannel(
+    plan.logicalTransferCount);
   DenseMap<ExternalTransferFamily *, unsigned> familyIndex;
   for (const ScheduledTransferSlice &slice : plan.slices) {
     ExternalTransferFamily *family = slice.family;
-    if (familyIndex.count(family))
+    if (!familyIndex.try_emplace(family, familyIndex.size()).second)
       continue;
-    size_t count = family->channelIds.size();
-    int64_t first = family->channelIds.valueAt(0);
-    for (size_t index = 1; index < count; ++index)
-      if (family->channelIds.valueAt(index)
-          != first + static_cast<int64_t>(index))
-        return funcOp.emitOpError(
-          "planned communication family has non-consecutive channels");
-    familyIndex[family] = families.size();
-    families.push_back({family, first, first + static_cast<int64_t>(count)});
+    for (size_t index = 0; index < family->channelIds.size(); ++index)
+      familyByChannel[family->channelIds.valueAt(index)] = family;
   }
-  llvm::sort(families, [](const ExpectedFamily &lhs,
-                          const ExpectedFamily &rhs) {
-    return lhs.firstChannel < rhs.firstChannel;
-  });
-  familyIndex.clear();
-  std::map<int64_t, unsigned> familyByFirstChannel;
-  int64_t nextChannel = 0;
-  for (auto [index, expected] : llvm::enumerate(families)) {
-    if (expected.firstChannel != nextChannel)
-      return funcOp.emitOpError(
-        "planned communication channels are not exactly contiguous");
-    nextChannel = expected.endChannel;
-    familyIndex[expected.family] = index;
-    familyByFirstChannel.emplace(expected.firstChannel, index);
-  }
-  if (static_cast<uint64_t>(nextChannel) != plan.logicalTransferCount)
-    return funcOp.emitOpError(
-      "planned communication channel count is inconsistent");
 
-  CoreTransferSequences expected;
+  DenseMap<int64_t, StaticIntSequenceChain> expected;
   for (const ScheduledTransferSlice &slice : plan.slices) {
     ExternalTransferFamily &family = *slice.family;
-    appendByCore(expected.sends, family.channelIds, family.sourceCores,
-                 slice.familyOffset, slice.transferCount);
-    appendByCore(expected.receives, family.channelIds, family.targetCores,
-                 slice.familyOffset, slice.transferCount);
-    appendEventsByCore(expected.events, family.channelIds, family.sourceCores,
+    appendEventsByCore(expected, family.channelIds, family.sourceCores,
                        slice.familyOffset, slice.transferCount, true);
-    appendEventsByCore(expected.events, family.channelIds, family.targetCores,
+    appendEventsByCore(expected, family.channelIds, family.targetCores,
                        slice.familyOffset, slice.transferCount, false);
   }
 
-  CoreTransferSequences actual;
+  DenseMap<int64_t, StaticIntSequenceChain> actual;
   SmallVector<std::unique_ptr<StaticIntSequence>> actualChannels;
   bool invalid = false;
   funcOp.walk([&](Operation *op) {
@@ -405,29 +327,28 @@ LogicalResult verifyRealizedCommunicationDeadlockFree(
       : cast<SpatChannelReceiveOp>(op).getOutput().getType();
     for (size_t index = 0; index < realized->channels.size(); ++index) {
       int64_t channel = realized->channels.valueAt(index);
-      auto upper = familyByFirstChannel.upper_bound(channel);
-      if (upper == familyByFirstChannel.begin()) {
+      if (channel < 0
+          || static_cast<uint64_t>(channel) >= familyByChannel.size()) {
         op->emitOpError("references an unknown logical channel");
         invalid = true;
         return;
       }
-      ExpectedFamily &expected = families[std::prev(upper)->second];
-      if (channel >= expected.endChannel) {
+      ExternalTransferFamily *family = familyByChannel[channel];
+      if (!family) {
         op->emitOpError("references an unknown logical channel");
         invalid = true;
         return;
       }
-      ExternalTransferFamily &family = *expected.family;
-      size_t familyOffset = channel - expected.firstChannel;
-      RequirementFamily &requirement = *family.requirement;
+      size_t familyOffset = channel - family->channelIds.valueAt(0);
+      RequirementFamily &requirement = *family->requirement;
       if (realized->parents.valueAt(index)
               != static_cast<int64_t>(requirement.exchange->exchangeId)
           || realized->counts.valueAt(index)
               != requirement.exchange->externalTransferCount
           || realized->sources.valueAt(index)
-              != family.sourceCores.valueAt(familyOffset)
+              != family->sourceCores.valueAt(familyOffset)
           || realized->targets.valueAt(index)
-              != family.targetCores.valueAt(familyOffset)
+              != family->targetCores.valueAt(familyOffset)
           || payloadType != requirement.publicationFragmentType) {
         op->emitOpError(
           "logical transfer metadata differs from its symbolic family");
@@ -439,24 +360,13 @@ LogicalResult verifyRealizedCommunicationDeadlockFree(
       return;
     actualChannels.push_back(
       std::make_unique<StaticIntSequence>(std::move(realized->channels)));
-    appendByCore(realized->send ? actual.sends : actual.receives,
-                 *actualChannels.back(),
-                 realized->send ? realized->sources : realized->targets,
-                 0, actualChannels.back()->size());
-    appendEventsByCore(actual.events, *actualChannels.back(),
+    appendEventsByCore(actual, *actualChannels.back(),
                        realized->send ? realized->sources : realized->targets,
                        0, actualChannels.back()->size(), realized->send);
   });
   if (invalid)
     return failure();
-  if (failed(compareSequences(
-        funcOp, expected.sends, actual.sends, "send"))
-      || failed(compareSequences(
-        funcOp, expected.receives, actual.receives, "receive"))
-      || failed(compareEventSequences(
-        funcOp, expected.events, actual.events)))
-    return failure();
-  return success();
+  return compareEventSequences(funcOp, expected, actual);
 }
 
 } // namespace onnx_mlir::spatial

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

@@ -104,11 +104,14 @@ private:
 };
 
 struct RequirementCoordinate {
+  unsigned specializationIndex = 0;
   unsigned leafIndex = 0;
   unsigned selectedPosition = 0;
 
   bool operator==(const RequirementCoordinate& other) const {
-    return leafIndex == other.leafIndex && selectedPosition == other.selectedPosition;
+    return specializationIndex == other.specializationIndex
+        && leafIndex == other.leafIndex
+        && selectedPosition == other.selectedPosition;
   }
 };
 
@@ -139,10 +142,10 @@ struct DeferredProjectionLeafTemplate {
   DeferredLeafForm form = DeferredLeafForm::DirectSource;
   mlir::Value sourceRoot;
   mlir::Value replacementRoot;
-  mlir::tensor::ExtractSliceOp leadingProjection;
   DeferredSliceTemplate leadingGeometry;
   DeferredSliceTemplate innerGeometry;
   mlir::RankedTensorType reconstructedType;
+  bool leadingRankReduced = false;
 };
 
 struct DeferredInsertAssemblyEntryTemplate {
@@ -160,6 +163,9 @@ struct DeferredInsertAssemblyTemplate {
 
 struct DeferredProgramTemplate {
   SpatDeferredCommunicationOp deferred;
+  unsigned specializationCount = 1;
+  mlir::Value specializationArgument;
+  mlir::RankedTensorType specializationFragmentType;
   mlir::Value scheduledLane;
   mlir::Value yieldedValue;
   llvm::SmallVector<DeferredProjectionLeafTemplate, 0> leaves;
@@ -189,9 +195,7 @@ struct ScheduledInfo {
   llvm::SmallVector<mlir::Block*> blocks;
   llvm::SmallVector<mlir::Operation*> stepAnchors;
   llvm::SmallVector<int64_t> cores;
-  llvm::SmallVector<int64_t> stepSourceIds;
-  llvm::SmallVector<int64_t> resultOffsets;
-  llvm::SmallVector<int64_t> resultCounts;
+  unsigned stepCount = 0;
   llvm::SmallVector<ProducedValue*> produced;
   llvm::SmallVector<unsigned> streamIds;
 

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

@@ -10,6 +10,8 @@
 
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 
 namespace onnx_mlir::spatial {
@@ -77,7 +79,6 @@ static FailureOr<Value> buildBlueprintReconstruction(
     sliceOffsets.push_back(builder.getIndexAttr((*sourceSlots)[fragmentIndex]));
     sliceSizes.push_back(builder.getIndexAttr(1));
     sliceStrides.push_back(builder.getIndexAttr(1));
-    SmallVector<int64_t> selectedShape {1};
     for (int64_t dim = 0; dim < rank; ++dim) {
       int64_t index = fragmentIndex * rank + dim;
       int64_t size = sizes[index];
@@ -87,21 +88,16 @@ static FailureOr<Value> buildBlueprintReconstruction(
       sliceOffsets.push_back(builder.getIndexAttr(sourceCoordinates[dim]));
       sliceSizes.push_back(builder.getIndexAttr(size));
       sliceStrides.push_back(builder.getIndexAttr(1));
-      selectedShape.push_back(size);
     }
-    auto selectedType = RankedTensorType::get(selectedShape, resultType.getElementType());
-    Value selected = tensor::ExtractSliceOp::create(
-        builder, loc, selectedType, sourceBlockArgs[operandIndex], sliceOffsets,
-        sliceSizes, sliceStrides);
-    SmallVector<int64_t> fragmentResultShape(selectedShape.begin() + 1,
-                                             selectedShape.end());
-    auto fragmentType = RankedTensorType::get(fragmentResultShape,
-                                              resultType.getElementType());
-    SmallVector<ReassociationIndices> reassociation {{0, 1}};
-    for (int64_t dim = 1; dim < rank; ++dim)
-      reassociation.push_back({dim + 1});
-    Value fragment = tensor::CollapseShapeOp::create(
-        builder, loc, fragmentType, selected, reassociation);
+    SmallVector<int64_t> selectedFragmentShape;
+    selectedFragmentShape.reserve(rank);
+    for (int64_t dim = 0; dim < rank; ++dim)
+      selectedFragmentShape.push_back(sizes[fragmentIndex * rank + dim]);
+    auto fragmentType = RankedTensorType::get(
+      selectedFragmentShape, resultType.getElementType());
+    Value fragment = extractMixedSliceOrIdentity(
+      builder, loc, sourceBlockArgs[operandIndex], fragmentType,
+      {sliceOffsets, sliceSizes, sliceStrides});
     SmallVector<OpFoldResult> targetOffsets, targetSizes, targetStrides;
     for (int64_t dim = 0; dim < rank; ++dim) {
       int64_t index = fragmentIndex * rank + dim;
@@ -116,56 +112,52 @@ static FailureOr<Value> buildBlueprintReconstruction(
   return result;
 }
 
-static FailureOr<Value> buildIndexSwitchSelection(OpBuilder &builder, Location loc,
-                                                  Value selector, ValueRange candidates,
-                                                  Operation *diagnosticOwner) {
-  if (candidates.empty())
-    return diagnosticOwner->emitOpError("direct selection requires at least one candidate"), failure();
-  Type type = candidates.front().getType();
-  if (llvm::any_of(candidates, [&](Value candidate) { return candidate.getType() != type; }))
-    return diagnosticOwner->emitOpError("direct selection requires identical candidate types"), failure();
-  if (candidates.size() == 1)
-    return candidates.front();
-
-  SmallVector<int64_t> cases;
-  for (int64_t index = 0; index < static_cast<int64_t>(candidates.size()) - 1; ++index)
-    cases.push_back(index);
-  auto selection = scf::IndexSwitchOp::create(
-    builder, loc, TypeRange {type}, selector, cases, cases.size());
-  auto buildYield = [&](Region &region, Value candidate) {
-    OpBuilder::InsertionGuard guard(builder);
-    Block *block = builder.createBlock(&region);
-    builder.setInsertionPointToEnd(block);
-    scf::YieldOp::create(builder, loc, candidate);
-  };
-  for (auto [region, candidate] : llvm::zip(selection.getCaseRegions(), candidates.drop_back()))
-    buildYield(region, candidate);
-  // The scheduled-lane verifier guarantees an in-range selector, so default is
-  // the final lane without an otherwise-unreachable extra branch.
-  buildYield(selection.getDefaultRegion(), candidates.back());
-  return selection.getResult(0);
-}
-
 static FailureOr<Value> buildSelectedDeferredSource(OpBuilder &builder, Location loc,
-                                                    SpatDeferredCommunicationOp transfer,
+                                                    Operation *diagnosticOwner,
                                                     Value scheduledLane,
                                                     ValueRange sourceBlockArgs,
                                                     ArrayRef<int64_t> sourceOperandForScheduledLane) {
   if (sourceBlockArgs.size() == 1)
     return sourceBlockArgs.front();
   if (!scheduledLane || sourceOperandForScheduledLane.empty())
-    return transfer.emitOpError("multiple deferred sources require the enclosing scheduled lane"), failure();
-  auto scheduled = transfer->getParentOfType<SpatScheduledComputeBatch>();
-  if (!scheduled || sourceOperandForScheduledLane.size() != static_cast<size_t>(scheduled.getLaneCount()))
-    return transfer.emitOpError("deferred source mapping must cover every scheduled lane"), failure();
-  SmallVector<Value> candidates;
-  candidates.reserve(sourceOperandForScheduledLane.size());
+    return diagnosticOwner->emitOpError("multiple deferred sources require the enclosing scheduled lane"), failure();
   for (int64_t sourceIndex : sourceOperandForScheduledLane) {
     if (sourceIndex < 0 || sourceIndex >= static_cast<int64_t>(sourceBlockArgs.size()))
-      return transfer.emitOpError("deferred source mapping operand is out of range"), failure();
-    candidates.push_back(sourceBlockArgs[sourceIndex]);
+      return diagnosticOwner->emitOpError("deferred source mapping operand is out of range"), failure();
   }
-  return buildIndexSwitchSelection(builder, loc, scheduledLane, candidates, transfer.getOperation());
+  StaticIntSequence sourceIndices = StaticIntSequence::fromValues(
+    sourceOperandForScheduledLane);
+  if (sourceIndices.getKind() == StaticIntSequenceKind::Uniform)
+    return sourceBlockArgs[sourceIndices.valueAt(0)];
+  Value selector;
+  if (sourceIndices.getKind() == StaticIntSequenceKind::Affine) {
+    int64_t step = sourceIndices.valueAt(1) - sourceIndices.valueAt(0);
+    selector = affineMulConst(
+      builder, loc, scheduledLane, step, diagnosticOwner);
+    selector = affineAddConst(
+      builder, loc, selector, sourceIndices.valueAt(0),
+      diagnosticOwner);
+  } else {
+    struct Run { int64_t value; size_t end; };
+    SmallVector<Run> runs;
+    sourceIndices.forEachEqualRun(
+      [&](int64_t value, size_t begin, size_t count) {
+        runs.push_back({value, begin + count});
+      });
+    selector = arith::ConstantIndexOp::create(
+      builder, loc, runs.back().value);
+    for (const Run &run : llvm::reverse(ArrayRef(runs).drop_back())) {
+      Value end = arith::ConstantIndexOp::create(builder, loc, run.end);
+      Value before = arith::CmpIOp::create(
+        builder, loc, arith::CmpIPredicate::ult, scheduledLane, end);
+      Value value = arith::ConstantIndexOp::create(builder, loc, run.value);
+      selector = arith::SelectOp::create(
+        builder, loc, before, value, selector);
+    }
+  }
+  return SpatDeferredSourceSelectOp::create(
+    builder, loc, sourceBlockArgs.front().getType(), selector,
+    sourceBlockArgs).getOutput();
 }
 
 static bool isDeferredPayloadCandidateOp(Operation *op) {
@@ -202,6 +194,10 @@ static FailureOr<Value> clonePayloadRoot(Value root, Block &body, const Deferred
   std::function<FailureOr<Value>(Value)> cloneScheduledLane = [&](Value value) -> FailureOr<Value> {
     if (mapping.contains(value)) return mapping.lookup(value);
     if (value == plan.scheduledLane) return value;
+    if (auto argument = dyn_cast<BlockArgument>(value);
+        argument && argument.getOwner() == &transfer.getBody().front()
+        && argument.getArgNumber() >= transfer.getSources().size())
+      return value;
     if (isa<BlockArgument>(value))
       return transfer.emitOpError("phase 1 payload shaping captures an unsupported block argument"), failure();
     Operation *op = value.getDefiningOp();
@@ -279,7 +275,8 @@ bool isDeferredFragmentAssemblyInput(
 LogicalResult prepareSingleCpuInput(OpBuilder &, Location loc, Value input, BlockArgument graphInput,
                                     const ComputeInstance &consumerInstance, const MergeScheduleResult &,
                                     ValueRange scheduledInputs, Block &block, unsigned firstInputArgument,
-                                    ArrayRef<ProducerValueKey> carriedKeys, Value graphLane, Value scheduledGraphLane,
+                                    const DenseMap<ProducerValueKey, MaterializedProducerRef> &availableValues,
+                                    Value graphLane, Value scheduledGraphLane,
                                     DeferredInputPlan &plan) {
   plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, {}, {}, {}, {}, 1, nullptr};
   if (isDeferredFragmentAssemblyInput(input, consumerInstance)) {
@@ -290,9 +287,14 @@ LogicalResult prepareSingleCpuInput(OpBuilder &, Location loc, Value input, Bloc
   auto producer = getProducerValueRef(input, &consumerInstance);
   if (!producer) { plan.availableValue = getBlockOperand(block, scheduledInputs, input, firstInputArgument); return success(); }
   ProducerValueKey key {producer->instance, producer->resultIndex};
-  auto carried = llvm::find(carriedKeys, key);
-  if (carried != carriedKeys.end()) {
-    plan.availableValue = block.getArgument(firstInputArgument + scheduledInputs.size() + std::distance(carriedKeys.begin(), carried));
+  auto batch = dyn_cast<SpatComputeBatch>(producer->instance.op);
+  bool hasCompleteValue = !batch
+                       || (producer->instance.laneStart == 0
+                           && producer->instance.laneCount
+                                == static_cast<uint32_t>(batch.getLaneCount()));
+  if (auto available = availableValues.find(key);
+      hasCompleteValue && available != availableValues.end()) {
+    plan.availableValue = available->second.payload;
     return success();
   }
   auto source = getOriginalProducerValue(*producer);
@@ -361,6 +363,29 @@ LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc
     if (needsIdentity) roots.push_back(plan.graphInput);
     llvm::sort(roots, [](Value a, Value b) { return a.getAsOpaquePointer() < b.getAsOpaquePointer(); });
     roots.erase(std::unique(roots.begin(), roots.end()), roots.end());
+    Value sharedSource;
+    if (!plan.blueprint) {
+      OpBuilder::InsertPoint restore = builder.saveInsertionPoint();
+      if (plan.scalarizedGraphLaneBase && plan.originalSources.size() > 1) {
+        Operation *loop = builder.getInsertionBlock()->getParentOp();
+        if (loop && !isa<scf::ForOp>(loop))
+          loop = loop->getParentOfType<scf::ForOp>();
+        if (loop)
+          builder.setInsertionPoint(loop);
+        else if (plan.scalarizedHoistBlock)
+          builder.setInsertionPointToEnd(plan.scalarizedHoistBlock);
+        else
+          return emitError(loc) << "phase 1 scalarized deferred source is missing a hoist point";
+      }
+      auto selected = buildSelectedDeferredSource(
+        builder, loc, plan.graphInput.getOwner()->getParentOp(),
+        plan.scheduledLane, plan.originalSources,
+        plan.sourceOperandForScheduledLane);
+      if (failed(selected))
+        return failure();
+      sharedSource = *selected;
+      builder.restoreInsertionPoint(restore);
+    }
     for (Value root : roots) {
       llvm::SmallPtrSet<Operation *, 16> laneDependencies;
       bool scalarize = plan.scalarizedGraphLaneBase
@@ -378,40 +403,68 @@ LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc
         else
           return emitError(loc) << "phase 1 scalarized deferred payload is missing a hoist point";
       }
-      SmallVector<Value> payloads;
       unsigned count = scalarize ? plan.scalarizedLaneCount : 1;
-      for (unsigned offset = 0; offset < count; ++offset) {
-      auto transfer = SpatDeferredCommunicationOp::create(builder, loc, root.getType(), plan.originalSources);
+      bool grouped = count > 1;
+      auto fragmentType = dyn_cast<RankedTensorType>(root.getType());
+      if (!fragmentType || !fragmentType.hasStaticShape())
+        return emitError(loc) << "phase 1 deferred payload requires a static ranked tensor";
+      Type outputType = fragmentType;
+      if (grouped) {
+        SmallVector<int64_t> groupedShape {static_cast<int64_t>(count)};
+        llvm::append_range(groupedShape, fragmentType.getShape());
+        outputType = RankedTensorType::get(groupedShape,
+                                           fragmentType.getElementType());
+      }
+      auto specialization = scalarize
+                          ? builder.getI64IntegerAttr(count)
+                          : IntegerAttr();
+      SmallVector<Value> transferSources;
+      if (plan.blueprint)
+        llvm::append_range(transferSources, plan.originalSources);
+      else
+        transferSources.push_back(sharedSource);
+      auto transfer = SpatDeferredCommunicationOp::create(
+        builder, loc, outputType, transferSources, specialization);
+      SmallVector<Type> bodyTypes(transfer.getSources().getTypes());
+      if (grouped)
+        bodyTypes.push_back(builder.getIndexType());
       Block *deferred = builder.createBlock(&transfer.getBody(), transfer.getBody().end(),
-        TypeRange {transfer.getSources().getTypes()}, SmallVector<Location>(transfer.getSources().size(), loc));
+        bodyTypes, SmallVector<Location>(bodyTypes.size(), loc));
       builder.setInsertionPointToStart(deferred);
+      ValueRange sourceArguments = deferred->getArguments().take_front(
+        transfer.getSources().size());
       auto selected = plan.blueprint
                         ? buildBlueprintReconstruction(builder, loc, plan.blueprint,
-                                                       deferred->getArguments())
-                        : buildSelectedDeferredSource(builder, loc, transfer,
-                                                      plan.scheduledLane,
-                                                      deferred->getArguments(),
-                                                      plan.sourceOperandForScheduledLane);
+                                                       sourceArguments)
+                        : FailureOr<Value>(sourceArguments.front());
       if (failed(selected)) return failure();
       Value boundGraphLane;
-      if (scalarize) {
-        boundGraphLane = affineAddConst(
-          builder, loc, plan.scalarizedGraphLaneBase, offset, transfer.getOperation());
-      }
+      if (grouped)
+        boundGraphLane = arith::AddIOp::create(
+          builder, loc, plan.scalarizedGraphLaneBase,
+          deferred->getArguments().back());
+      else if (scalarize)
+        boundGraphLane = plan.scalarizedGraphLaneBase;
       auto payload = clonePayloadRoot(root, body, plan, builder, transfer, *selected, boundGraphLane);
       if (failed(payload)) return failure();
       SpatYieldOp::create(builder, loc, *payload);
-      payloads.push_back(transfer.getOutput());
       builder.setInsertionPointAfter(transfer);
-      }
-      if (scalarize) {
+      if (grouped) {
         builder.restoreInsertionPoint(restore);
-        auto selected = buildIndexSwitchSelection(
-          builder, loc, plan.scalarizedLocalLane, payloads, root.getDefiningOp());
-        if (failed(selected)) return failure();
-        mapper.map(root, *selected);
+        SmallVector<OpFoldResult> offsets {
+          plan.scalarizedLocalLane};
+        SmallVector<OpFoldResult> sizes {builder.getIndexAttr(1)};
+        SmallVector<OpFoldResult> strides {builder.getIndexAttr(1)};
+        for (int64_t dimension : fragmentType.getShape()) {
+          offsets.push_back(builder.getIndexAttr(0));
+          sizes.push_back(builder.getIndexAttr(dimension));
+          strides.push_back(builder.getIndexAttr(1));
+        }
+        mapper.map(root, extractMixedSliceOrIdentity(
+          builder, loc, transfer.getOutput(), fragmentType,
+          {offsets, sizes, strides}));
       } else {
-        mapper.map(root, payloads.front());
+        mapper.map(root, transfer.getOutput());
       }
       collectClosure(root, body, plan, absorbed);
     }

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

@@ -31,7 +31,7 @@ LogicalResult prepareSingleCpuInput(OpBuilder &builder, Location loc, Value inpu
                                     const MergeScheduleResult &schedule,
                                     ValueRange scheduledInputs, Block &block,
                                     unsigned firstInputArgument,
-                                    ArrayRef<ProducerValueKey> carriedKeys,
+                                    const DenseMap<ProducerValueKey, MaterializedProducerRef> &availableValues,
                                     Value graphLane, Value scheduledGraphLane,
                                     DeferredInputPlan &plan);
 

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

@@ -12,141 +12,6 @@ namespace onnx_mlir::spatial {
 using namespace mlir;
 namespace {
 
-static LogicalResult validateScalarLinearization(ScheduledInfo &info) {
-  auto scheduled = cast<SpatScheduledCompute>(info.op);
-  for (unsigned index = 1; index < info.blocks.size(); ++index) {
-    auto previous = dyn_cast<SpatBlockYieldOp>(
-      info.blocks[index - 1]->getTerminator());
-    if (!previous
-        || previous.getOutputs().size()
-             != info.blocks[index]->getNumArguments())
-      return scheduled.emitOpError(
-        "phase 2 cannot linearize malformed scalar scheduled blocks");
-    for (auto [argument, value] : llvm::zip(
-           info.blocks[index]->getArguments(), previous.getOutputs())) {
-      if (argument.getType() == value.getType())
-        continue;
-      for (Operation *user : argument.getUsers())
-        if (!isa<SpatBlockYieldOp>(user))
-          return scheduled.emitOpError(
-            "phase 2 cannot linearize a live mismatched carried value");
-    }
-  }
-  return success();
-}
-
-static LogicalResult linearizeScalar(ScheduledInfo &info,
-                                     IRRewriter &rewriter) {
-  auto scheduled = cast<SpatScheduledCompute>(info.op);
-  if (failed(validateScalarLinearization(info)))
-    return failure();
-  Block *first = info.blocks.front();
-  SmallVector<SmallVector<Value>> incoming(info.blocks.size());
-  for (unsigned index = 1; index < info.blocks.size(); ++index) {
-    auto previous = cast<SpatBlockYieldOp>(
-      info.blocks[index - 1]->getTerminator());
-    incoming[index].assign(previous.getOutputs().begin(),
-                           previous.getOutputs().end());
-  }
-  IRMapping carried;
-  for (unsigned index = 1; index < info.blocks.size(); ++index)
-    for (auto [argument, value] : llvm::zip(
-           info.blocks[index]->getArguments(), incoming[index])) {
-      Value resolved = carried.lookupOrDefault(value);
-      if (argument.getType() == resolved.getType()) {
-        carried.map(argument, resolved);
-        argument.replaceAllUsesWith(resolved);
-      }
-    }
-  for (unsigned index = 1; index < info.blocks.size(); ++index)
-    for (Operation &op : llvm::make_early_inc_range(
-           info.blocks[index]->without_terminator()))
-      op.moveBefore(first->getTerminator());
-  for (Block *block : info.blocks)
-    cast<SpatBlockYieldOp>(block->getTerminator()).erase();
-  SmallVector<Value> outputs(scheduled.getNumResults());
-  for (ProducedValue *produced : info.produced) {
-    unsigned result = info.resultOffsets[produced->step]
-        + produced->resultIndex;
-    outputs[result] = produced->payload;
-  }
-  if (llvm::any_of(outputs, [](Value output) { return !output; }))
-    return scheduled.emitOpError(
-      "phase 2 cannot recover every scheduled scalar result");
-  rewriter.setInsertionPointToEnd(first);
-  SpatYieldOp::create(rewriter, scheduled.getLoc(), outputs);
-  scheduled->setAttr("scheduled.realized", rewriter.getBoolAttr(true));
-  for (unsigned index = 1; index < info.blocks.size(); ++index) {
-    if (llvm::any_of(info.blocks[index]->getArguments(),
-                     [](BlockArgument argument) {
-                       return !argument.use_empty();
-                     }))
-      return scheduled.emitOpError(
-        "phase 2 scalar linearization left a live block argument");
-    info.blocks[index]->erase();
-  }
-  info.blocks.assign(1, first);
-  return success();
-}
-
-static LogicalResult linearizeBatch(ScheduledInfo &info,
-                                    IRRewriter &rewriter) {
-  auto scheduled = cast<SpatScheduledComputeBatch>(info.op);
-  Block *first = info.blocks.front();
-  SmallVector<SpatInParallelOp> terminators;
-  for (Block *block : info.blocks) {
-    auto parallel = dyn_cast<SpatInParallelOp>(block->getTerminator());
-    if (!parallel)
-      return scheduled.emitOpError(
-        "phase 2 cannot linearize a batch block without spat.in_parallel");
-    terminators.push_back(parallel);
-  }
-  for (unsigned index = 1; index < info.blocks.size(); ++index)
-    for (auto [argument, firstArgument] : llvm::zip(
-           info.blocks[index]->getArguments(), first->getArguments())) {
-      if (argument.getType() != firstArgument.getType())
-        return scheduled.emitOpError(
-          "phase 2 cannot linearize incompatible batch block arguments");
-      argument.replaceAllUsesWith(firstArgument);
-    }
-  for (unsigned index = 1; index < info.blocks.size(); ++index)
-    for (Operation &op : llvm::make_early_inc_range(
-           info.blocks[index]->without_terminator()))
-      op.moveBefore(first->getTerminator());
-  rewriter.setInsertionPoint(first->getTerminator());
-  auto combined = SpatInParallelOp::create(rewriter, scheduled.getLoc());
-  Block &combinedBlock = combined.getRegion().front();
-  for (SpatInParallelOp parallel : terminators)
-    for (Operation &op : llvm::make_early_inc_range(
-           parallel.getRegion().front()))
-      op.moveBefore(&combinedBlock, combinedBlock.end());
-  for (SpatInParallelOp parallel : terminators)
-    parallel.erase();
-  scheduled->setAttr("scheduled.realized", rewriter.getBoolAttr(true));
-  for (unsigned index = 1; index < info.blocks.size(); ++index) {
-    if (llvm::any_of(info.blocks[index]->getArguments(),
-                     [](BlockArgument argument) {
-                       return !argument.use_empty();
-                     }))
-      return scheduled.emitOpError(
-        "phase 2 batch linearization left a live block argument");
-    info.blocks[index]->erase();
-  }
-  info.blocks.assign(1, first);
-  return success();
-}
-
-static LogicalResult linearizeScheduled(
-    DeferredTransferPlan &plan, IRRewriter &rewriter) {
-  for (ScheduledInfo &info : plan.scheduled) {
-    LogicalResult result = info.isBatch()
-      ? linearizeBatch(info, rewriter) : linearizeScalar(info, rewriter);
-    if (failed(result))
-      return failure();
-  }
-  return success();
-}
-
 static LogicalResult replaceFinalGraphPublications(
     func::FuncOp funcOp, DeferredTransferPlan &plan) {
   for (Operation &op : funcOp.getOps()) {
@@ -157,15 +22,30 @@ static LogicalResult replaceFinalGraphPublications(
       continue;
     for (auto [resultIndex, result] : llvm::enumerate(op.getResults())) {
       SmallVector<OpOperand *> externalUses;
-      for (OpOperand &use : result.getUses())
-        if (!isa<SpatGraphCompute, SpatGraphComputeBatch>(use.getOwner()))
-          externalUses.push_back(&use);
+      for (OpOperand &use : result.getUses()) {
+        Operation *user = use.getOwner();
+        if (isa<SpatGraphCompute, SpatGraphComputeBatch,
+                SpatDeferredCommunicationOp>(user))
+          continue;
+        if (auto blueprint = dyn_cast<SpatBlueprintOp>(user)) {
+          bool blueprintEscapes = llvm::any_of(
+            blueprint.getOutput().getUses(), [](OpOperand &blueprintUse) {
+              return !isa<SpatGraphCompute, SpatGraphComputeBatch,
+                          SpatDeferredCommunicationOp>(
+                blueprintUse.getOwner());
+            });
+          if (!blueprintEscapes)
+            continue;
+        }
+        externalUses.push_back(&use);
+      }
       if (externalUses.empty())
         continue;
       SmallVector<Value> exact;
       for (ProducedValue *produced :
            plan.producedByGraph.lookup(graphId.getInt()))
         if (produced->resultIndex == resultIndex
+            && produced->published
             && produced->published.getType() == result.getType()
             && !llvm::is_contained(exact, produced->published))
           exact.push_back(produced->published);
@@ -187,23 +67,36 @@ static LogicalResult replaceFinalGraphPublications(
 
 static LogicalResult eraseOldGraph(func::FuncOp funcOp,
                                    IRRewriter &rewriter) {
-  SmallVector<Operation *> graphOps;
+  SmallVector<Operation *> oldGraph;
   for (Operation &op : funcOp.getOps())
-    if (isa<SpatGraphCompute, SpatGraphComputeBatch>(op))
-      graphOps.push_back(&op);
-  for (Operation *op : llvm::reverse(graphOps)) {
+    if (isa<SpatGraphCompute, SpatGraphComputeBatch, SpatBlueprintOp>(op))
+      oldGraph.push_back(&op);
+  for (Operation *op : llvm::reverse(oldGraph)) {
+    if (auto blueprint = dyn_cast<SpatBlueprintOp>(op)) {
+      if (blueprint.getOutput().use_empty())
+        rewriter.eraseOp(blueprint);
+      continue;
+    }
     if (!op->use_empty())
       return op->emitOpError(
         "phase 2 cannot erase an old graph compute with live results");
     rewriter.eraseOp(op);
   }
-  SmallVector<SpatBlueprintOp> deadBlueprints;
-  funcOp.walk([&](SpatBlueprintOp blueprint) {
-    if (blueprint.getOutput().use_empty())
-      deadBlueprints.push_back(blueprint);
+  return success();
+}
+
+static LogicalResult eraseDeferredSourceSelectors(
+    func::FuncOp funcOp, IRRewriter &rewriter) {
+  SmallVector<SpatDeferredSourceSelectOp> selectors;
+  funcOp.walk([&](SpatDeferredSourceSelectOp selector) {
+    selectors.push_back(selector);
   });
-  for (SpatBlueprintOp blueprint : deadBlueprints)
-    rewriter.eraseOp(blueprint);
+  for (SpatDeferredSourceSelectOp selector : llvm::reverse(selectors)) {
+    if (!selector.getOutput().use_empty())
+      return selector.emitOpError(
+        "phase 2 left a live deferred source selection");
+    rewriter.eraseOp(selector);
+  }
   return success();
 }
 
@@ -223,8 +116,10 @@ static LogicalResult verifyDominance(func::FuncOp funcOp) {
 
 } // namespace
 
-LogicalResult realizeDeferredCommunication(func::FuncOp funcOp) {
-  auto transfers = buildDeferredTransferPlan(funcOp);
+LogicalResult realizeDeferredCommunication(
+    func::FuncOp funcOp,
+    const ScheduledComputeMaterializationResult &materialization) {
+  auto transfers = buildDeferredTransferPlan(funcOp, materialization);
   if (failed(transfers))
     return funcOp.emitOpError(
       "phase 2 failed to build symbolic transfer families");
@@ -240,7 +135,8 @@ LogicalResult realizeDeferredCommunication(func::FuncOp funcOp) {
       "phase 2 failed to build sparse boundary programs");
 
   IRRewriter rewriter(funcOp.getContext());
-  if (failed(linearizeScheduled(*transfers, rewriter)))
+  if (failed(retargetDeferredPublications(funcOp, *transfers))
+      || failed(replaceFinalGraphPublications(funcOp, *transfers)))
     return failure();
   ConstantPool constants(funcOp, rewriter);
   DeferredEmissionContext context(rewriter, constants);
@@ -258,8 +154,7 @@ LogicalResult realizeDeferredCommunication(func::FuncOp funcOp) {
         "phase 2 cannot erase deferred communication with live uses");
     rewriter.eraseOp(op);
   }
-  if (failed(retargetDeferredPublications(funcOp, *transfers))
-      || failed(replaceFinalGraphPublications(funcOp, *transfers))
+  if (failed(eraseDeferredSourceSelectors(funcOp, rewriter))
       || failed(eraseOldGraph(funcOp, rewriter))
       || failed(verifyDominance(funcOp))
       || failed(verifyRealizedCommunicationDeadlockFree(funcOp, *schedule)))

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

@@ -4,6 +4,10 @@
 
 namespace onnx_mlir::spatial {
 
-mlir::LogicalResult realizeDeferredCommunication(mlir::func::FuncOp funcOp);
+struct ScheduledComputeMaterializationResult;
+
+mlir::LogicalResult realizeDeferredCommunication(
+    mlir::func::FuncOp funcOp,
+    const ScheduledComputeMaterializationResult &materialization);
 
 } // namespace onnx_mlir::spatial

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

@@ -10,6 +10,20 @@ namespace onnx_mlir::spatial {
 using namespace mlir;
 namespace {
 
+using TransferEmissionSignature =
+  std::tuple<ScheduledInfo*, Value, Type, bool, bool, bool>;
+
+static TransferEmissionSignature getTransferEmissionSignature(
+    const ExternalTransferFamily& family) {
+  ProducedValue* producer = family.requirement->producer;
+  return {producer->scheduled,
+          producer->payload,
+          family.requirement->publicationFragmentType,
+          family.requirement->graphLanes.has_value(),
+          family.requirement->producerProjection.has_value(),
+          producer->scheduled->isBatch()};
+}
+
 struct StreamThreshold {
   unsigned stream = 0;
   unsigned completedStep = 0;
@@ -20,11 +34,10 @@ struct TransferGroup {
   SmallVector<ScheduledTransferSlice> ordered;
   SmallVector<StreamThreshold> dependencies;
   unsigned unsatisfied = 0;
-  bool ready = false;
   bool scheduled = false;
-  BoundaryCost cost;
   std::tuple<unsigned, unsigned, unsigned, uint64_t> originalPriority;
-  std::optional<TransferEmissionSignature> firstSignature;
+  std::tuple<unsigned, unsigned, unsigned,
+             std::tuple<unsigned, unsigned, unsigned, uint64_t>> staticPriority;
 };
 
 struct StreamProgress {
@@ -35,28 +48,16 @@ struct StreamProgress {
 };
 
 static size_t hashSignature(const TransferEmissionSignature& signature) {
-  return llvm::hash_combine(signature.scheduled,
-                            signature.payload.getAsOpaquePointer(),
-                            signature.fragmentType.getAsOpaquePointer(),
-                            signature.hasGraphLane,
-                            signature.hasProducerProjection,
-                            signature.sourceIsBatch);
+  return llvm::hash_combine(std::get<0>(signature),
+                            std::get<1>(signature).getAsOpaquePointer(),
+                            std::get<2>(signature).getAsOpaquePointer(),
+                            std::get<3>(signature),
+                            std::get<4>(signature),
+                            std::get<5>(signature));
 }
 
 static bool betterStaticPriority(const TransferGroup& lhs, const TransferGroup& rhs) {
-  auto left = std::tuple(lhs.cost.instructionCount,
-                         lhs.cost.branchRegions,
-                         std::numeric_limits<unsigned>::max()
-                           - lhs.cost.absorbedTransfers,
-                         lhs.cost.lookupEntries,
-                         lhs.originalPriority);
-  auto right = std::tuple(rhs.cost.instructionCount,
-                          rhs.cost.branchRegions,
-                          std::numeric_limits<unsigned>::max()
-                            - rhs.cost.absorbedTransfers,
-                          rhs.cost.lookupEntries,
-                          rhs.originalPriority);
-  return left < right;
+  return lhs.staticPriority < rhs.staticPriority;
 }
 
 static bool orderPermutationCycles(SmallVectorImpl<ScheduledTransferSlice>& slices) {
@@ -148,9 +149,29 @@ static void orderGroupSlices(TransferGroup& group) {
       });
     llvm::append_range(group.ordered, slicesBySignature[id]);
   }
-  group.cost = estimateCanonicalBoundaryCost({}, group.ordered);
-  if (!group.ordered.empty())
-    group.firstSignature = getTransferEmissionSignature(*group.ordered.front().family);
+  unsigned instructionCount = 0;
+  unsigned absorbedTransfers = 0;
+  unsigned transferCount = 0;
+  std::optional<TransferEmissionSignature> previousSignature;
+  RequirementFamily* previousRequirement = nullptr;
+  for (const ScheduledTransferSlice& slice : group.ordered) {
+    TransferEmissionSignature signature =
+      getTransferEmissionSignature(*slice.family);
+    if (!previousSignature || *previousSignature != signature)
+      ++instructionCount;
+    else
+      absorbedTransfers += slice.transferCount;
+    previousSignature = signature;
+    if (previousRequirement != slice.family->requirement)
+      ++instructionCount;
+    previousRequirement = slice.family->requirement;
+    transferCount += slice.transferCount;
+  }
+  group.staticPriority =
+    {instructionCount,
+     std::numeric_limits<unsigned>::max() - absorbedTransfers,
+     transferCount,
+     group.originalPriority};
 }
 
 static SmallVector<TransferGroup> buildGroups(DeferredTransferPlan& plan) {
@@ -213,39 +234,9 @@ static unsigned tailExtension(const TransferGroup& group,
 
 } // namespace
 
-TransferEmissionSignature getTransferEmissionSignature(const ExternalTransferFamily& family) {
-  ProducedValue* producer = family.requirement->producer;
-  return {producer->scheduled,
-          producer->payload,
-          family.requirement->publicationFragmentType,
-          family.requirement->graphLanes.has_value(),
-          family.requirement->producerProjection.has_value(),
-          producer->scheduled->isBatch()};
-}
-
-BoundaryCost estimateCanonicalBoundaryCost(
-    ArrayRef<ScheduledTransferSlice> existingTail,
-    ArrayRef<ScheduledTransferSlice> candidate) {
-  BoundaryCost cost;
-  std::optional<TransferEmissionSignature> previous;
-  RequirementFamily* previousRequirement = nullptr;
-  if (!existingTail.empty()) {
-    previous = getTransferEmissionSignature(*existingTail.back().family);
-    previousRequirement = existingTail.back().family->requirement;
-  }
-  for (const ScheduledTransferSlice& slice : candidate) {
-    TransferEmissionSignature signature = getTransferEmissionSignature(*slice.family);
-    if (!previous || !(*previous == signature))
-      ++cost.instructionCount;
-    else
-      cost.absorbedTransfers += slice.transferCount;
-    previous = signature;
-    if (previousRequirement != slice.family->requirement)
-      ++cost.instructionCount;
-    previousRequirement = slice.family->requirement;
-    cost.lookupEntries += slice.transferCount;
-  }
-  return cost;
+bool haveSameTransferEmissionSignature(
+    const ExternalTransferFamily& lhs, const ExternalTransferFamily& rhs) {
+  return getTransferEmissionSignature(lhs) == getTransferEmissionSignature(rhs);
 }
 
 FailureOr<ScheduledCommunicationPlan> scheduleDeferredCommunication(func::FuncOp funcOp, DeferredTransferPlan& plan) {
@@ -269,16 +260,12 @@ FailureOr<ScheduledCommunicationPlan> scheduleDeferredCommunication(func::FuncOp
   DenseMap<size_t, std::unique_ptr<Heap>> bySignature;
   auto addReady = [&](unsigned index) {
     TransferGroup& group = groups[index];
-    if (group.ready || group.scheduled)
-      return;
-    group.ready = true;
     ready.push(index);
-    if (group.firstSignature) {
-      auto& bucket = bySignature[hashSignature(*group.firstSignature)];
-      if (!bucket)
-        bucket = std::make_unique<Heap>(compare);
-      bucket->push(index);
-    }
+    auto& bucket = bySignature[hashSignature(
+      getTransferEmissionSignature(*group.ordered.front().family))];
+    if (!bucket)
+      bucket = std::make_unique<Heap>(compare);
+    bucket->push(index);
   };
   for (auto [index, group] : llvm::enumerate(groups))
     if (group.unsatisfied == 0)
@@ -331,7 +318,7 @@ FailureOr<ScheduledCommunicationPlan> scheduleDeferredCommunication(func::FuncOp
         unsigned candidate = bucket->second->top();
         bucket->second->pop();
         TransferGroup& group = groups[candidate];
-        if (!group.ready || group.scheduled)
+        if (group.scheduled)
           continue;
         inspected.push_back(candidate);
         unsigned extension = tailExtension(
@@ -350,12 +337,11 @@ FailureOr<ScheduledCommunicationPlan> scheduleDeferredCommunication(func::FuncOp
     while (!chosen && !ready.empty()) {
       unsigned candidate = ready.top();
       ready.pop();
-      if (groups[candidate].ready && !groups[candidate].scheduled)
+      if (!groups[candidate].scheduled)
         chosen = candidate;
     }
     if (chosen) {
       TransferGroup& group = groups[*chosen];
-      group.ready = false;
       group.scheduled = true;
       ++finishedGroups;
       for (ScheduledTransferSlice slice : group.ordered) {

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

@@ -21,35 +21,8 @@ struct ScheduledCommunicationPlan {
   uint64_t logicalTransferCount = 0;
 };
 
-struct TransferEmissionSignature {
-  ScheduledInfo* scheduled = nullptr;
-  mlir::Value payload;
-  mlir::Type fragmentType;
-  bool hasGraphLane = false;
-  bool hasProducerProjection = false;
-  bool sourceIsBatch = false;
-
-  bool operator==(const TransferEmissionSignature& other) const {
-    return scheduled == other.scheduled && payload == other.payload
-        && fragmentType == other.fragmentType
-        && hasGraphLane == other.hasGraphLane
-        && hasProducerProjection == other.hasProducerProjection
-        && sourceIsBatch == other.sourceIsBatch;
-  }
-};
-
-struct BoundaryCost {
-  unsigned instructionCount = 0;
-  unsigned branchRegions = 0;
-  unsigned absorbedTransfers = 0;
-  unsigned lookupEntries = 0;
-};
-
-TransferEmissionSignature getTransferEmissionSignature(const ExternalTransferFamily& family);
-
-BoundaryCost estimateCanonicalBoundaryCost(
-  mlir::ArrayRef<ScheduledTransferSlice> existingTail,
-  mlir::ArrayRef<ScheduledTransferSlice> candidate);
+bool haveSameTransferEmissionSignature(
+  const ExternalTransferFamily& lhs, const ExternalTransferFamily& rhs);
 
 mlir::FailureOr<ScheduledCommunicationPlan> scheduleDeferredCommunication(mlir::func::FuncOp funcOp,
                                                                           DeferredTransferPlan& plan);

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

@@ -1,6 +1,7 @@
 #include "DeferredProjectionAnalysis.hpp"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/Matchers.h"
 
@@ -109,27 +110,20 @@ static FailureOr<std::optional<unsigned>> sourceArgument(
       && argument.getArgNumber() < deferred.getSources().size())
     return std::optional<unsigned>(argument.getArgNumber());
   auto result = dyn_cast<OpResult>(value);
-  auto selection = result
-    ? dyn_cast<scf::IndexSwitchOp>(result.getOwner()) : scf::IndexSwitchOp();
-  if (!selection || result.getResultNumber() != 0
-      || selection.getNumResults() != 1)
-    return std::optional<unsigned>();
-  auto selector = evaluateDeferredIndex(selection.getArg(), environment);
-  if (failed(selector))
-    return failure();
-  Region *selected = &selection.getDefaultRegion();
-  for (auto [caseValue, region] :
-       llvm::zip(selection.getCases(), selection.getCaseRegions()))
-    if (caseValue == *selector) {
-      selected = &region;
-      break;
-    }
-  auto yield = selected->hasOneBlock()
-    ? dyn_cast<scf::YieldOp>(selected->front().getTerminator())
-    : scf::YieldOp();
-  return yield && yield.getResults().size() == 1
-    ? sourceArgument(yield.getResults().front(), deferred, environment)
-    : FailureOr<std::optional<unsigned>>(failure());
+  auto compactSelection = result
+    ? dyn_cast<SpatDeferredSourceSelectOp>(result.getOwner())
+    : SpatDeferredSourceSelectOp();
+  if (compactSelection && result.getResultNumber() == 0) {
+    auto selector = evaluateDeferredIndex(
+      compactSelection.getSelector(), environment);
+    if (failed(selector) || *selector < 0
+        || *selector >= static_cast<int64_t>(
+             compactSelection.getSources().size()))
+      return failure();
+    return sourceArgument(
+      compactSelection.getSources()[*selector], deferred, environment);
+  }
+  return std::optional<unsigned>();
 }
 
 static void collectSourceArguments(Value value,
@@ -144,17 +138,13 @@ static void collectSourceArguments(Value value,
     return;
   }
   auto result = dyn_cast<OpResult>(value);
-  auto selection = result
-    ? dyn_cast<scf::IndexSwitchOp>(result.getOwner()) : scf::IndexSwitchOp();
-  if (!selection || result.getResultNumber() != 0)
-    return;
-  for (Region &region : selection.getCaseRegions())
-    collectSourceArguments(
-      cast<scf::YieldOp>(region.front().getTerminator()).getResults().front(),
-      deferred, indices);
-  collectSourceArguments(
-    cast<scf::YieldOp>(selection.getDefaultRegion().front().getTerminator())
-      .getResults().front(), deferred, indices);
+  auto compactSelection = result
+    ? dyn_cast<SpatDeferredSourceSelectOp>(result.getOwner())
+    : SpatDeferredSourceSelectOp();
+  if (compactSelection && result.getResultNumber() == 0) {
+    for (Value source : compactSelection.getSources())
+      collectSourceArguments(source, deferred, indices);
+  }
 }
 
 static Value getEnclosingScheduledLane(
@@ -168,34 +158,6 @@ static Value getEnclosingScheduledLane(
   return block && !block->empty() ? block->getArgument(0) : Value();
 }
 
-static bool isAllowedStaticIndexExpression(
-    Value value, Value scheduledLane,
-    llvm::SmallDenseSet<Value, 16> &visiting) {
-  if (value == scheduledLane)
-    return true;
-  Attribute constant;
-  if (matchPattern(value, m_Constant(&constant)))
-    return true;
-  if (isa<BlockArgument>(value) || !value.getDefiningOp()
-      || !visiting.insert(value).second)
-    return false;
-  Operation *op = value.getDefiningOp();
-  bool allowed = op->getNumRegions() == 0
-      && (isPureIndexComputationOp(op) || isCompileTimeOp(op))
-      && llvm::all_of(op->getOperands(), [&](Value operand) {
-           return isAllowedStaticIndexExpression(
-             operand, scheduledLane, visiting);
-         });
-  visiting.erase(value);
-  return allowed;
-}
-
-static bool isAllowedStaticIndexExpression(Value value,
-                                           Value scheduledLane) {
-  llvm::SmallDenseSet<Value, 16> visiting;
-  return isAllowedStaticIndexExpression(value, scheduledLane, visiting);
-}
-
 static bool originatesFromDeferredSource(
     Value value, SpatDeferredCommunicationOp deferred,
     llvm::SmallDenseSet<Value, 16> &visited) {
@@ -205,10 +167,9 @@ static bool originatesFromDeferredSource(
     return argument.getOwner() == &deferred.getBody().front()
         && argument.getArgNumber() < deferred.getSources().size();
   Operation *op = value.getDefiningOp();
-  return op && (isa<scf::IndexSwitchOp>(op)
-      || llvm::any_of(op->getOperands(), [&](Value operand) {
-           return originatesFromDeferredSource(operand, deferred, visited);
-         }));
+  return op && llvm::any_of(op->getOperands(), [&](Value operand) {
+    return originatesFromDeferredSource(operand, deferred, visited);
+  });
 }
 
 static bool originatesFromDeferredSource(
@@ -217,43 +178,6 @@ static bool originatesFromDeferredSource(
   return originatesFromDeferredSource(value, deferred, visited);
 }
 
-static LogicalResult verifyCanonicalSourceSelector(
-    scf::IndexSwitchOp selection, SpatDeferredCommunicationOp deferred,
-    Value scheduledLane, int64_t laneCount) {
-  if (selection->getBlock() != &deferred.getBody().front()
-      || selection.getNumResults() != 1
-      || !selection.getArg().getType().isIndex()
-      || !isAllowedStaticIndexExpression(selection.getArg(), scheduledLane))
-    return selection.emitOpError(
-      "is not a canonical deferred source selector");
-  if (laneCount < 2
-      || selection.getCases().size() != static_cast<size_t>(laneCount - 1))
-    return selection.emitOpError(
-      "must cover every non-default scheduled lane");
-  for (auto [index, caseValue] : llvm::enumerate(selection.getCases()))
-    if (caseValue != static_cast<int64_t>(index))
-      return selection.emitOpError(
-        "must use consecutive scheduled-lane cases starting at zero");
-  auto verifyRegion = [&](Region &region) -> LogicalResult {
-    auto yield = region.hasOneBlock()
-      ? dyn_cast<scf::YieldOp>(region.front().getTerminator()) : scf::YieldOp();
-    if (!yield || yield.getResults().size() != 1)
-      return selection.emitOpError(
-        "source-selector region must yield one result");
-    for (Operation &op : region.front().without_terminator())
-      if (!isa<tensor::CastOp>(op))
-        return selection.emitOpError(
-          "source-selector regions may contain only tensor.cast");
-    llvm::SmallSet<unsigned, 4> indices;
-    collectSourceArguments(yield.getResults().front(), deferred, indices);
-    return success(indices.size() == 1);
-  };
-  for (Region &region : selection.getCaseRegions())
-    if (failed(verifyRegion(region)))
-      return failure();
-  return verifyRegion(selection.getDefaultRegion());
-}
-
 static bool isInsideDeferredLoop(Operation *op,
                                  SpatDeferredCommunicationOp deferred) {
   for (Operation *parent = op->getParentOp(); parent && parent != deferred;
@@ -263,6 +187,83 @@ static bool isInsideDeferredLoop(Operation *op,
   return false;
 }
 
+static FailureOr<SmallVector<unsigned>>
+getPossibleDeferredSourceOperandIndices(
+    Value sourceRoot, SpatDeferredCommunicationOp deferred) {
+  llvm::SmallSet<unsigned, 4> indices;
+  collectSourceArguments(sourceRoot, deferred, indices);
+  if (indices.empty())
+    return failure();
+  return SmallVector<unsigned>(indices.begin(), indices.end());
+}
+
+static LogicalResult validateDeferredProgram(
+    SpatDeferredCommunicationOp deferred, Block &body, Value scheduledLane,
+    Value specialization, int64_t laneCount, int64_t specializationCount) {
+  auto specializes = [&](Value value, auto &&accept) {
+    for (int64_t specializationIndex = 0;
+         specializationIndex < specializationCount; ++specializationIndex)
+      for (int64_t lane = 0; lane < laneCount; ++lane) {
+        StaticIndexEnvironment environment;
+        if (scheduledLane)
+          environment.bindings[scheduledLane] = lane;
+        if (specialization)
+          environment.bindings[specialization] = specializationIndex;
+        auto result = evaluateDeferredIndex(value, environment);
+        if (failed(result) || !accept(*result))
+          return false;
+      }
+    return true;
+  };
+  WalkResult result = body.walk([&](Operation *op) -> WalkResult {
+    auto reject = [&](StringRef message) {
+      op->emitOpError(message);
+      return WalkResult::interrupt();
+    };
+    if (isa<SpatYieldOp, scf::YieldOp>(op)
+        || (isa<linalg::YieldOp>(op)
+            && isa<linalg::TransposeOp>(op->getParentOp())))
+      return WalkResult::advance();
+    if (auto selection = dyn_cast<SpatDeferredSourceSelectOp>(op)) {
+      if (selection->getBlock() != &body)
+        return reject("must be a top-level deferred source selector");
+      if (selection.getSources().empty()
+          || !selection.getSelector().getType().isIndex()
+          || llvm::any_of(selection.getSources(), [&](Value source) {
+               return source.getType() != selection.getOutput().getType()
+                   || failed(getPossibleDeferredSourceOperandIndices(
+                        source, deferred));
+             }))
+        return reject("has invalid deferred source operands or types");
+      if (!specializes(selection.getSelector(), [&](int64_t index) {
+            return index >= 0
+                && index < static_cast<int64_t>(selection.getSources().size());
+          }))
+        return reject("source selector does not specialize to a valid source");
+      return WalkResult::advance();
+    }
+    if (auto loop = dyn_cast<scf::ForOp>(op)) {
+      for (Value bound : {loop.getLowerBound(), loop.getUpperBound()})
+        if (!specializes(bound, [](int64_t) { return true; }))
+          return reject("has a loop bound that does not specialize");
+      if (!specializes(loop.getStep(), [](int64_t step) { return step > 0; }))
+        return reject("has a non-positive or non-static deferred loop step");
+      return WalkResult::advance();
+    }
+    if (op->getNumRegions() != 0 && !isa<linalg::TransposeOp>(op))
+      return reject("contains an unsupported region operation");
+    if (!isShapingOnlyOp(op) && !isCompileTimeOp(op)
+        && !isPureIndexComputationOp(op))
+      return reject("contains an unsupported deferred operation");
+    for (Value operand : op->getOperands())
+      if (isInsideDeferredLoop(op, deferred)
+          && originatesFromDeferredSource(operand, deferred))
+        return reject("projects a deferred source inside a residual loop");
+    return WalkResult::advance();
+  });
+  return failure(result.wasInterrupted());
+}
+
 static bool haveLeadingUnitDifference(RankedTensorType larger,
                                       RankedTensorType smaller) {
   return larger && smaller && larger.hasStaticShape() && smaller.hasStaticShape()
@@ -314,7 +315,7 @@ analyzeInsertAssembly(const DeferredProgramTemplate &program) {
     if (!transform)
       return std::optional<DeferredInsertAssemblyTemplate>();
     DeferredInsertAssemblyEntryTemplate entry;
-    entry.coordinate = {leaf->second, 0};
+    entry.coordinate = {0, leaf->second, 0};
     entry.sourceTransform = *transform;
     entry.sourceType = insertedType;
     entry.targetGeometry = {SmallVector<OpFoldResult>(insert.getMixedOffsets()),
@@ -345,66 +346,6 @@ analyzeInsertAssembly(const DeferredProgramTemplate &program) {
 
 } // namespace
 
-LogicalResult verifyDeferredProgramContract(
-    SpatDeferredCommunicationOp deferred) {
-  if (!deferred.getBody().hasOneBlock())
-    return deferred.emitOpError(
-      "deferred program must have exactly one body block");
-  Block &body = deferred.getBody().front();
-  auto terminator = dyn_cast<SpatYieldOp>(body.getTerminator());
-  if (!terminator || terminator.getOutputs().size() != 1)
-    return deferred.emitOpError(
-      "deferred program must have exactly one yielded value");
-  Value scheduledLane;
-  int64_t laneCount = 1;
-  if (auto scheduled = deferred->getParentOfType<SpatScheduledComputeBatch>()) {
-    scheduledLane = getEnclosingScheduledLane(deferred, scheduled);
-    laneCount = scheduled.getLaneCount();
-  }
-  bool invalid = false;
-  deferred.getBody().walk([&](Operation *op) {
-    auto reject = [&](StringRef message) {
-      op->emitOpError(message);
-      invalid = true;
-    };
-    if (invalid || isa<SpatYieldOp, scf::YieldOp>(op))
-      return;
-    if (auto selection = dyn_cast<scf::IndexSwitchOp>(op)) {
-      if (failed(verifyCanonicalSourceSelector(
-            selection, deferred, scheduledLane, laneCount)))
-        invalid = true;
-      return;
-    }
-    if (auto loop = dyn_cast<scf::ForOp>(op)) {
-      for (Value bound : {loop.getLowerBound(), loop.getUpperBound(),
-                          loop.getStep()})
-        if (!isAllowedStaticIndexExpression(bound, scheduledLane))
-          return reject("has a non-static deferred loop bound");
-      for (int64_t lane = 0; lane < laneCount; ++lane) {
-        StaticIndexEnvironment environment;
-        if (scheduledLane)
-          environment.bindings[scheduledLane] = lane;
-        auto lower = evaluateDeferredIndex(loop.getLowerBound(), environment);
-        auto upper = evaluateDeferredIndex(loop.getUpperBound(), environment);
-        auto step = evaluateDeferredIndex(loop.getStep(), environment);
-        if (failed(lower) || failed(upper) || failed(step) || *step <= 0)
-          return reject("has a loop bound that does not specialize");
-      }
-      return;
-    }
-    if (op->getNumRegions() != 0)
-      return reject("contains an unsupported region operation");
-    if (!isShapingOnlyOp(op) && !isCompileTimeOp(op)
-        && !isPureIndexComputationOp(op))
-      return reject("contains an unsupported deferred operation");
-    for (Value operand : op->getOperands())
-      if (isInsideDeferredLoop(op, deferred)
-          && originatesFromDeferredSource(operand, deferred))
-        return reject("projects a deferred source inside a residual loop");
-  });
-  return success(!invalid);
-}
-
 FailureOr<int64_t> evaluateDeferredIndex(
     Value value, const StaticIndexEnvironment &environment) {
   llvm::SmallDenseSet<Value, 16> visiting;
@@ -421,18 +362,11 @@ FailureOr<int64_t> evaluateDeferredIndex(
   return failure();
 }
 
-FailureOr<SmallVector<unsigned>> getPossibleDeferredSourceOperandIndices(
-    Value sourceRoot, SpatDeferredCommunicationOp deferred) {
-  llvm::SmallSet<unsigned, 4> indices;
-  collectSourceArguments(sourceRoot, deferred, indices);
-  if (indices.empty())
-    return failure();
-  return SmallVector<unsigned>(indices.begin(), indices.end());
-}
-
 DeferredLaneValueEvaluator::DeferredLaneValueEvaluator(
-    const DeferredProgramTemplate &program, unsigned laneCount)
-  : program(program), laneCount(laneCount) {}
+    const DeferredProgramTemplate &program, unsigned laneCount,
+    unsigned specializationIndex)
+  : program(program), laneCount(laneCount),
+    specializationIndex(specializationIndex) {}
 
 FailureOr<StaticIntSequence> DeferredLaneValueEvaluator::evaluate(Value value) {
   if (auto it = values.find(value); it != values.end())
@@ -448,6 +382,9 @@ FailureOr<StaticIntSequence> DeferredLaneValueEvaluator::evaluate(Value value) {
     StaticIndexEnvironment environment;
     if (program.scheduledLane)
       environment.bindings[program.scheduledLane] = lane;
+    if (program.specializationArgument)
+      environment.bindings[program.specializationArgument] =
+        specializationIndex;
     auto result = evaluateDeferredIndex(value, environment);
     if (failed(result))
       return failure();
@@ -483,6 +420,9 @@ DeferredLaneValueEvaluator::resolveSourceOperandIndices(Value sourceRoot) {
     StaticIndexEnvironment environment;
     if (program.scheduledLane)
       environment.bindings[program.scheduledLane] = lane;
+    if (program.specializationArgument)
+      environment.bindings[program.specializationArgument] =
+        specializationIndex;
     auto index = sourceArgument(sourceRoot, program.deferred, environment);
     if (failed(index) || !*index)
       return failure();
@@ -495,6 +435,9 @@ DeferredLaneValueEvaluator::resolveSourceOperandIndices(Value sourceRoot) {
 
 FailureOr<DeferredProgramTemplate> analyzeDeferredProgramTemplate(
     SpatDeferredCommunicationOp deferred) {
+  if (!deferred.getBody().hasOneBlock())
+    return deferred.emitOpError(
+      "deferred program must have exactly one body block"), failure();
   Block &body = deferred.getBody().front();
   auto yield = dyn_cast<SpatYieldOp>(body.getTerminator());
   if (!yield || yield.getOutputs().size() != 1)
@@ -503,8 +446,29 @@ FailureOr<DeferredProgramTemplate> analyzeDeferredProgramTemplate(
   DeferredProgramTemplate program;
   program.deferred = deferred;
   program.yieldedValue = yield.getOutputs().front();
+  auto specialization = deferred->getAttrOfType<IntegerAttr>(
+    "specialization_count");
+  program.specializationCount = specialization
+                              ? specialization.getInt()
+                              : 1;
+  if (program.specializationCount > 1) {
+    program.specializationArgument = body.getArguments().back();
+    program.specializationFragmentType = dyn_cast<RankedTensorType>(
+      program.yieldedValue.getType());
+    if (!program.specializationFragmentType)
+      return deferred.emitOpError(
+        "grouped deferred fragment must be a ranked tensor"), failure();
+  }
   if (auto scheduled = deferred->getParentOfType<SpatScheduledComputeBatch>())
     program.scheduledLane = getEnclosingScheduledLane(deferred, scheduled);
+  int64_t laneCount = 1;
+  if (auto scheduled = deferred->getParentOfType<SpatScheduledComputeBatch>())
+    laneCount = scheduled.getLaneCount();
+  if (failed(validateDeferredProgram(
+        deferred, body, program.scheduledLane,
+        program.specializationArgument, laneCount,
+        program.specializationCount)))
+    return failure();
 
   llvm::SmallDenseSet<Value, 32> visited;
   std::function<LogicalResult(Value)> visit = [&](Value value) -> LogicalResult {
@@ -529,7 +493,6 @@ FailureOr<DeferredProgramTemplate> analyzeDeferredProgramTemplate(
                                     : DeferredLeafForm::ScalarProjection;
         leaf.sourceRoot = slice.getSource();
         leaf.replacementRoot = value;
-        leaf.leadingProjection = slice;
         leaf.leadingGeometry = {
           SmallVector<OpFoldResult>(slice.getMixedOffsets()),
           SmallVector<OpFoldResult>(slice.getMixedSizes()),
@@ -543,6 +506,19 @@ FailureOr<DeferredProgramTemplate> analyzeDeferredProgramTemplate(
             SmallVector<OpFoldResult>(
               ArrayRef(slice.getMixedStrides()).drop_front())};
         leaf.reconstructedType = cast<RankedTensorType>(value.getType());
+        if (graphProjection
+            && slice.getSourceType().getRank()
+                 == leaf.reconstructedType.getRank() + 1
+            && slice.getMixedSizes().size()
+                 == static_cast<size_t>(slice.getSourceType().getRank())) {
+          ArrayRef<OpFoldResult> innerSizes =
+            ArrayRef(slice.getMixedSizes()).drop_front();
+          leaf.leadingRankReduced = llvm::equal(
+            innerSizes, leaf.reconstructedType.getShape(),
+            [](OpFoldResult size, int64_t dimension) {
+              return getConstantIntValue(size) == dimension;
+            });
+        }
         program.leaves.push_back(std::move(leaf));
         return success();
       }
@@ -553,7 +529,7 @@ FailureOr<DeferredProgramTemplate> analyzeDeferredProgramTemplate(
         return deferred.emitOpError(
           "deferred source is not a ranked tensor");
       program.leaves.push_back({DeferredLeafForm::DirectSource, value, value,
-                                {}, {}, {}, type});
+                                {}, {}, type});
       return success();
     }
     if (value.getType().isIndex() || isa<IntegerType>(value.getType()))
@@ -586,7 +562,7 @@ FailureOr<DeferredProgramTemplate> analyzeDeferredProgramTemplate(
 FailureOr<const GraphBatchPublicationMap *> getGraphBatchPublicationMap(
     SpatGraphComputeBatch graphBatch, unsigned resultIndex,
     GraphBatchPublicationCache &cache) {
-  GraphBatchPublicationKey key {graphBatch.getOperation(), resultIndex};
+  std::pair<Operation *, unsigned> key {graphBatch.getOperation(), resultIndex};
   if (auto it = cache.find(key); it != cache.end())
     return &it->second;
   auto resultType = dyn_cast<RankedTensorType>(

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

@@ -15,16 +15,14 @@ mlir::FailureOr<int64_t> evaluateDeferredIndex(
 mlir::FailureOr<int64_t> evaluateDeferredIndex(
   mlir::OpFoldResult value, const StaticIndexEnvironment &environment);
 
-mlir::LogicalResult verifyDeferredProgramContract(
-  SpatDeferredCommunicationOp deferred);
-
 mlir::FailureOr<DeferredProgramTemplate> analyzeDeferredProgramTemplate(
   SpatDeferredCommunicationOp deferred);
 
 class DeferredLaneValueEvaluator {
 public:
   DeferredLaneValueEvaluator(const DeferredProgramTemplate &program,
-                             unsigned laneCount);
+                             unsigned laneCount,
+                             unsigned specializationIndex = 0);
 
   mlir::FailureOr<StaticIntSequence> evaluate(mlir::Value value);
   mlir::FailureOr<StaticIntSequence> evaluate(mlir::OpFoldResult value);
@@ -34,14 +32,11 @@ public:
 private:
   const DeferredProgramTemplate &program;
   unsigned laneCount;
+  unsigned specializationIndex;
   llvm::DenseMap<mlir::Value, StaticIntSequence> values;
   llvm::DenseMap<mlir::Value, StaticIntSequence> sourceOperands;
 };
 
-mlir::FailureOr<llvm::SmallVector<unsigned>>
-getPossibleDeferredSourceOperandIndices(
-  mlir::Value sourceRoot, SpatDeferredCommunicationOp deferred);
-
 struct GraphBatchPublicationMap {
   mlir::RankedTensorType physicalResultType;
   mlir::RankedTensorType publicationFragmentType;
@@ -49,39 +44,12 @@ struct GraphBatchPublicationMap {
   llvm::SmallVector<int64_t> physicalSlotToGraphLane;
 };
 
-struct GraphBatchPublicationKey {
-  mlir::Operation *graphBatch = nullptr;
-  unsigned resultIndex = 0;
-  bool operator==(const GraphBatchPublicationKey &other) const {
-    return graphBatch == other.graphBatch && resultIndex == other.resultIndex;
-  }
-};
-
 using GraphBatchPublicationCache =
-  llvm::DenseMap<GraphBatchPublicationKey, GraphBatchPublicationMap>;
+  llvm::DenseMap<std::pair<mlir::Operation *, unsigned>,
+                 GraphBatchPublicationMap>;
 
 mlir::FailureOr<const GraphBatchPublicationMap *> getGraphBatchPublicationMap(
   SpatGraphComputeBatch graphBatch, unsigned resultIndex,
   GraphBatchPublicationCache &cache);
 
 } // namespace onnx_mlir::spatial
-
-namespace llvm {
-template <> struct DenseMapInfo<onnx_mlir::spatial::GraphBatchPublicationKey> {
-  static onnx_mlir::spatial::GraphBatchPublicationKey getEmptyKey() {
-    return {DenseMapInfo<mlir::Operation *>::getEmptyKey(), 0};
-  }
-  static onnx_mlir::spatial::GraphBatchPublicationKey getTombstoneKey() {
-    return {DenseMapInfo<mlir::Operation *>::getTombstoneKey(), 0};
-  }
-  static unsigned getHashValue(
-      const onnx_mlir::spatial::GraphBatchPublicationKey &key) {
-    return hash_combine(key.graphBatch, key.resultIndex);
-  }
-  static bool isEqual(
-      const onnx_mlir::spatial::GraphBatchPublicationKey &lhs,
-      const onnx_mlir::spatial::GraphBatchPublicationKey &rhs) {
-    return lhs == rhs;
-  }
-};
-} // namespace llvm

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

@@ -1,318 +1,338 @@
-#include "mlir/Dialect/Arith/IR/Arith.h"
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/IR/IRMapping.h"
-
 #include "DeferredResultRealization.hpp"
 #include "DeferredBoundaryRealization.hpp"
 #include "DeferredProjectionAnalysis.hpp"
-#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
-#include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/IRMapping.h"
+#include "llvm/ADT/DenseSet.h"
+#include <array>
+#include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
-
 namespace onnx_mlir::spatial {
-using namespace mlir;
-namespace {
-
-static Value getSequencePosition(
-  const LaneSet& familyLanes, Value lane, Operation* anchor, DeferredEmissionContext& context, Location loc) {
-  unsigned begin = familyLanes.intervals().front().begin;
-  if (!lane)
-    return context.constants.getIndex(0);
-  return affineAddConst(context.rewriter, loc, lane, -static_cast<int64_t>(begin), anchor);
-}
-
-static OpFoldResult materializeGeometryValue(
-  const StaticIntSequence& sequence, Value lane, Operation* anchor, DeferredEmissionContext& context, Location loc) {
-  if (sequence.getKind() == StaticIntSequenceKind::Uniform)
-    return context.rewriter.getIndexAttr(sequence.valueAt(0));
-  return emitStaticIntLookup(sequence, lane, anchor, context.constants, context.rewriter, loc);
-}
-
-static MixedSliceGeometry materializeGeometry(const DeferredResultPlan::SliceGeometry& geometry,
-                                              Value lane,
-                                              Operation* anchor,
-                                              DeferredEmissionContext& context,
-                                              Location loc) {
-  MixedSliceGeometry result;
-  for (const StaticIntSequence& value : geometry.offsets)
-    result.offsets.push_back(materializeGeometryValue(value, lane, anchor, context, loc));
-  for (const StaticIntSequence& value : geometry.sizes)
-    result.sizes.push_back(materializeGeometryValue(value, lane, anchor, context, loc));
-  for (const StaticIntSequence& value : geometry.strides)
-    result.strides.push_back(materializeGeometryValue(value, lane, anchor, context, loc));
-  return result;
-}
-
-static RequirementFamily*
-findRequirement(const DeferredResultPlan& plan, RequirementCoordinate coordinate, unsigned representativeLane) {
-  RequirementFamily* match = nullptr;
-  for (RequirementFamily* requirement : plan.requirements)
-    if (requirement->coordinate == coordinate && requirement->targetLanes.contains(representativeLane)) {
-      assert(!match && "requirement coordinate is not unique");
-      match = requirement;
-    }
-  return match;
-}
-
-static bool isIdentityGeometry(const DeferredResultPlan::SliceGeometry& geometry, RankedTensorType type) {
-  if (geometry.offsets.size() != static_cast<size_t>(type.getRank()))
-    return false;
-  for (auto [dimension, values] :
-       llvm::enumerate(llvm::zip_equal(geometry.offsets, geometry.sizes, geometry.strides))) {
-    const StaticIntSequence& offset = std::get<0>(values);
-    const StaticIntSequence& size = std::get<1>(values);
-    const StaticIntSequence& stride = std::get<2>(values);
-    if (offset.getKind() != StaticIntSequenceKind::Uniform || size.getKind() != StaticIntSequenceKind::Uniform
-        || stride.getKind() != StaticIntSequenceKind::Uniform || offset.valueAt(0) != 0
-        || size.valueAt(0) != type.getDimSize(dimension) || stride.valueAt(0) != 1)
-      return false;
-  }
-  return true;
-}
-
-static FailureOr<Value> applyInnerGeometry(Value fragment,
-                                           const DeferredProjectionLeafTemplate& leaf,
-                                           const DeferredResultPlan::SliceGeometry& geometry,
-                                           Value lane,
-                                           DeferredEmissionContext& context) {
-  if (leaf.form != DeferredLeafForm::GraphBatchProjection)
-    return fragment;
-  auto fragmentType = dyn_cast<RankedTensorType>(fragment.getType());
-  if (!fragmentType || geometry.offsets.size() != static_cast<size_t>(fragmentType.getRank()))
-    return failure();
-  if (isIdentityGeometry(geometry, fragmentType))
-    return fragment;
-  SmallVector<int64_t> shape(leaf.reconstructedType.getShape().drop_front());
-  auto resultType = RankedTensorType::get(shape, fragmentType.getElementType());
-  Value sourceRoot = leaf.sourceRoot;
-  MixedSliceGeometry mixed =
-    materializeGeometry(geometry, lane, sourceRoot.getParentBlock()->getParentOp(), context, sourceRoot.getLoc());
-  return extractMixedSliceOrIdentity(context.rewriter, leaf.sourceRoot.getLoc(), fragment, resultType, mixed);
-}
-
-static FailureOr<Value> reconstructLeaf(const DeferredResultPlan& plan,
-                                        unsigned leafIndex,
-                                        const LaneSet& activeLanes,
-                                        Value lane,
-                                        DeferredEmissionContext& context) {
-  DeferredExchangePlan& exchange = *plan.exchange;
-  const DeferredProjectionLeafTemplate& leaf = exchange.program.leaves[leafIndex];
-  unsigned representative = activeLanes.intervals().front().begin;
-  if (Value assembled = context.projectionAssemblies.lookup({&exchange, leafIndex})) {
-    for (RequirementFamily* requirement : plan.requirements) {
-      if (requirement->coordinate.leafIndex != leafIndex || !requirement->targetLanes.contains(representative))
-        continue;
-      Value local = context.receives.lookup(requirement);
-      if (!local)
-        continue;
-      auto shaped = applyInnerGeometry(local, leaf, plan.innerGeometry[leafIndex], lane, context);
-      auto source = succeeded(shaped)
-                    ? addLeadingUnitTensorDimension(context.rewriter, exchange.deferred.getLoc(), *shaped)
-                    : FailureOr<Value>(failure());
-      if (failed(source))
-        return failure();
-      auto sourceType = cast<RankedTensorType>(source->getType());
-      MixedSliceGeometry geometry;
-      geometry.offsets.assign(sourceType.getRank(), context.rewriter.getIndexAttr(0));
-      geometry.offsets.front() = context.rewriter.getIndexAttr(requirement->coordinate.selectedPosition);
-      for (int64_t dimension : sourceType.getShape())
-        geometry.sizes.push_back(context.rewriter.getIndexAttr(dimension));
-      geometry.strides.assign(sourceType.getRank(), context.rewriter.getIndexAttr(1));
-      assembled = insertMixedSlice(context.rewriter, exchange.deferred.getLoc(), *source, assembled, geometry);
-    }
-    return assembled;
-  }
-  unsigned positionCount = 0;
-  for (RequirementFamily* requirement : plan.requirements)
-    if (requirement->coordinate.leafIndex == leafIndex && requirement->targetLanes.contains(representative))
-      positionCount = std::max(positionCount, requirement->coordinate.selectedPosition + 1);
-  if (positionCount == 0)
-    return failure();
-  SmallVector<Value> expanded;
-  for (unsigned position = 0; position < positionCount; ++position) {
-    RequirementFamily* requirement = findRequirement(plan, {leafIndex, position}, representative);
-    if (!requirement)
-      return failure();
-    auto fragment = materializeDeferredRequirement(*requirement, activeLanes, lane, context);
-    if (failed(fragment))
-      return failure();
-    auto shaped = applyInnerGeometry(*fragment, leaf, plan.innerGeometry[leafIndex], lane, context);
-    if (failed(shaped))
-      return failure();
-    if (positionCount == 1 && shaped->getType() == leaf.reconstructedType)
-      return *shaped;
-    auto value = addLeadingUnitTensorDimension(context.rewriter, exchange.deferred.getLoc(), *shaped);
-    if (failed(value))
-      return failure();
-    expanded.push_back(*value);
-  }
-  if (expanded.size() == 1 && expanded.front().getType() == leaf.reconstructedType)
-    return expanded.front();
-  constexpr size_t maxConcatInputs = 64;
-  while (expanded.size() > 1) {
-    SmallVector<Value> next;
-    next.reserve((expanded.size() + maxConcatInputs - 1) / maxConcatInputs);
-    for (size_t index = 0; index < expanded.size(); index += maxConcatInputs) {
-      ValueRange inputs = ValueRange(expanded).slice(index, std::min(maxConcatInputs, expanded.size() - index));
-      if (inputs.size() == 1) {
-        next.push_back(inputs.front());
-        continue;
-      }
-      auto first = cast<RankedTensorType>(inputs.front().getType());
-      SmallVector<int64_t> shape(first.getShape());
-      shape.front() = 0;
-      for (Value input : inputs)
-        shape.front() += cast<RankedTensorType>(input.getType()).getDimSize(0);
-      auto resultType = RankedTensorType::get(shape, first.getElementType());
-      next.push_back(
-        tensor::ConcatOp::create(context.rewriter, exchange.deferred.getLoc(), resultType, 0, inputs).getResult());
-    }
-    expanded = std::move(next);
-  }
-  return expanded.front().getType() == leaf.reconstructedType ? FailureOr<Value>(expanded.front())
-                                                              : FailureOr<Value>(failure());
-}
-
-} // namespace
-
-FailureOr<DeferredResultPlan>
-buildDeferredResultPlan(DeferredExchangePlan& exchange) {
-  DeferredResultPlan result;
-  result.exchange = &exchange;
-  for (RequirementFamily& requirement : exchange.requirements)
-    result.requirements.push_back(&requirement);
-  DeferredLaneValueEvaluator evaluator(
-    exchange.program, exchange.targetLaneCount);
-  auto buildGeometry = [&](const DeferredSliceTemplate& source,
-                           DeferredResultPlan::SliceGeometry& target) {
-    auto append = [&](ArrayRef<OpFoldResult> values,
-                      SmallVectorImpl<StaticIntSequence>& sequences) {
-      for (OpFoldResult value : values) {
-        auto sequence = evaluator.evaluate(value);
-        if (failed(sequence))
-          return failure();
-        sequences.push_back(*sequence);
-      }
-      return success();
-    };
-    return success(succeeded(append(source.offsets, target.offsets))
-                   && succeeded(append(source.sizes, target.sizes))
-                   && succeeded(append(source.strides, target.strides)));
-  };
-  for (const DeferredProjectionLeafTemplate& leaf : exchange.program.leaves) {
-    DeferredResultPlan::SliceGeometry geometry;
-    if (failed(buildGeometry(leaf.innerGeometry, geometry)))
-      return failure();
-    result.innerGeometry.push_back(std::move(geometry));
-  }
-  if (exchange.program.insertAssembly)
-    for (const DeferredInsertAssemblyEntryTemplate& entry :
-         exchange.program.insertAssembly->entries) {
-      DeferredResultPlan::SliceGeometry geometry;
-      if (failed(buildGeometry(entry.targetGeometry, geometry)))
-        return failure();
-      result.assemblyGeometry.push_back(std::move(geometry));
-    }
-  for (Operation* op : exchange.program.residualOps)
-    op->walk([&](Operation* nested) {
-      for (Value operand : nested->getOperands()) {
-        if ((!operand.getType().isIndex()
-             && !isa<IntegerType>(operand.getType()))
-            || result.residualValues.count(operand))
-          continue;
-        if (auto sequence = evaluator.evaluate(operand); succeeded(sequence))
-          result.residualValues.try_emplace(operand, *sequence);
-      }
-    });
-  return result;
-}
-
-FailureOr<Value> materializeDeferredRequirement(RequirementFamily& requirement,
-                                                const LaneSet& activeLanes,
-                                                Value lane,
-                                                DeferredEmissionContext& context) {
-  if (Value received = context.receives.lookup(&requirement))
-    return received;
-  ProducedValue& producer = *requirement.producer;
-  Value payload = producer.payload;
-  Location loc = requirement.exchange->deferred.getLoc();
-  Value position = getSequencePosition(
-    requirement.targetLanes, lane, requirement.exchange->deferred, context, loc);
-  if (requirement.producerProjection) {
-    auto fragmentType = dyn_cast<RankedTensorType>(
-      requirement.publicationFragmentType);
-    if (!fragmentType)
-      return failure();
-    MixedSliceGeometry geometry = materializeGeometry(
-      *requirement.producerProjection, position,
-      requirement.exchange->deferred, context, loc);
-    return extractMixedSliceOrIdentity(
-      context.rewriter, loc, payload, fragmentType, geometry);
-  }
-  if (payload.getType() == requirement.publicationFragmentType)
-    return payload;
-  auto payloadType = dyn_cast<RankedTensorType>(payload.getType());
-  auto fragmentType = dyn_cast<RankedTensorType>(requirement.publicationFragmentType);
-  if (!payloadType || !fragmentType || !requirement.producerLocalOffsets
-      || payloadType.getRank() != fragmentType.getRank() + 1)
-    return failure();
-  Value offset = emitStaticIntLookup(*requirement.producerLocalOffsets,
-                                     position,
-                                     requirement.exchange->deferred,
-                                     context.constants,
-                                     context.rewriter,
-                                     loc);
+using namespace mlir; namespace {
+static MixedSliceGeometry leadingSlice(OpBuilder &builder, RankedTensorType type, OpFoldResult position) {
   MixedSliceGeometry geometry;
-  geometry.offsets.assign(payloadType.getRank(), context.rewriter.getIndexAttr(0));
-  geometry.sizes.push_back(context.rewriter.getIndexAttr(1));
-  geometry.strides.assign(payloadType.getRank(), context.rewriter.getIndexAttr(1));
-  geometry.offsets.front() = offset;
-  for (int64_t dimension : fragmentType.getShape())
-    geometry.sizes.push_back(context.rewriter.getIndexAttr(dimension));
-  SmallVector<int64_t> unitShape {1};
-  llvm::append_range(unitShape, fragmentType.getShape());
-  auto unitType = RankedTensorType::get(unitShape, fragmentType.getElementType());
-  Value unit = extractMixedSliceOrIdentity(context.rewriter, loc, payload, unitType, geometry);
-  return removeLeadingUnitTensorDimension(context.rewriter, loc, unit, fragmentType);
+  geometry.offsets.assign(type.getRank() + 1, builder.getIndexAttr(0));
+  geometry.offsets.front() = position;
+  geometry.sizes.push_back(builder.getIndexAttr(1));
+  for (int64_t dimension : type.getShape()) geometry.sizes.push_back(builder.getIndexAttr(dimension));
+  geometry.strides.assign(type.getRank() + 1, builder.getIndexAttr(1));
+  return geometry;
+}
+static LogicalResult verifyCollectionCoverage(
+    DeferredExchangePlan &exchange, FragmentCollectionPlan &collection,
+    unsigned slotCount) {
+  LaneSet all = LaneSet::all(exchange.targetLaneCount);
+  SmallVector<LaneSet, 0> coverage(slotCount);
+  for (const FragmentCollectionPlan::Requirement &entry :
+       collection.requirements) {
+    if (entry.position >= coverage.size()
+        || !coverage[entry.position]
+              .intersect(entry.family->targetLanes).empty())
+      return exchange.deferred.emitOpError(
+               "fragment collection has overlapping coverage at slot ")
+             << entry.position;
+    coverage[entry.position] =
+      coverage[entry.position].unite(entry.family->targetLanes);
+  }
+  for (auto [position, lanes] : llvm::enumerate(coverage))
+    if (!(lanes == all))
+      return exchange.deferred.emitOpError(
+               "fragment collection has incomplete coverage at slot ")
+             << position;
+  return success();
 }
 
-FailureOr<Value> realizeDeferredResult(const DeferredResultPlan& plan,
-                                       const LaneSet& activeLanes,
-                                       Value lane,
-                                       DeferredEmissionContext& context) {
-  DeferredExchangePlan& exchange = *plan.exchange;
-  if (Value assembled = context.assemblies.lookup(&exchange))
-    return assembled;
+static LogicalResult buildLeafCollections(DeferredExchangePlan &exchange,
+                                          DeferredResultPlan &plan) {
+  unsigned specializationCount = exchange.program.specializationCount;
+  for (auto [leafIndex, leaf] : llvm::enumerate(exchange.program.leaves)) {
+    SmallVector<RequirementFamily *> requirements;
+    unsigned positionCount = 0;
+    for (RequirementFamily &requirement : exchange.requirements)
+      if (requirement.coordinate.leafIndex == leafIndex) {
+        requirements.push_back(&requirement);
+        positionCount = std::max(
+          positionCount, requirement.coordinate.selectedPosition + 1);
+      }
+    if (requirements.empty())
+      return exchange.deferred.emitOpError(
+               "cannot form fragment collection for leaf ")
+             << leafIndex << ": no requirements";
+    auto fragmentType = dyn_cast<RankedTensorType>(
+      requirements.front()->publicationFragmentType);
+    if (!fragmentType || llvm::any_of(requirements,
+          [&](RequirementFamily *item) {
+            return item->publicationFragmentType != fragmentType;
+          }))
+      return exchange.deferred.emitOpError(
+               "cannot form fragment collection for leaf ")
+             << leafIndex
+             << ": publication fragment types differ or are unranked";
+    RankedTensorType normalized = fragmentType;
+    if (leaf.form == DeferredLeafForm::GraphBatchProjection)
+      normalized = leaf.leadingRankReduced
+        ? leaf.reconstructedType
+        : RankedTensorType::get(
+            leaf.reconstructedType.getShape().drop_front(),
+            leaf.reconstructedType.getElementType());
+    bool direct = positionCount == 1 && normalized == leaf.reconstructedType;
+    bool leading = leaf.reconstructedType.getRank() == normalized.getRank() + 1
+                && leaf.reconstructedType.getDimSize(0) == positionCount
+                && leaf.reconstructedType.getShape().drop_front()
+                     == normalized.getShape();
+    if (!direct && !leading)
+      return exchange.deferred.emitOpError(
+               "cannot form fragment collection for leaf ")
+             << leafIndex << ": " << positionCount << " fragments of "
+             << normalized << " do not reconstruct " << leaf.reconstructedType;
+    FragmentCollectionKind kind = specializationCount == 1
+                                ? FragmentCollectionKind::Leaf
+                                : FragmentCollectionKind::GroupedLeaf;
+    SmallVector<int64_t> shape;
+    if (specializationCount > 1)
+      shape.push_back(specializationCount);
+    llvm::append_range(shape, leaf.reconstructedType.getShape());
+    FragmentCollectionPlan collection;
+    collection.key = {&exchange, kind, static_cast<unsigned>(leafIndex)};
+    collection.collectionType = RankedTensorType::get(
+      shape, leaf.reconstructedType.getElementType());
+    collection.positionCount = positionCount;
+    for (RequirementFamily *requirement : requirements)
+      collection.requirements.push_back({
+        requirement,
+        requirement->coordinate.specializationIndex * positionCount
+          + requirement->coordinate.selectedPosition});
+    if (failed(verifyCollectionCoverage(
+          exchange, collection, specializationCount * positionCount)))
+      return failure();
+    plan.collections.push_back(std::move(collection));
+  }
+  return success();
+}
+
+static bool supportsAssemblyTransform(
+    Type publicationType, const DeferredInsertAssemblyEntryTemplate &entry) {
+  auto publication = dyn_cast<RankedTensorType>(publicationType);
+  RankedTensorType source = entry.sourceType;
+  if (!publication || !source)
+    return false;
+  switch (entry.sourceTransform) {
+  case DeferredAssemblySourceTransform::Identity:
+    return publication == source;
+  case DeferredAssemblySourceTransform::AddLeadingUnitDimension:
+    return source.getRank() == publication.getRank() + 1
+        && source.getDimSize(0) == 1
+        && source.getShape().drop_front() == publication.getShape();
+  case DeferredAssemblySourceTransform::RemoveLeadingUnitDimension:
+    return publication.getRank() == source.getRank() + 1
+        && publication.getDimSize(0) == 1
+        && publication.getShape().drop_front() == source.getShape();
+  }
+  llvm_unreachable("unknown deferred assembly source transform");
+}
+
+static LogicalResult buildInsertAssemblyCollection(
+    DeferredExchangePlan &exchange, DeferredResultPlan &plan) {
+  if (exchange.program.specializationCount != 1)
+    return exchange.deferred.emitOpError(
+      "grouped specialization insert assembly is unsupported");
+  const DeferredInsertAssemblyTemplate &assembly =
+    *exchange.program.insertAssembly;
+  FragmentCollectionPlan collection;
+  collection.key = {&exchange, FragmentCollectionKind::InsertAssembly, 0};
+  collection.collectionType = assembly.resultType;
+  collection.positionCount = assembly.entries.size();
+  llvm::DenseSet<RequirementFamily *> collected;
+  for (auto [position, entry] : llvm::enumerate(assembly.entries)) {
+    bool found = false;
+    for (RequirementFamily &requirement : exchange.requirements) {
+      if (!(requirement.coordinate == entry.coordinate))
+        continue;
+      if (!supportsAssemblyTransform(
+            requirement.publicationFragmentType, entry))
+        return exchange.deferred.emitOpError(
+                 "insert assembly source transform does not match publication type at entry ")
+               << position;
+      if (!collected.insert(&requirement).second)
+        return exchange.deferred.emitOpError(
+                 "insert assembly requirement is owned by multiple entries at entry ")
+               << position;
+      collection.requirements.push_back(
+        {&requirement, static_cast<unsigned>(position)});
+      found = true;
+    }
+    if (!found)
+      return exchange.deferred.emitOpError(
+               "insert assembly entry has no requirement at entry ")
+             << position;
+  }
+  if (collected.size() != exchange.requirements.size())
+    return exchange.deferred.emitOpError(
+      "insert assembly does not own every deferred requirement");
+  if (failed(verifyCollectionCoverage(
+        exchange, collection, assembly.entries.size())))
+    return failure();
+  plan.collections.push_back(std::move(collection));
+  return success();
+}
+
+using TemplateGeometryMember = SmallVector<OpFoldResult> DeferredSliceTemplate::*;
+static constexpr std::array<TemplateGeometryMember, 3> templateGeometryMembers{
+  &DeferredSliceTemplate::offsets, &DeferredSliceTemplate::sizes,
+  &DeferredSliceTemplate::strides};
+
+template <typename GetValue>
+static FailureOr<StaticIntGrid> buildGrid(
+    const DeferredProgramTemplate &program, unsigned laneCount,
+    unsigned rowCount, bool specializeRows, GetValue getValue) {
+  SmallVector<StaticIntSequence> rows;
+  for (unsigned row = 0; row < rowCount; ++row) {
+    DeferredLaneValueEvaluator evaluator(
+      program, laneCount, specializeRows ? row : 0);
+    auto sequence = evaluator.evaluate(getValue(row));
+    if (failed(sequence)) return failure();
+    rows.push_back(std::move(*sequence));
+  }
+  return StaticIntGrid::fromRows(rows);
+}
+
+template <typename GetGeometry>
+static FailureOr<DeferredGridSliceGeometry> buildGeometryGrids(
+    const DeferredProgramTemplate &program, unsigned laneCount,
+    unsigned rowCount, bool specializeRows, GetGeometry getGeometry) {
+  DeferredGridSliceGeometry result;
+  const DeferredSliceTemplate &first = getGeometry(0);
+  for (auto [group, sourceMember] : llvm::enumerate(templateGeometryMembers)) {
+    TemplateGeometryMember member = sourceMember;
+    for (unsigned dimension = 0;
+         dimension < (first.*member).size(); ++dimension) {
+      auto grid = buildGrid(program, laneCount, rowCount, specializeRows,
+        [&](unsigned row) { return (getGeometry(row).*member)[dimension]; });
+      if (failed(grid)) return failure();
+      result[group].push_back(std::move(*grid));
+    }
+  }
+  return result;
+}
+static Value cloneResidual(DeferredExchangePlan &exchange, IRMapping &mapping, DeferredEmissionContext &context) {
+  for (Operation *op : exchange.program.residualOps) {
+    Operation *copy = context.rewriter.clone(*op, mapping);
+    for (auto [oldValue, newValue] : llvm::zip(op->getResults(), copy->getResults())) mapping.map(oldValue, newValue);
+  }
+  return mapping.lookupOrDefault(exchange.program.yieldedValue);
+}
+static FailureOr<Value> realizeOne(const DeferredResultPlan &plan, Value lane, DeferredEmissionContext &context) {
+  DeferredExchangePlan &exchange = *plan.exchange;
+  if (exchange.program.insertAssembly) {
+    Value collection = context.fragmentCollections.lookup(
+      {&exchange, FragmentCollectionKind::InsertAssembly, 0});
+    return collection && collection.getType() == exchange.program.insertAssembly->resultType
+      ? FailureOr<Value>(collection) : FailureOr<Value>(failure());
+  }
   IRMapping mapping;
-  if (exchange.program.scheduledLane)
-    mapping.map(exchange.program.scheduledLane, lane);
-  for (auto [value, sequence] : plan.residualValues) {
-    if (mapping.contains(value))
-      continue;
-    Value selected = emitStaticIntLookup(sequence,
-                                         lane ? lane : context.constants.getIndex(0),
-                                         exchange.deferred,
-                                         context.constants,
-                                         context.rewriter,
-                                         exchange.deferred.getLoc());
-    if (!value.getType().isIndex())
-      selected = arith::IndexCastOp::create(context.rewriter, exchange.deferred.getLoc(), value.getType(), selected);
+  if (exchange.program.scheduledLane) mapping.map(exchange.program.scheduledLane, lane);
+  for (auto &[value, grid] : plan.residualValues) {
+    if (mapping.contains(value)) continue;
+    Value selected = grid.emitLookup(
+      context.constants.getIndex(0),
+      lane ? lane : context.constants.getIndex(0), exchange.deferred,
+      context.constants, context.rewriter, exchange.deferred.getLoc());
+    if (!value.getType().isIndex()) selected = arith::IndexCastOp::create(context.rewriter, exchange.deferred.getLoc(), value.getType(), selected);
     mapping.map(value, selected);
   }
-  for (unsigned leaf = 0; leaf < exchange.program.leaves.size(); ++leaf) {
-    auto value = reconstructLeaf(plan, leaf, activeLanes, lane, context);
-    if (failed(value))
+  for (unsigned index = 0; index < exchange.program.leaves.size(); ++index) {
+    Value value = context.fragmentCollections.lookup(
+      {&exchange, FragmentCollectionKind::Leaf, index});
+    if (!value || value.getType() != exchange.program.leaves[index].reconstructedType) {
+      exchange.deferred.emitOpError("failed to reconstruct deferred result leaf ") << index;
       return failure();
-    mapping.map(exchange.program.leaves[leaf].replacementRoot, *value);
+    }
+    mapping.map(exchange.program.leaves[index].replacementRoot, value);
   }
-  for (Operation* op : exchange.program.residualOps) {
-    Operation* copy = context.rewriter.clone(*op, mapping);
-    for (auto [oldValue, newValue] : llvm::zip(op->getResults(), copy->getResults()))
-      mapping.map(oldValue, newValue);
-  }
-  Value result = mapping.lookupOrDefault(exchange.program.yieldedValue);
-  return result && result.getType() == exchange.deferred.getOutput().getType() ? FailureOr<Value>(result)
-                                                                               : FailureOr<Value>(failure());
+  Value result = cloneResidual(exchange, mapping, context);
+  Type expected = exchange.program.specializationCount > 1 ? Type(exchange.program.specializationFragmentType) : exchange.deferred.getOutput().getType();
+  return result && result.getType() == expected ? FailureOr<Value>(result) : FailureOr<Value>(failure());
+}
+} // namespace
+FailureOr<DeferredResultPlan> buildDeferredResultPlan(DeferredExchangePlan &exchange) {
+  DeferredResultPlan plan;
+  plan.exchange = &exchange;
+  if (failed(exchange.program.insertAssembly
+      ? buildInsertAssemblyCollection(exchange, plan)
+      : buildLeafCollections(exchange, plan))) return failure();
+  unsigned specializations = exchange.program.specializationCount;
+  for (const auto &leaf : exchange.program.leaves) {
+    auto geometry = buildGeometryGrids(exchange.program,
+      exchange.targetLaneCount, specializations, true,
+      [&](unsigned) -> const DeferredSliceTemplate & {
+        return leaf.innerGeometry;
+      });
+    if (failed(geometry)) return failure();
+    plan.innerGeometry.push_back(std::move(*geometry));
+  }
+  if (exchange.program.insertAssembly) {
+    const auto &entries = exchange.program.insertAssembly->entries;
+    auto geometry = buildGeometryGrids(exchange.program,
+      exchange.targetLaneCount, entries.size(), false,
+      [&](unsigned row) -> const DeferredSliceTemplate & {
+        return entries[row].targetGeometry;
+      });
+    if (failed(geometry)) return failure();
+    plan.assemblyGeometry = std::move(*geometry);
+  }
+  llvm::DenseSet<Value> residualValues;
+  for (Operation *op : exchange.program.residualOps)
+    op->walk([&](Operation *nested) {
+      for (Value operand : nested->getOperands())
+        if (operand.getType().isIndex() || isa<IntegerType>(operand.getType()))
+          residualValues.insert(operand);
+    });
+  for (Value value : residualValues) {
+    auto grid = buildGrid(exchange.program, exchange.targetLaneCount,
+      specializations, true,
+      [&](unsigned) { return OpFoldResult(value); });
+    if (succeeded(grid))
+      plan.residualValues.try_emplace(value, std::move(*grid));
+  }
+  return plan;
+}
+FailureOr<Value> realizeDeferredResult(const DeferredResultPlan &plan, Value lane, DeferredEmissionContext &context) {
+  DeferredExchangePlan &exchange = *plan.exchange;
+  if (exchange.program.specializationCount == 1) return realizeOne(plan, lane, context);
+  auto outputType = dyn_cast<RankedTensorType>(exchange.deferred.getOutput().getType());
+  RankedTensorType fragmentType = exchange.program.specializationFragmentType;
+  if (!outputType || !fragmentType || outputType.getRank() != fragmentType.getRank() + 1 || outputType.getDimSize(0) != exchange.program.specializationCount || outputType.getShape().drop_front() != fragmentType.getShape()) return failure();
+  if (exchange.program.insertAssembly) return failure();
+  SmallVector<Value> leafStacks;
+  for (auto [index, leaf] : llvm::enumerate(exchange.program.leaves)) {
+      FragmentCollectionKey key{&exchange, FragmentCollectionKind::GroupedLeaf, static_cast<unsigned>(index)};
+      Value collection = context.fragmentCollections.lookup(key);
+      SmallVector<int64_t> shape{exchange.program.specializationCount};
+      llvm::append_range(shape, leaf.reconstructedType.getShape());
+      if (!collection || collection.getType() != RankedTensorType::get(shape, leaf.reconstructedType.getElementType())) return failure();
+      leafStacks.push_back(collection);
+  }
+  Location loc = exchange.deferred.getLoc();
+  Value initial = tensor::EmptyOp::create(context.rewriter, loc, outputType.getShape(), outputType.getElementType());
+  auto loop = buildNormalizedScfFor(context.rewriter, loc, context.constants.getIndex(0), context.constants.getIndex(exchange.program.specializationCount), context.constants.getIndex(1), ValueRange{initial},
+                                    [&](OpBuilder &, Location, Value specialization, ValueRange iterArgs, SmallVectorImpl<Value> &yielded) -> LogicalResult {
+                                      IRMapping mapping;
+                                      if (exchange.program.scheduledLane) mapping.map(exchange.program.scheduledLane, lane);
+                                      if (exchange.program.specializationArgument) mapping.map(exchange.program.specializationArgument, specialization);
+                                      Value runtimeLane = lane ? lane : context.constants.getIndex(0);
+                                      for (auto &[value, grid] : plan.residualValues) {
+                                          Value selected = grid.emitLookup(specialization, runtimeLane, exchange.deferred, context.constants, context.rewriter, loc);
+                                          if (!value.getType().isIndex()) selected = arith::IndexCastOp::create(context.rewriter, loc, value.getType(), selected);
+                                          mapping.map(value, selected);
+                                      }
+                                      for (auto [index, leaf] : llvm::enumerate(exchange.program.leaves)) { Value selected = extractMixedSliceOrIdentity(context.rewriter, loc, leafStacks[index], leaf.reconstructedType, leadingSlice(context.rewriter, leaf.reconstructedType, specialization)); if (!selected) return failure(); mapping.map(leaf.replacementRoot, selected); }
+                                      Value fragment = cloneResidual(exchange, mapping, context);
+                                      if (!fragment || fragment.getType() != fragmentType) return failure();
+                                      yielded.push_back(insertMixedSlice(context.rewriter, loc, fragment, iterArgs.front(), leadingSlice(context.rewriter, fragmentType, specialization)));
+                                      return success();
+                                    });
+  return failed(loop) || loop->results.size() != 1 ? FailureOr<Value>(failure()) : FailureOr<Value>(loop->results.front());
 }
-
 } // namespace onnx_mlir::spatial

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

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

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

@@ -10,190 +10,44 @@ namespace onnx_mlir::spatial {
 using namespace mlir;
 namespace {
 
-static FailureOr<SmallVector<int64_t>> getI64Array(Operation* op, StringRef name) {
-  auto attr = op->getAttrOfType<DenseI64ArrayAttr>(name);
-  if (!attr)
-    return op->emitOpError() << "phase 2 requires '" << name << "' metadata";
-  return SmallVector<int64_t>(attr.asArrayRef());
+static FailureOr<unsigned> getStepIndex(
+    ScheduledInfo& info, SpatDeferredCommunicationOp deferred) {
+  Operation *position = deferred;
+  Block *block = info.blocks.front();
+  while (position && position->getBlock() != block)
+    position = position->getParentOp();
+  if (!position)
+    return failure();
+  for (unsigned step : llvm::reverse(
+         llvm::seq<unsigned>(0, info.stepAnchors.size())))
+    if (info.stepAnchors[step] == position
+        || info.stepAnchors[step]->isBeforeInBlock(position))
+      return step;
+  return failure();
 }
 
-static FailureOr<SmallVector<int64_t>> getLaneTable(Operation* op, StringRef name, size_t expected) {
-  if (auto array = op->getAttrOfType<DenseI64ArrayAttr>(name)) {
-    if (array.size() != static_cast<int64_t>(expected))
-      return op->emitOpError() << "phase 2 metadata '" << name << "' has the wrong size";
-    return SmallVector<int64_t>(array.asArrayRef());
-  }
-  auto elements = op->getAttrOfType<DenseIntElementsAttr>(name);
-  if (!elements || elements.getNumElements() != static_cast<int64_t>(expected))
-    return op->emitOpError() << "phase 2 requires a correctly-sized '" << name << "' lane table";
-  SmallVector<int64_t> values;
-  for (APInt value : elements.getValues<APInt>())
-    values.push_back(value.getSExtValue());
-  return values;
-}
-
-static Block* getScheduledBlock(SpatDeferredCommunicationOp deferred, Operation* scheduled) {
-  Block* block = deferred->getBlock();
-  while (block && block->getParentOp() != scheduled) {
-    Operation* parent = block->getParentOp();
-    block = parent ? parent->getBlock() : nullptr;
-  }
-  return block;
-}
-
-static FailureOr<unsigned> getStepIndex(ScheduledInfo& info, Block* block) {
-  auto it = llvm::find(info.blocks, block);
-  return it == info.blocks.end() ? FailureOr<unsigned>(failure())
-                                 : FailureOr<unsigned>(std::distance(info.blocks.begin(), it));
-}
-
-static FailureOr<Value>
-getScalarStepResult(SpatScheduledCompute scheduled, Block& block, unsigned resultIndex, unsigned resultCount) {
-  auto yield = dyn_cast<SpatBlockYieldOp>(block.getTerminator());
-  if (!yield || resultIndex >= resultCount || yield.getOutputs().size() < resultCount)
-    return scheduled.emitOpError("phase 2 cannot recover a scalar scheduled step result"), failure();
-  return yield.getOutputs()[yield.getOutputs().size() - resultCount + resultIndex];
-}
-
-struct BatchPublication {
-  Value payload;
-  tensor::ParallelInsertSliceOp insertion;
-};
-
-struct ProducedValueGeometry {
-  int64_t laneStart = 0;
-  int64_t laneCount = 1;
-  int64_t publishedSlotStart = 0;
-  int64_t publishedSlotCount = 1;
-  Value payload;
-  Value published;
-};
-
-static FailureOr<BatchPublication>
-getBatchStepResult(SpatScheduledComputeBatch scheduled, Block& block, unsigned globalResultIndex) {
-  auto parallel = dyn_cast<SpatInParallelOp>(block.getTerminator());
-  if (!parallel)
-    return scheduled.emitOpError("phase 2 cannot recover a batched scheduled step result"), failure();
-  unsigned resultBase = 1 + scheduled.getWeights().size() + scheduled.getInputs().size();
-  for (Operation& op : parallel.getRegion().front()) {
-    auto insert = dyn_cast<tensor::ParallelInsertSliceOp>(op);
-    auto destination = insert ? dyn_cast<BlockArgument>(insert.getDest()) : BlockArgument();
-    if (destination && destination.getOwner() == &block && destination.getArgNumber() == resultBase + globalResultIndex)
-      return BatchPublication {insert.getSource(), insert};
-  }
-  return scheduled.emitOpError("phase 2 cannot find the batched scheduled result insertion"), failure();
-}
-
-static FailureOr<ProducedValueGeometry> verifyScheduledPublicationGeometry(
-  ScheduledInfo& info, unsigned step, unsigned globalResult, unsigned lane,
-  int64_t laneStart, int64_t laneCount, Value payload,
-  tensor::ParallelInsertSliceOp insertion = {}) {
-  ProducedValueGeometry result;
-  result.laneStart = laneStart;
-  result.laneCount = info.isBatch() ? laneCount : std::max<int64_t>(laneCount, 1);
-  result.payload = payload;
-  result.published = info.op->getResult(globalResult);
-  auto payloadType = dyn_cast<RankedTensorType>(payload.getType());
-  auto publishedType = dyn_cast<RankedTensorType>(result.published.getType());
-  if (laneStart < 0 || result.laneCount <= 0)
-    return info.op->emitOpError("phase 2 scheduled publication has an invalid source-lane range"), failure();
-
-  if (!info.isBatch()) {
-    result.publishedSlotCount = result.laneCount;
-    if (payload.getType() != result.published.getType()
-        || (result.laneCount > 1
-            && (!publishedType || !publishedType.hasStaticShape()
-                || publishedType.getRank() == 0
-                || publishedType.getDimSize(0) != result.laneCount)))
-      return info.op->emitOpError("phase 2 scalar scheduled publication is incompatible with its result"), failure();
-    return result;
-  }
-
-  if (!insertion || !payloadType || !payloadType.hasStaticShape()
-      || !publishedType || !publishedType.hasStaticShape()
-      || insertion.getDest() != info.blocks[step]->getArgument(
-           1 + info.op->getNumOperands() + globalResult))
-    return info.op->emitOpError("phase 2 batched publication is malformed"), failure();
-  unsigned rank = publishedType.getRank();
-  if (rank == 0 || payloadType.getRank() != rank
-      || insertion.getMixedOffsets().size() != rank
-      || insertion.getMixedSizes().size() != rank
-      || insertion.getMixedStrides().size() != rank)
-    return info.op->emitOpError("phase 2 batched publication geometry rank is invalid"), failure();
-
-  StaticIndexEnvironment environment;
-  environment.bindings[info.blocks[step]->getArgument(0)] = lane;
-  SmallVector<int64_t> offsets, sizes, strides;
-  for (OpFoldResult value : insertion.getMixedOffsets()) {
-    auto evaluated = evaluateDeferredIndex(value, environment);
-    if (failed(evaluated))
-      return info.op->emitOpError("phase 2 batched publication offset is not static"), failure();
-    offsets.push_back(*evaluated);
-  }
-  for (OpFoldResult value : insertion.getMixedSizes()) {
-    auto evaluated = evaluateDeferredIndex(value, environment);
-    if (failed(evaluated))
-      return info.op->emitOpError("phase 2 batched publication size is not static"), failure();
-    sizes.push_back(*evaluated);
-  }
-  for (OpFoldResult value : insertion.getMixedStrides()) {
-    auto evaluated = evaluateDeferredIndex(value, environment);
-    if (failed(evaluated))
-      return info.op->emitOpError("phase 2 batched publication stride is not static"), failure();
-    strides.push_back(*evaluated);
-  }
-  if (offsets.front() < 0 || sizes.front() <= 0 || sizes.front() != result.laneCount
-      || strides.front() != 1
-      || offsets.front() + sizes.front() > publishedType.getDimSize(0))
-    return info.op->emitOpError("phase 2 batched publication leading geometry is invalid"), failure();
-  for (unsigned dimension = 1; dimension < rank; ++dimension)
-    if (offsets[dimension] != 0 || strides[dimension] != 1)
-      return info.op->emitOpError("phase 2 batched publication inner geometry is invalid"), failure();
-  for (unsigned dimension = 0; dimension < rank; ++dimension)
-    if (sizes[dimension] != payloadType.getDimSize(dimension))
-      return info.op->emitOpError("phase 2 batched publication payload shape is incompatible"), failure();
-  result.publishedSlotStart = offsets.front();
-  result.publishedSlotCount = sizes.front();
-  return result;
-}
-
-static LogicalResult collectScheduledOperations(func::FuncOp funcOp, DeferredTransferPlan& plan) {
+static LogicalResult collectScheduledOperations(
+    const ScheduledComputeMaterializationResult &materialization,
+    DeferredTransferPlan &plan) {
   unsigned nextStream = 0;
-  for (Operation& op : funcOp.getOps()) {
-    if (!isa<SpatScheduledCompute, SpatScheduledComputeBatch>(op))
-      continue;
+  for (const ScheduledMaterializationRecord &record :
+       materialization.materializedSchedules) {
+    Operation &op = *record.scheduledOp;
     ScheduledInfo info;
     info.op = &op;
     Region& body = isa<SpatScheduledCompute>(op) ? cast<SpatScheduledCompute>(op).getBody()
                                                  : cast<SpatScheduledComputeBatch>(op).getBody();
-    for (Block& block : body) {
-      info.blocks.push_back(&block);
-      info.stepAnchors.push_back(&block.front());
-    }
-    auto sourceIds = getI64Array(&op, "scheduled.step_source_ids");
-    auto offsets = getI64Array(&op, "scheduled.step_result_offsets");
-    auto counts = getI64Array(&op, "scheduled.step_result_counts");
-    if (failed(sourceIds) || failed(offsets) || failed(counts))
-      return failure();
-    info.stepSourceIds = std::move(*sourceIds);
-    info.resultOffsets = std::move(*offsets);
-    info.resultCounts = std::move(*counts);
-    if (info.blocks.size() != info.stepSourceIds.size() || info.blocks.size() != info.resultOffsets.size()
-        || info.blocks.size() != info.resultCounts.size())
-      return op.emitOpError("phase 2 scheduled metadata does not match its block count");
-    if (auto scalar = dyn_cast<SpatScheduledCompute>(op)) {
-      auto core = scalar->getAttrOfType<IntegerAttr>(kCoreIdAttrName);
-      if (!core)
-        return scalar.emitOpError("phase 2 requires scalar coreId metadata");
-      info.cores.push_back(core.getInt());
-    }
-    else {
-      auto cores = op.getAttrOfType<DenseI32ArrayAttr>(kCoreIdsAttrName);
-      if (!cores)
-        return op.emitOpError("phase 2 requires batch coreIds metadata");
-      for (int32_t core : cores.asArrayRef())
-        info.cores.push_back(core);
-    }
+    if (!body.hasOneBlock())
+      return op.emitOpError("phase 2 requires canonical one-block scheduled IR");
+    info.blocks.push_back(&body.front());
+    info.stepCount = record.runs.empty() ? record.stepPlans.size()
+                                        : record.runs.size();
+    info.stepAnchors = record.stepAnchors;
+    if (llvm::any_of(info.stepAnchors,
+                     [](Operation *anchor) { return !anchor; }))
+      return op.emitOpError("phase 2 scheduled step anchor is missing");
+    for (size_t core : record.cpus)
+      info.cores.push_back(core);
     for (size_t lane = 0; lane < info.cores.size(); ++lane)
       info.streamIds.push_back(nextStream++);
     plan.scheduled.push_back(std::move(info));
@@ -201,67 +55,47 @@ static LogicalResult collectScheduledOperations(func::FuncOp funcOp, DeferredTra
   plan.stepCounts.resize(nextStream);
   for (ScheduledInfo& info : plan.scheduled)
     for (unsigned stream : info.streamIds)
-      plan.stepCounts[stream] = info.blocks.size();
+      plan.stepCounts[stream] = info.stepCount;
   return success();
 }
 
-static LogicalResult collectProducedValues(DeferredTransferPlan& plan) {
-  for (ScheduledInfo& info : plan.scheduled) {
-    SmallVector<int64_t> laneStarts, laneCounts;
-    size_t tableSize = info.blocks.size() * info.cores.size();
-    auto starts = info.isBatch() ? getLaneTable(info.op, "scheduled.source_lane_starts", tableSize)
-                                 : getI64Array(info.op, "scheduled.source_lane_starts");
-    auto counts = info.isBatch() ? getLaneTable(info.op, "scheduled.source_lane_counts", tableSize)
-                                 : getI64Array(info.op, "scheduled.source_lane_counts");
-    if (failed(starts) || failed(counts))
-      return failure();
-    laneStarts = std::move(*starts);
-    laneCounts = std::move(*counts);
-    for (unsigned step = 0; step < info.blocks.size(); ++step) {
-      if (info.resultOffsets[step] < 0 || info.resultCounts[step] < 0)
-        return info.op->emitOpError("phase 2 scheduled result metadata must be non-negative");
-      for (unsigned result = 0; result < static_cast<unsigned>(info.resultCounts[step]); ++result) {
-        unsigned globalResult = info.resultOffsets[step] + result;
-        if (!info.isBatch()) {
-          auto payload = getScalarStepResult(
-            cast<SpatScheduledCompute>(info.op), *info.blocks[step], result, info.resultCounts[step]);
-          if (failed(payload))
-            return failure();
-          auto geometry = verifyScheduledPublicationGeometry(
-            info, step, globalResult, 0, laneStarts[step], laneCounts[step], *payload);
-          if (failed(geometry))
-            return failure();
-          auto produced = std::make_unique<ProducedValue>(ProducedValue {
-            &info, step, result, info.stepSourceIds[step], info.cores.front(),
-            geometry->laneStart, geometry->laneCount, 0,
-            geometry->publishedSlotStart, geometry->publishedSlotCount,
-            geometry->payload, geometry->published});
-          info.produced.push_back(produced.get());
-          plan.producedByGraph[produced->graphId].push_back(produced.get());
-          plan.producedStorage.push_back(std::move(produced));
-          continue;
-        }
-        auto publication =
-          getBatchStepResult(cast<SpatScheduledComputeBatch>(info.op), *info.blocks[step], globalResult);
-        if (failed(publication))
-          return failure();
-        for (unsigned lane = 0; lane < info.cores.size(); ++lane) {
-          size_t laneIndex = step * info.cores.size() + lane;
-          auto geometry = verifyScheduledPublicationGeometry(
-            info, step, globalResult, lane, laneStarts[laneIndex],
-            laneCounts[laneIndex], publication->payload, publication->insertion);
-          if (failed(geometry))
-            return failure();
-          auto produced = std::make_unique<ProducedValue>(ProducedValue {
-            &info, step, result, info.stepSourceIds[step], info.cores[lane],
-            geometry->laneStart, geometry->laneCount, lane,
-            geometry->publishedSlotStart, geometry->publishedSlotCount,
-            geometry->payload, geometry->published});
-          info.produced.push_back(produced.get());
-          plan.producedByGraph[produced->graphId].push_back(produced.get());
-          plan.producedStorage.push_back(std::move(produced));
-        }
-      }
+static LogicalResult collectProducedValues(
+    const ScheduledComputeMaterializationResult &materialization,
+    DeferredTransferPlan &plan) {
+  if (plan.scheduled.size() != materialization.materializedSchedules.size())
+    return failure();
+  for (auto [recordIndex, record] :
+       llvm::enumerate(materialization.materializedSchedules)) {
+    ScheduledInfo &info = plan.scheduled[recordIndex];
+    DenseMap<ComputeInstance, std::pair<unsigned, unsigned>> coordinates;
+    if (record.runs.empty()) {
+      for (auto [step, stepPlan] : llvm::enumerate(record.stepPlans))
+        for (auto [lane, instance] :
+             llvm::enumerate(stepPlan.stepTuple.instances))
+          coordinates[instance] = {step, lane};
+    } else {
+      for (auto [step, run] : llvm::enumerate(record.runs))
+        for (const ComputeInstance &instance : run.instances)
+          coordinates[instance] = {step, 0};
+    }
+
+    for (const MaterializedStepValue &value : record.stepValues) {
+      auto coordinate = coordinates.find(value.instance);
+      if (coordinate == coordinates.end()
+          || coordinate->second.second >= info.cores.size()
+          || value.graphId < 0 || value.laneStart < 0
+          || value.laneCount <= 0 || !value.payload)
+        return info.op->emitOpError(
+          "phase 2 received an invalid scheduled materialization record");
+      unsigned step = coordinate->second.first;
+      unsigned lane = coordinate->second.second;
+      auto produced = std::make_unique<ProducedValue>(ProducedValue {
+        &info, step, value.resultIndex, value.graphId, info.cores[lane],
+        value.laneStart, value.laneCount, lane, value.payloadLaneStart,
+        value.payloadLaneCount, value.payload, value.published});
+      info.produced.push_back(produced.get());
+      plan.producedByGraph[produced->graphId].push_back(produced.get());
+      plan.producedStorage.push_back(std::move(produced));
     }
   }
   return success();
@@ -293,206 +127,176 @@ static FailureOr<ProducedValue*> findProducer(DeferredTransferPlan& plan,
   return match;
 }
 
-struct RequirementPoint {
-  struct SliceGeometry {
-    SmallVector<int64_t> offsets;
-    SmallVector<int64_t> sizes;
-    SmallVector<int64_t> strides;
-  };
-
-  ProducedValue* producer = nullptr;
-  Type fragmentType;
-  std::optional<int64_t> graphLane;
-  std::optional<int64_t> localOffset;
-  std::optional<SliceGeometry> producerProjection;
-
-  bool sameFamily(const RequirementPoint& other) const {
-    return producer == other.producer && fragmentType == other.fragmentType;
-  }
-};
-
-static FailureOr<SmallVector<int64_t>> evaluateGeometryValues(
-    ArrayRef<OpFoldResult> values,
-    DeferredLaneValueEvaluator& evaluator,
-    unsigned lane) {
-  SmallVector<int64_t> result;
-  result.reserve(values.size());
-  for (OpFoldResult value : values) {
-    auto sequence = evaluator.evaluate(value);
-    if (failed(sequence))
-      return failure();
-    result.push_back(sequence->valueAt(lane));
-  }
-  return result;
-}
-
-static FailureOr<std::optional<RequirementPoint>>
-resolveRequirementPoint(DeferredTransferPlan& plan,
-                        DeferredExchangePlan& exchange,
-                        const DeferredProjectionLeafTemplate& leaf,
-                        DeferredLaneValueEvaluator& evaluator,
-                        unsigned lane,
-                        unsigned position,
-                        GraphBatchPublicationCache& publicationCache) {
-  auto sourceIndices = evaluator.resolveSourceOperandIndices(leaf.sourceRoot);
-  if (failed(sourceIndices))
-    return failure();
-  unsigned sourceIndex = sourceIndices->valueAt(lane);
-  auto source = dyn_cast<OpResult>(exchange.deferred.getSources()[sourceIndex]);
-  if (!source)
-    return exchange.deferred.emitOpError("phase 2 requires graph-result deferred sources"), failure();
-  auto graphId = source.getOwner()->getAttrOfType<IntegerAttr>("scheduled.graph_id");
-  if (!graphId)
-    return exchange.deferred.emitOpError("phase 2 cannot identify graph producer"), failure();
-
-  RequirementPoint point;
-  if (auto batch = dyn_cast<SpatGraphComputeBatch>(source.getOwner())) {
-    auto publication = getGraphBatchPublicationMap(batch, source.getResultNumber(), publicationCache);
-    if (failed(publication))
-      return failure();
-    int64_t physicalSlot = position;
-    if (leaf.form == DeferredLeafForm::GraphBatchProjection) {
-      auto offset = evaluator.evaluate(leaf.leadingGeometry.offsets.front());
-      auto size = evaluator.evaluate(leaf.leadingGeometry.sizes.front());
-      auto stride = evaluator.evaluate(leaf.leadingGeometry.strides.front());
-      if (failed(offset) || failed(size) || failed(stride))
-        return failure();
-      if (position >= static_cast<unsigned>(size->valueAt(lane)))
-        return std::optional<RequirementPoint>();
-      physicalSlot = offset->valueAt(lane) + static_cast<int64_t>(position) * stride->valueAt(lane);
-    }
-    else if (position >= (*publication)->physicalSlotToGraphLane.size()) {
-      return std::optional<RequirementPoint>();
-    }
-    if (physicalSlot < 0 || physicalSlot >= static_cast<int64_t>((*publication)->physicalSlotToGraphLane.size()))
-      return exchange.deferred.emitOpError("projection physical slot is outside publication map"), failure();
-    point.graphLane = (*publication)->physicalSlotToGraphLane[physicalSlot];
-    point.fragmentType = (*publication)->publicationFragmentType;
-    auto producer = findProducer(plan, exchange.deferred, graphId.getInt(), source.getResultNumber(), point.graphLane);
-    if (failed(producer))
-      return failure();
-    point.producer = *producer;
-    point.localOffset = *point.graphLane - point.producer->laneStart;
-  }
-  else {
-    if (position != 0 || !isa<SpatGraphCompute>(source.getOwner()))
-      return std::optional<RequirementPoint>();
-    point.fragmentType = leaf.form == DeferredLeafForm::ScalarProjection
-                           ? Type(leaf.reconstructedType)
-                           : source.getType();
-    auto producer = findProducer(plan, exchange.deferred, graphId.getInt(), source.getResultNumber(), std::nullopt);
-    if (failed(producer))
-      return failure();
-    point.producer = *producer;
-    if (leaf.form == DeferredLeafForm::ScalarProjection) {
-      RequirementPoint::SliceGeometry geometry;
-      auto offsets = evaluateGeometryValues(leaf.leadingGeometry.offsets, evaluator, lane);
-      auto sizes = evaluateGeometryValues(leaf.leadingGeometry.sizes, evaluator, lane);
-      auto strides = evaluateGeometryValues(leaf.leadingGeometry.strides, evaluator, lane);
-      if (failed(offsets) || failed(sizes) || failed(strides))
-        return failure();
-      geometry.offsets = std::move(*offsets);
-      geometry.sizes = std::move(*sizes);
-      geometry.strides = std::move(*strides);
-      point.producerProjection = std::move(geometry);
-    }
-  }
-  return std::optional<RequirementPoint>(point);
-}
-
-static void appendRequirementFamily(DeferredExchangePlan& exchange,
-                                    RequirementCoordinate coordinate,
-                                    unsigned begin,
-                                    ArrayRef<RequirementPoint> points) {
-  RequirementFamily family;
-  family.exchange = &exchange;
-  family.coordinate = coordinate;
-  family.targetLanes = LaneSet::range(begin, begin + points.size());
-  family.producer = points.front().producer;
-  family.publicationFragmentType = points.front().fragmentType;
-  auto sequence = [&](auto member) -> std::optional<StaticIntSequence> {
-    if (!(points.front().*member))
-      return std::nullopt;
-    SmallVector<int64_t> values;
-    for (const RequirementPoint& point : points)
-      values.push_back(*(point.*member));
-    return StaticIntSequence::fromValues(values);
-  };
-  family.graphLanes = sequence(&RequirementPoint::graphLane);
-  family.producerLocalOffsets = sequence(&RequirementPoint::localOffset);
-  if (points.front().producerProjection) {
-    family.producerProjection.emplace();
-    auto appendGeometry = [&](auto member,
-                              SmallVectorImpl<StaticIntSequence>& target) {
-      for (size_t dimension = 0;
-           dimension < ((*points.front().producerProjection).*member).size();
-           ++dimension) {
-        SmallVector<int64_t> values;
-        values.reserve(points.size());
-        for (const RequirementPoint& point : points)
-          values.push_back(((*point.producerProjection).*member)[dimension]);
-        target.push_back(StaticIntSequence::fromValues(values));
-      }
-    };
-    appendGeometry(&RequirementPoint::SliceGeometry::offsets,
-                   family.producerProjection->offsets);
-    appendGeometry(&RequirementPoint::SliceGeometry::sizes,
-                   family.producerProjection->sizes);
-    appendGeometry(&RequirementPoint::SliceGeometry::strides,
-                   family.producerProjection->strides);
-  }
-  exchange.requirements.push_back(std::move(family));
-}
-
 static LogicalResult buildRequirementFamilies(DeferredTransferPlan& plan,
                                               DeferredExchangePlan& exchange,
                                               GraphBatchPublicationCache& publicationCache) {
-  DeferredLaneValueEvaluator evaluator(exchange.program, exchange.targetLaneCount);
-  for (auto leafItem : llvm::enumerate(exchange.program.leaves)) {
-    unsigned leafIndex = leafItem.index();
-    const DeferredProjectionLeafTemplate& leaf = leafItem.value();
-    unsigned positionCount = 1;
-    if (leaf.form == DeferredLeafForm::GraphBatchProjection) {
-      auto sizes = evaluator.evaluate(leaf.leadingGeometry.sizes.front());
-      if (failed(sizes))
+  for (unsigned specialization = 0;
+       specialization < exchange.program.specializationCount;
+       ++specialization) {
+    DeferredLaneValueEvaluator evaluator(
+      exchange.program, exchange.targetLaneCount, specialization);
+    for (auto leafItem : llvm::enumerate(exchange.program.leaves)) {
+      unsigned leafIndex = leafItem.index();
+      const DeferredProjectionLeafTemplate &leaf = leafItem.value();
+      auto sourceIndices = evaluator.resolveSourceOperandIndices(leaf.sourceRoot);
+      if (failed(sourceIndices))
         return failure();
-      for (unsigned lane = 0; lane < exchange.targetLaneCount; ++lane)
-        positionCount = std::max<unsigned>(positionCount, sizes->valueAt(lane));
-    }
-    else {
-      auto sources = evaluator.resolveSourceOperandIndices(leaf.sourceRoot);
-      if (failed(sources))
-        return failure();
-      for (unsigned lane = 0; lane < exchange.targetLaneCount; ++lane) {
-        auto source = dyn_cast<OpResult>(exchange.deferred.getSources()[sources->valueAt(lane)]);
-        auto type = source ? dyn_cast<RankedTensorType>(source.getType()) : RankedTensorType();
-        if (source && isa<SpatGraphComputeBatch>(source.getOwner()) && type)
-          positionCount = std::max<unsigned>(positionCount, type.getDimSize(0));
-      }
-    }
-    for (unsigned position = 0; position < positionCount; ++position) {
-      unsigned runBegin = 0;
-      SmallVector<RequirementPoint> run;
-      auto flush = [&] {
-        if (!run.empty())
-          appendRequirementFamily(exchange, {static_cast<unsigned>(leafIndex), position}, runBegin, run);
-        run.clear();
-      };
-      for (unsigned lane = 0; lane < exchange.targetLaneCount; ++lane) {
-        auto point = resolveRequirementPoint(plan, exchange, leaf, evaluator, lane, position, publicationCache);
-        if (failed(point))
-          return failure();
-        if (!*point) {
-          flush();
-          continue;
+      DeferredStaticSliceGeometry geometry;
+      auto evaluate = [&](ArrayRef<OpFoldResult> values,
+                          SmallVectorImpl<StaticIntSequence> &target) {
+        for (OpFoldResult value : values) {
+          auto sequence = evaluator.evaluate(value);
+          if (failed(sequence))
+            return failure();
+          target.push_back(std::move(*sequence));
+        }
+        return success();
+      };
+      if (failed(evaluate(leaf.leadingGeometry.offsets, geometry.offsets))
+          || failed(evaluate(leaf.leadingGeometry.sizes, geometry.sizes))
+          || failed(evaluate(leaf.leadingGeometry.strides, geometry.strides)))
+        return failure();
+      struct Source {
+        OpResult value;
+        int64_t graphId;
+        const GraphBatchPublicationMap *publication = nullptr;
+        ProducedValue *scalarProducer = nullptr;
+      };
+      SmallVector<Source> sources;
+      unsigned positionCount = 1;
+      for (unsigned lane = 0; lane < exchange.targetLaneCount; ++lane) {
+        Value value = exchange.deferred.getSources()[sourceIndices->valueAt(lane)];
+        while (auto selection = value.getDefiningOp<SpatDeferredSourceSelectOp>()) {
+          auto selectors = evaluator.evaluate(selection.getSelector());
+          if (failed(selectors))
+            return exchange.deferred.emitOpError(
+              "phase 2 cannot evaluate deferred source selection"), failure();
+          int64_t selected = selectors->valueAt(lane);
+          if (selected < 0 || selected >= static_cast<int64_t>(selection.getSources().size()))
+            return exchange.deferred.emitOpError(
+              "phase 2 deferred source selection is out of range"), failure();
+          value = selection.getSources()[selected];
+        }
+        auto source = dyn_cast<OpResult>(value);
+        if (!source)
+          return exchange.deferred.emitOpError(
+            "phase 2 requires graph-result deferred sources"), failure();
+        auto graphId = source.getOwner()->getAttrOfType<IntegerAttr>("scheduled.graph_id");
+        if (!graphId)
+          return exchange.deferred.emitOpError(
+            "phase 2 cannot identify graph producer"), failure();
+        Source resolved {source, graphId.getInt()};
+        if (auto batch = dyn_cast<SpatGraphComputeBatch>(source.getOwner())) {
+          auto publication = getGraphBatchPublicationMap(
+            batch, source.getResultNumber(), publicationCache);
+          if (failed(publication))
+            return failure();
+          if (leaf.form != DeferredLeafForm::GraphBatchProjection)
+            positionCount = std::max<unsigned>(
+              positionCount, (*publication)->physicalSlotToGraphLane.size());
+        }
+        else if (isa<SpatGraphCompute>(source.getOwner())) {
+          auto producer = findProducer(plan, exchange.deferred, graphId.getInt(),
+                                       source.getResultNumber(), std::nullopt);
+          if (failed(producer))
+            return failure();
+          resolved.scalarProducer = *producer;
+        }
+        sources.push_back(resolved);
+      }
+      for (Source &source : sources)
+        if (auto batch = dyn_cast<SpatGraphComputeBatch>(source.value.getOwner()))
+          source.publication = &publicationCache.find(
+            {batch, source.value.getResultNumber()})->second;
+      if (leaf.form == DeferredLeafForm::GraphBatchProjection)
+        for (unsigned lane = 0; lane < exchange.targetLaneCount; ++lane)
+          positionCount = std::max<unsigned>(
+            positionCount, geometry.sizes.front().valueAt(lane));
+
+      for (unsigned position = 0; position < positionCount; ++position) {
+        SmallVector<ProducedValue *> producers(exchange.targetLaneCount);
+        SmallVector<Type> fragmentTypes(exchange.targetLaneCount);
+        SmallVector<int64_t> graphLanes(exchange.targetLaneCount, -1);
+        SmallVector<int64_t> localOffsets(exchange.targetLaneCount, -1);
+        for (unsigned lane = 0; lane < exchange.targetLaneCount; ++lane) {
+          Source source = sources[lane];
+          if (source.publication) {
+            int64_t slot = position;
+            if (leaf.form == DeferredLeafForm::GraphBatchProjection) {
+              if (position >= geometry.sizes.front().valueAt(lane))
+                continue;
+              slot = geometry.offsets.front().valueAt(lane)
+                   + static_cast<int64_t>(position)
+                       * geometry.strides.front().valueAt(lane);
+            }
+            else if (position >= source.publication->physicalSlotToGraphLane.size())
+              continue;
+            if (slot < 0 || slot >= static_cast<int64_t>(
+                  source.publication->physicalSlotToGraphLane.size()))
+              return exchange.deferred.emitOpError(
+                "projection physical slot is outside publication map"),
+                failure();
+            graphLanes[lane] = source.publication->physicalSlotToGraphLane[slot];
+            fragmentTypes[lane] = source.publication->publicationFragmentType;
+            auto producer = findProducer(
+              plan, exchange.deferred, source.graphId,
+              source.value.getResultNumber(),
+              graphLanes[lane]);
+            if (failed(producer))
+              return failure();
+            producers[lane] = *producer;
+            localOffsets[lane] = graphLanes[lane] - (*producer)->laneStart;
+            if (!(*producer)->scheduled->isBatch())
+              localOffsets[lane] += (*producer)->publishedSlotStart;
+          }
+          else if (position == 0 && source.scalarProducer) {
+            fragmentTypes[lane] = leaf.form == DeferredLeafForm::ScalarProjection
+                                ? Type(leaf.reconstructedType)
+                                : source.value.getType();
+            producers[lane] = source.scalarProducer;
+          }
+        }
+
+        StaticIntSequence graphLaneSequence =
+          StaticIntSequence::fromValues(graphLanes);
+        StaticIntSequence localOffsetSequence =
+          StaticIntSequence::fromValues(localOffsets);
+        for (unsigned begin = 0; begin < exchange.targetLaneCount;) {
+          if (!producers[begin]) {
+            ++begin;
+            continue;
+          }
+          unsigned end = begin + 1;
+          while (end < exchange.targetLaneCount
+                 && producers[end] == producers[begin]
+                 && fragmentTypes[end] == fragmentTypes[begin])
+            ++end;
+          RequirementFamily family;
+          family.exchange = &exchange;
+          family.coordinate = {specialization, leafIndex, position};
+          family.targetLanes = LaneSet::range(begin, end);
+          family.producer = producers[begin];
+          family.publicationFragmentType = fragmentTypes[begin];
+          if (graphLanes[begin] >= 0) {
+            family.graphLanes = graphLaneSequence.slice(begin, end - begin);
+            family.producerLocalOffsets =
+              localOffsetSequence.slice(begin, end - begin);
+          }
+          else if (leaf.form == DeferredLeafForm::ScalarProjection) {
+            family.producerProjection.emplace();
+            auto slice = [&](ArrayRef<StaticIntSequence> source,
+                             SmallVectorImpl<StaticIntSequence> &target) {
+              for (const StaticIntSequence &sequence : source)
+                target.push_back(sequence.slice(begin, end - begin));
+            };
+            slice(geometry.offsets, family.producerProjection->offsets);
+            slice(geometry.sizes, family.producerProjection->sizes);
+            slice(geometry.strides, family.producerProjection->strides);
+          }
+          exchange.requirements.push_back(std::move(family));
+          begin = end;
         }
-        if (!run.empty() && !run.front().sameFamily(**point))
-          flush();
-        if (run.empty())
-          runBegin = lane;
-        run.push_back(**point);
       }
-      flush();
     }
   }
   return success();
@@ -563,7 +367,8 @@ static LogicalResult buildExchanges(func::FuncOp funcOp, DeferredTransferPlan& p
     if (!targetOp)
       targetOp = deferred->getParentOfType<SpatScheduledComputeBatch>();
     ScheduledInfo* target = scheduledByOp.lookup(targetOp);
-    auto step = target ? getStepIndex(*target, getScheduledBlock(deferred, targetOp)) : FailureOr<unsigned>(failure());
+    auto step = target ? getStepIndex(*target, deferred)
+                       : FailureOr<unsigned>(failure());
     auto program = analyzeDeferredProgramTemplate(deferred);
     if (!target || failed(step) || failed(program))
       return deferred.emitOpError("phase 2 cannot normalize deferred communication");
@@ -586,6 +391,13 @@ static LogicalResult
 retargetBlueprint(DeferredTransferPlan& plan, SpatBlueprintOp blueprint, GraphBatchPublicationCache& publicationCache) {
   if (blueprint.getMode() != "fragment_assembly")
     return success();
+  bool escapesScheduledGraph = llvm::any_of(
+    blueprint.getOutput().getUses(), [](OpOperand &use) {
+      return !isa<SpatGraphCompute, SpatGraphComputeBatch,
+                  SpatDeferredCommunicationOp>(use.getOwner());
+    });
+  if (!escapesScheduledGraph)
+    return success();
   auto operandIndices = blueprint.getFragmentOperandIndices();
   auto sourceSlots = blueprint.getFragmentSourceSlots();
   if (!operandIndices || !sourceSlots)
@@ -608,6 +420,9 @@ retargetBlueprint(DeferredTransferPlan& plan, SpatBlueprintOp blueprint, GraphBa
       auto producer = findProducer(plan, blueprint, graphId.getInt(), result.getResultNumber(), graphLane);
       if (failed(producer))
         return failure();
+      if (!(*producer)->published)
+        return blueprint.emitOpError(
+          "phase 2 Blueprint source has no scheduled publication"), failure();
       source = (*producer)->published;
       slot = (*producer)->publishedSlotStart + graphLane - (*producer)->laneStart;
       if (slot < (*producer)->publishedSlotStart
@@ -636,9 +451,12 @@ retargetBlueprint(DeferredTransferPlan& plan, SpatBlueprintOp blueprint, GraphBa
 
 } // namespace
 
-FailureOr<DeferredTransferPlan> buildDeferredTransferPlan(func::FuncOp funcOp) {
+FailureOr<DeferredTransferPlan> buildDeferredTransferPlan(
+    func::FuncOp funcOp,
+    const ScheduledComputeMaterializationResult &materialization) {
   DeferredTransferPlan plan;
-  if (failed(collectScheduledOperations(funcOp, plan)) || failed(collectProducedValues(plan))
+  if (failed(collectScheduledOperations(materialization, plan))
+      || failed(collectProducedValues(materialization, plan))
       || failed(buildExchanges(funcOp, plan)))
     return failure();
   return std::move(plan);

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

@@ -3,6 +3,7 @@
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 
 #include "DeferredCommunicationModel.hpp"
+#include "ScheduledComputeMaterialization.hpp"
 
 namespace onnx_mlir::spatial {
 
@@ -14,7 +15,9 @@ struct DeferredTransferPlan {
   llvm::SmallVector<unsigned> stepCounts;
 };
 
-mlir::FailureOr<DeferredTransferPlan> buildDeferredTransferPlan(mlir::func::FuncOp funcOp);
+mlir::FailureOr<DeferredTransferPlan>
+buildDeferredTransferPlan(mlir::func::FuncOp funcOp,
+                          const ScheduledComputeMaterializationResult &materialization);
 
 mlir::LogicalResult retargetDeferredPublications(mlir::func::FuncOp funcOp, DeferredTransferPlan& plan);
 

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

@@ -44,11 +44,9 @@ struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, Operatio
       signalPassFailure();
       return;
     }
-
     // Phase 1 is intentionally dumped before its verifier: malformed deferred
     // payloads must be diagnosed from the producer-owned body.
-    dumpModule(moduleOp, "spatial2_scheduled_no_comm", /*assumeVerified=*/true);
-
+    dumpModule(moduleOp, "spatial3_scheduled_no_comm", /*assumeVerified=*/true);
     if (failed(verifyMaterializedScheduleMapping(funcOp,
                                                  schedule,
                                                  materialization->peftClassPlans,
@@ -68,14 +66,6 @@ struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, Operatio
       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();
@@ -83,30 +73,35 @@ struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, Operatio
     }
 
     SpatialDataflowExportStage exportMode = getSpatialDataflowExportStage();
-    if (shouldExportSpatialDataflowStage(exportMode, SpatialDataflowExportStage::Spatial2)
-        && failed(exportSpatialDataflowCsvScheduled(funcOp, "spatial2_scheduled_no_comm", "spatial2"))) {
+    if (shouldExportSpatialDataflowStage(exportMode, SpatialDataflowExportStage::Spatial3)
+        && failed(exportSpatialDataflowCsvScheduled(
+          funcOp, materialization->materializedSchedules,
+          "spatial3_scheduled_no_comm", "spatial3"))) {
       signalPassFailure();
       return;
     }
 
-    dumpScheduledComputeReportAndModule(moduleOp,
-                                        funcOp,
-                                        schedule,
-                                        materialization->peftClassPlans,
-                                        materialization->materializedSchedules);
-    if (failed(realizeDeferredCommunication(funcOp))) {
+    dumpScheduledComputeReport(moduleOp,
+                               funcOp,
+                               schedule,
+                               materialization->peftClassPlans,
+                               materialization->materializedSchedules);
+    if (failed(realizeDeferredCommunication(funcOp, *materialization))) {
       moduleOp.emitError("MergeComputeNodes phase 2 communication realization failed");
       signalPassFailure();
       return;
     }
-    dumpModule(moduleOp, "spatial3_scheduled", /*assumeVerified=*/true);
-    if (failed(verifyScheduledSpatialInvariants(funcOp))) {
+    dumpModule(moduleOp, "spatial4_scheduled", /*assumeVerified=*/true);
+    if (failed(verifyScheduledResultsLive(materialization->materializedSchedules))
+        || failed(verifyScheduledSpatialInvariants(funcOp))) {
       moduleOp.emitError("scheduled Spatial phase 2 verification failed");
       signalPassFailure();
       return;
     }
-    if (shouldExportSpatialDataflowStage(exportMode, SpatialDataflowExportStage::Spatial3)
-        && failed(exportSpatialDataflowCsvScheduled(funcOp, "spatial3_scheduled", "spatial3"))) {
+    if (shouldExportSpatialDataflowStage(exportMode, SpatialDataflowExportStage::Spatial4)
+        && failed(exportSpatialDataflowCsvScheduled(
+          funcOp, materialization->materializedSchedules,
+          "spatial4_scheduled", "spatial4"))) {
       signalPassFailure();
     }
   }

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

@@ -6,6 +6,7 @@
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 
 #include <map>
@@ -18,11 +19,6 @@ namespace spatial {
 using namespace mlir;
 namespace {
 
-struct BatchFragmentSpec {
-  RankedTensorType resultType;
-  RankedTensorType sourceSliceType;
-};
-
 static SmallVector<OpFoldResult> remapMixedOffsets(ArrayRef<OpFoldResult> mixedOffsets, IRMapping &mapper) {
   SmallVector<OpFoldResult> remapped;
   remapped.reserve(mixedOffsets.size());
@@ -35,135 +31,34 @@ static SmallVector<OpFoldResult> remapMixedOffsets(ArrayRef<OpFoldResult> mixedO
   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);
+  appendComputeBlockArguments(argTypes, argLocs, scheduled.getWeights(),
+                              scheduled.getInputs(), 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() != sourceType.getRank() + 1 || destType.getDimSize(0) != batch.getLaneCount()
-        || destType.getElementType() != sourceType.getElementType())
-      return batch.emitOpError("graph_compute_batch result must be a leading physical-slot dimension followed by its fragment");
-    if (!llvm::equal(destType.getShape().drop_front(), sourceType.getShape()))
-      return batch.emitOpError("graph_compute_batch result trailing shape must match its published fragment");
-    if (!insert.hasUnitStride())
-      return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
-    auto offsets = insert.getMixedOffsets();
-    auto sizes = insert.getMixedSizes();
-    auto strides = insert.getMixedStrides();
-    if (offsets.size() != static_cast<size_t>(destType.getRank()) || sizes.size() != static_cast<size_t>(destType.getRank())
-        || strides.size() != static_cast<size_t>(destType.getRank()))
-      return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
-    if (!isa<Value>(offsets.front()) || !valueTransitivelyDependsOn(cast<Value>(offsets.front()), *batch.getLaneArgument()))
-      return batch.emitOpError("graph_compute_batch publication must select its physical slot in dimension zero");
-    for (unsigned dim = 1; dim < offsets.size(); ++dim) {
-      auto offset = dyn_cast<Attribute>(offsets[dim]);
-      auto integer = dyn_cast_or_null<IntegerAttr>(offset);
-      if (!integer || integer.getInt() != 0)
-        return batch.emitOpError("graph_compute_batch publication must have zero trailing offsets");
-    }
-    auto staticIndex = [](OpFoldResult value) -> std::optional<int64_t> {
-      auto attr = dyn_cast<Attribute>(value);
-      auto integer = dyn_cast_or_null<IntegerAttr>(attr);
-      return integer ? std::optional<int64_t>(integer.getInt()) : std::nullopt;
-    };
-    if (staticIndex(sizes.front()) != 1)
-      return batch.emitOpError("graph_compute_batch publication sizes must be [1] plus the fragment shape");
-    for (auto [size, dim] : llvm::zip_equal(ArrayRef<OpFoldResult>(sizes).drop_front(), sourceType.getShape()))
-      if (staticIndex(size) != dim)
-        return batch.emitOpError("graph_compute_batch publication sizes must be [1] plus the fragment shape");
-    return BatchFragmentSpec {spatial::getGraphBatchPhysicalResultType(fragmentLaneCount, sourceType), sourceType};
-  }
-
-  return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
-}
-
-
 static SourceLaneSelector buildSourceLaneSelector(PatternRewriter &rewriter,
                                                   const ComputeStepTuple &stepTuple,
                                                   Operation *constantAnchor,
@@ -223,123 +118,6 @@ static FailureOr<Value> buildSourceLaneStartForScheduledLane(OpBuilder &builder,
   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) &&
-            !isDeferredFragmentAssemblyInput(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) &&
-            !isDeferredFragmentAssemblyInput(input, instance))
-          appendUnique(peftClassPlan.inputs, input);
-    }
-  }
-
-  return success();
-}
-
 static void cloneComputeBody(OpBuilder &builder, Block &source, IRMapping &mapper,
                              SmallVectorImpl<Value> &yieldedValues,
                              const llvm::SmallPtrSetImpl<Operation *> &absorbed) {
@@ -351,18 +129,22 @@ static void cloneComputeBody(OpBuilder &builder, Block &source, IRMapping &mappe
     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,
-                                                            SmallVectorImpl<Value> &yieldedValues) {
+static LogicalResult materializeResultfulBatchRun(
+    PatternRewriter &rewriter, SpatComputeBatch batch,
+    const ScheduledInstanceRun &run, ValueRange scheduledWeights,
+    ValueRange scheduledInputs, Block &block,
+    const MergeScheduleResult &schedule,
+    const DenseMap<ProducerValueKey, MaterializedProducerRef> &availableValues,
+    SmallVectorImpl<Value> &yieldedValues) {
+  const ComputeInstance &first = run.instances.front();
+  const ComputeInstance &last = run.instances.back();
+  uint32_t runLaneCount = last.laneStart + last.laneCount - first.laneStart;
+  ComputeInstance runInstance {batch.getOperation(), first.laneStart,
+                               runLaneCount};
   SmallVector<Value> initResults;
   SmallVector<BatchFragmentSpec> fragmentSpecs;
   for (unsigned resultIndex = 0; resultIndex < batch.getNumResults(); ++resultIndex) {
-    auto spec = getBatchFragmentSpec(batch, resultIndex, instance.laneCount);
+    auto spec = getBatchFragmentSpec(batch, resultIndex, runLaneCount);
     if (failed(spec))
       return failure();
     fragmentSpecs.push_back(*spec);
@@ -372,8 +154,8 @@ static LogicalResult materializeResultfulBatchChunkAsScalar(PatternRewriter &rew
     initResults.push_back(*empty);
   }
 
-  Value lower = getOrCreateIndexConstant(rewriter, batch.getOperation(), instance.laneStart);
-  Value upper = getOrCreateIndexConstant(rewriter, batch.getOperation(), instance.laneStart + instance.laneCount);
+  Value lower = getOrCreateIndexConstant(rewriter, batch.getOperation(), first.laneStart);
+  Value upper = getOrCreateIndexConstant(rewriter, batch.getOperation(), first.laneStart + runLaneCount);
   Value step = getOrCreateIndexConstant(rewriter, batch.getOperation(), 1);
   auto loop = buildNormalizedScfFor(
     rewriter,
@@ -386,11 +168,12 @@ static LogicalResult materializeResultfulBatchChunkAsScalar(PatternRewriter &rew
 
       IRMapping mapper;
       mapper.map(*batch.getLaneArgument(), originalLane);
-      Value localLane = arith::SubIOp::create(builder,
-                                              bodyLoc,
-                                              originalLane,
-                                              getOrCreateIndexConstant(rewriter, batch.getOperation(), instance.laneStart))
-                          .getResult();
+      Value localLane = runLaneCount == 1
+        ? getOrCreateIndexConstant(rewriter, batch.getOperation(), 0)
+        : arith::SubIOp::create(
+            builder, bodyLoc, originalLane,
+            getOrCreateIndexConstant(
+              rewriter, batch.getOperation(), first.laneStart));
       for (auto [index, weight] : llvm::enumerate(batch.getWeights()))
         mapper.map(*batch.getWeightArgument(index), getBlockOperand(block, scheduledWeights, weight));
       SmallVector<DeferredInputPlan> inputPlans;
@@ -400,19 +183,19 @@ static LogicalResult materializeResultfulBatchChunkAsScalar(PatternRewriter &rew
                                      input.getLoc(),
                                      input,
                                      *batch.getInputArgument(index),
-                                     instance,
+                                     runInstance,
                                      schedule,
                                      scheduledInputs,
                                      block,
                                      scheduledWeights.size(),
-                                     ArrayRef<ProducerValueKey> {},
+                                     availableValues,
                                      *batch.getLaneArgument(),
                                      originalLane,
                                      plan)))
           return failure();
         plan.scalarizedLocalLane = localLane;
         plan.scalarizedGraphLaneBase = lower;
-        plan.scalarizedLaneCount = instance.laneCount;
+        plan.scalarizedLaneCount = runLaneCount;
         plan.scalarizedHoistBlock = &block;
         inputPlans.push_back(std::move(plan));
       }
@@ -474,19 +257,25 @@ static LogicalResult materializeSingleCpuPeftClass(
   auto instancesIt = peftClassPlan.instancesByCpu.find(cpu);
   assert(instancesIt != peftClassPlan.instancesByCpu.end() && "missing single-cpu schedule");
   const SmallVector<ComputeInstance> &instances = instancesIt->second;
+  Block *block = createScheduledComputeBlock(
+    rewriter, scheduled, instances.front().op->getLoc());
+  DenseMap<ProducerValueKey, MaterializedProducerRef> availableValues;
+  DenseMap<ProducerValueKey, unsigned> publicationIndices;
+  for (const ScheduledPublication &publication : record.publications)
+    publicationIndices[publication.producer] = publication.scheduledResultIndex;
+  SmallVector<Value> publishedOutputs(scheduled.getNumResults());
+  record.runs = buildScheduledInstanceRuns(instances);
+  for (const ScheduledInstanceRun &run : record.runs) {
+    const ComputeInstance &instance = run.instances.front();
+    for (const ComputeInstance &member : run.instances) {
+      GraphComputeBlockKey key = getGraphComputeBlockKey(member);
+      graphComputeToBlockMap[key] = block;
+      record.computeKeys.push_back(key);
+      record.blocks.push_back(block);
+    }
 
-  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);
+    Operation *previous = block->empty() ? nullptr : &block->back();
+    rewriter.setInsertionPointToEnd(block);
     SmallVector<Value> yieldedValues;
     if (auto compute = dyn_cast<SpatCompute>(instance.op)) {
       IRMapping mapper;
@@ -504,7 +293,7 @@ static LogicalResult materializeSingleCpuPeftClass(
                                      scheduled.getInputs(),
                                      *block,
                                      scheduled.getWeights().size(),
-                                     carriedKeys,
+                                     availableValues,
                                      {},
                                      {},
                                      plan)))
@@ -517,45 +306,76 @@ static LogicalResult materializeSingleCpuPeftClass(
       cloneComputeBody(rewriter, compute.getBody().front(), mapper, yieldedValues, absorbed);
     } else {
       auto batch = cast<SpatComputeBatch>(instance.op);
-      if (failed(materializeResultfulBatchChunkAsScalar(rewriter,
-                                                        batch,
-                                                        instance,
-                                                        scheduled.getWeights(),
-                                                        scheduled.getInputs(),
-                                                        *block,
-                                                        schedule,
-                                                        yieldedValues)))
+      if (failed(materializeResultfulBatchRun(rewriter, batch, run,
+                                              scheduled.getWeights(),
+                                              scheduled.getInputs(), *block,
+                                              schedule, availableValues,
+                                              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);
+    if (yieldedValues.size()
+        != getComputeInstanceResultValueCount(instance))
+      return scheduled.emitOpError(
+        "scheduled scalar step produced an unexpected result count");
+    for (auto [resultIndex, value] : llvm::enumerate(yieldedValues)) {
+      int64_t runLaneStart = instance.laneStart;
+      const ComputeInstance &last = run.instances.back();
+      int64_t runLaneCount = last.laneStart + last.laneCount - runLaneStart;
+      availableValues[{{instance.op, static_cast<uint32_t>(runLaneStart),
+                        static_cast<uint32_t>(runLaneCount)}, resultIndex}] =
+        {value, 0, runLaneCount};
+      for (const ComputeInstance &member : run.instances) {
+        ProducerValueKey key {member, resultIndex};
+        int64_t payloadLaneStart = member.laneStart - runLaneStart;
+        availableValues[key] = {value, payloadLaneStart, member.laneCount};
+        Value published;
+        if (auto it = publicationIndices.find(key);
+            it != publicationIndices.end()) {
+          Value publication = value;
+          if (run.instances.size() != 1) {
+            auto spec = getBatchFragmentSpec(
+              cast<SpatComputeBatch>(member.op), resultIndex,
+              member.laneCount);
+            if (failed(spec))
+              return failure();
+            MixedSliceGeometry geometry;
+            auto payloadType = cast<RankedTensorType>(value.getType());
+            geometry.offsets.assign(payloadType.getRank(),
+                                    rewriter.getIndexAttr(0));
+            geometry.offsets.front() =
+              rewriter.getIndexAttr(payloadLaneStart);
+            geometry.sizes.reserve(payloadType.getRank());
+            for (int64_t dimension : spec->resultType.getShape())
+              geometry.sizes.push_back(rewriter.getIndexAttr(dimension));
+            geometry.strides.assign(payloadType.getRank(),
+                                    rewriter.getIndexAttr(1));
+            publication = extractMixedSliceOrIdentity(
+              rewriter, member.op->getLoc(), value, spec->resultType,
+              geometry);
+          }
+          publishedOutputs[it->second] = publication;
+          published = scheduled.getResult(it->second);
+        }
+        auto graphId = member.op->getAttrOfType<IntegerAttr>(
+          "scheduled.graph_id");
+        if (!graphId)
+          return member.op->emitOpError(
+            "scheduled materialization requires graph identity metadata");
+        record.stepValues.push_back(
+          {member, static_cast<unsigned>(resultIndex), value,
+           graphId.getInt(), member.laneStart, member.laneCount,
+           payloadLaneStart, member.laneCount, published});
+      }
     }
+    Operation *anchor = previous ? previous->getNextNode() : &block->front();
+    record.stepAnchors.push_back(anchor);
   }
+  if (llvm::any_of(publishedOutputs, [](Value value) { return !value; }))
+    return scheduled.emitOpError(
+      "scheduled scalar materialization did not produce every declared publication");
+  rewriter.setInsertionPointToEnd(block);
+  SpatYieldOp::create(rewriter, scheduled.getLoc(), publishedOutputs);
   return success();
 }
 
@@ -566,7 +386,7 @@ static SmallVector<OpFoldResult> buildScheduledOutputInsertOffsets(OpBuilder &bu
                                                                    Location loc,
                                                                    Value scheduledLane,
                                                                    int64_t lanesPerScheduledLane,
-                                                                   RankedTensorType localFragmentType,
+                                                                   int64_t destinationRank,
                                                                    Operation *constantAnchor) {
   SmallVector<OpFoldResult> offsets;
   Value scheduledOutputLane = scheduledLane;
@@ -575,10 +395,82 @@ static SmallVector<OpFoldResult> buildScheduledOutputInsertOffsets(OpBuilder &bu
       builder, loc, scheduledLane, lanesPerScheduledLane, constantAnchor);
   }
   offsets.push_back(scheduledOutputLane);
-  offsets.append(localFragmentType.getRank() - 1, OpFoldResult(builder.getIndexAttr(0)));
+  offsets.append(destinationRank - 1, OpFoldResult(builder.getIndexAttr(0)));
   return offsets;
 }
 
+struct BatchPublication {
+  Value fragment;
+  BlockArgument destination;
+  SmallVector<OpFoldResult> offsets;
+  SmallVector<OpFoldResult> sizes;
+  SmallVector<OpFoldResult> strides;
+};
+
+static LogicalResult collectMultiCpuPublications(
+  PatternRewriter &rewriter,
+  SpatScheduledComputeBatch scheduled,
+  const ComputeStepTuple &stepTuple,
+  const ComputeInstance &representative,
+  ArrayRef<Value> localFragments,
+  const DenseMap<ProducerValueKey, unsigned> &publicationIndices,
+  Block &block,
+  ScheduledMaterializationRecord &record,
+  SmallVectorImpl<BatchPublication> &publications) {
+  Value scheduledLane = block.getArgument(0);
+  for (auto [resultIndex, localFragment] : llvm::enumerate(localFragments)) {
+    auto publicationIt = publicationIndices.find({representative, resultIndex});
+    Value published = publicationIt == publicationIndices.end()
+      ? Value()
+      : scheduled.getResult(publicationIt->second);
+    for (auto [lane, instance] : llvm::enumerate(stepTuple.instances)) {
+      auto instancePublication = publicationIndices.find({instance, resultIndex});
+      if ((instancePublication == publicationIndices.end()) != !published
+          || (instancePublication != publicationIndices.end()
+              && instancePublication->second != publicationIt->second))
+        return scheduled.emitOpError("scheduled batch tuple has inconsistent publication ownership");
+      auto graphId = instance.op->getAttrOfType<IntegerAttr>("scheduled.graph_id");
+      if (!graphId)
+        return instance.op->emitOpError("scheduled materialization requires graph identity metadata");
+      record.stepValues.push_back({instance, static_cast<unsigned>(resultIndex),
+                                   localFragment, graphId.getInt(),
+                                   instance.laneStart, instance.laneCount,
+                                   static_cast<int64_t>(lane) * instance.laneCount,
+                                   instance.laneCount, published});
+    }
+    if (!published)
+      continue;
+
+    auto localType = cast<RankedTensorType>(localFragment.getType());
+    auto destination = block.getArgument(
+      getScheduledBatchResultArgBase(scheduled) + publicationIt->second);
+    auto destinationType = cast<RankedTensorType>(destination.getType());
+    if (destinationType.getRank() != localType.getRank()
+        && destinationType.getRank() != localType.getRank() + 1)
+      return scheduled.emitOpError(
+        "scheduled publication source must match or rank-reduce into its destination");
+
+    int64_t lanesPerScheduledLane = isa<SpatCompute>(representative.op)
+      ? 1 : representative.laneCount;
+    SmallVector<OpFoldResult> offsets = buildScheduledOutputInsertOffsets(
+      rewriter, scheduled.getLoc(), scheduledLane, lanesPerScheduledLane,
+      destinationType.getRank(), scheduled.getOperation());
+    SmallVector<OpFoldResult> sizes;
+    SmallVector<OpFoldResult> strides;
+    if (destinationType.getRank() == localType.getRank() + 1) {
+      sizes.push_back(rewriter.getIndexAttr(1));
+      strides.push_back(rewriter.getIndexAttr(1));
+    }
+    for (int64_t dim : localType.getShape()) {
+      sizes.push_back(rewriter.getIndexAttr(dim));
+      strides.push_back(rewriter.getIndexAttr(1));
+    }
+    publications.push_back({localFragment, destination, std::move(offsets),
+                            std::move(sizes), std::move(strides)});
+  }
+  return success();
+}
+
 static LogicalResult materializeMultiCpuPeftClass(
   PatternRewriter &rewriter,
   SpatScheduledComputeBatch scheduled,
@@ -588,25 +480,31 @@ static LogicalResult materializeMultiCpuPeftClass(
   ScheduledMaterializationRecord &record) {
   std::map<std::vector<uint32_t>, Value> laneStartTableCache;
   ArrayRef<ScheduledStepPlan> stepPlans = record.stepPlans;
+  DenseMap<ProducerValueKey, unsigned> publicationIndices;
+  for (const ScheduledPublication &publication : record.publications)
+    publicationIndices[publication.producer] = publication.scheduledResultIndex;
+  SmallVector<Type> blockArgTypes {rewriter.getIndexType()};
+  SmallVector<Location> blockArgLocs {scheduled.getLoc()};
+  for (Value weight : scheduled.getWeights()) {
+    blockArgTypes.push_back(weight.getType());
+    blockArgLocs.push_back(weight.getLoc());
+  }
+  for (Value input : scheduled.getInputs()) {
+    blockArgTypes.push_back(input.getType());
+    blockArgLocs.push_back(input.getLoc());
+  }
+  for (Type resultType : scheduled.getResultTypes()) {
+    blockArgTypes.push_back(resultType);
+    blockArgLocs.push_back(scheduled.getLoc());
+  }
+  Block *block = rewriter.createBlock(
+    &scheduled.getBody(), scheduled.getBody().end(), blockArgTypes,
+    blockArgLocs);
+  SmallVector<BatchPublication> publications;
   for (const ScheduledStepPlan &stepPlan : stepPlans) {
     const ComputeStepTuple &stepTuple = stepPlan.stepTuple;
     SourceLaneSelector sourceLaneSelector =
       buildSourceLaneSelector(rewriter, stepTuple, scheduled.getOperation(), laneStartTableCache);
-    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;
@@ -614,7 +512,8 @@ static LogicalResult materializeMultiCpuPeftClass(
       record.blocks.push_back(block);
     }
 
-    rewriter.setInsertionPointToStart(block);
+    Operation *previous = block->empty() ? nullptr : &block->back();
+    rewriter.setInsertionPointToEnd(block);
     Value scheduledLane = block->getArgument(0);
     const ComputeInstance &representative = stepTuple.instances.front();
     SmallVector<Value> finalLocalFragments;
@@ -654,18 +553,7 @@ static LogicalResult materializeMultiCpuPeftClass(
         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());
+        finalLocalFragments.push_back(yielded);
       }
     } else {
       auto batch = cast<SpatComputeBatch>(representative.op);
@@ -779,45 +667,30 @@ static LogicalResult materializeMultiCpuPeftClass(
       finalLocalFragments.assign(loop->results.begin(), loop->results.end());
     }
 
-    struct Publication {
-      Value fragment;
-      SmallVector<OpFoldResult> offsets;
-      SmallVector<OpFoldResult> sizes;
-      SmallVector<OpFoldResult> strides;
-    };
-    SmallVector<Publication> publications;
-    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,
-        scheduled.getOperation());
-      SmallVector<OpFoldResult> sizes;
-      SmallVector<OpFoldResult> strides;
-      for (int64_t dim : localFragmentType.getShape()) {
-        sizes.push_back(rewriter.getIndexAttr(dim));
-        strides.push_back(rewriter.getIndexAttr(1));
-      }
-      publications.push_back(
-        {localFragment, std::move(offsets), std::move(sizes),
-         std::move(strides)});
-    }
-    auto inParallel = SpatInParallelOp::create(rewriter, scheduled.getLoc());
-    rewriter.setInsertionPointToStart(&inParallel.getRegion().front());
-    for (auto [resultIndex, publication] : llvm::enumerate(publications))
-      tensor::ParallelInsertSliceOp::create(
-        rewriter,
-        scheduled.getLoc(),
-        publication.fragment,
-        block->getArgument(getScheduledBatchResultArgBase(scheduled) + stepPlan.resultOffset + resultIndex),
-        publication.offsets,
-        publication.sizes,
-        publication.strides);
+    if (failed(collectMultiCpuPublications(
+          rewriter, scheduled, stepTuple, representative,
+          finalLocalFragments, publicationIndices, *block, record,
+          publications)))
+      return failure();
+    Operation *anchor = previous ? previous->getNextNode()
+                                 : (block->empty() ? nullptr : &block->front());
+    if (!anchor)
+      return scheduled.emitOpError(
+        "scheduled batch step did not materialize a body operation");
+    record.stepAnchors.push_back(anchor);
   }
+  rewriter.setInsertionPointToEnd(block);
+  if (publications.empty()) {
+    SpatYieldOp::create(rewriter, scheduled.getLoc(), ValueRange {});
+    return success();
+  }
+  auto inParallel = SpatInParallelOp::create(rewriter, scheduled.getLoc());
+  rewriter.setInsertionPointToEnd(&inParallel.getRegion().front());
+  for (const BatchPublication &publication : publications)
+    tensor::ParallelInsertSliceOp::create(
+      rewriter, scheduled.getLoc(), publication.fragment,
+      publication.destination, publication.offsets, publication.sizes,
+      publication.strides);
   return success();
 }
 
@@ -878,6 +751,7 @@ materializeScheduledCompute(func::FuncOp funcOp,
     record.canonicalPeftClassId = peftClassPlan.canonicalClassId;
     record.cpus = peftClassPlan.cpus;
     record.stepPlans = buildScheduledStepPlans(peftClassPlan);
+    record.publications = peftClassPlan.publications;
 
     if (peftClassPlan.cpus.size() == 1) {
       auto scheduled = SpatScheduledCompute::create(
@@ -886,39 +760,8 @@ materializeScheduledCompute(func::FuncOp funcOp,
         TypeRange(peftClassPlan.resultTypes),
         peftClassPlan.weights,
         peftClassPlan.inputs);
+      scheduled->setAttr("scheduled.realized", rewriter.getBoolAttr(true));
       scheduled->setAttr(kCoreIdAttrName, rewriter.getI32IntegerAttr(static_cast<int32_t>(peftClassPlan.cpus.front())));
-      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;
-      SmallVector<int64_t> stepSourceIds;
-      size_t resultOffset = 0;
-      for (const ComputeInstance &instance : peftClassPlan.instancesByCpu.lookup(peftClassPlan.cpus.front())) {
-        stepSources.push_back(rewriter.getStringAttr(getInstanceName(instance)));
-        stepSourceIds.push_back(graphIds.lookup(instance.op));
-        sourceLaneSelectors.push_back(rewriter.getStringAttr(isa<SpatCompute>(instance.op) ? "scalar" : "affine"));
-        size_t resultCount = getComputeInstanceResultValueCount(instance);
-        stepResultOffsets.push_back(static_cast<int64_t>(resultOffset));
-        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_source_ids", rewriter.getDenseI64ArrayAttr(stepSourceIds));
-      scheduled->setAttr("scheduled.step_result_offsets", rewriter.getDenseI64ArrayAttr(stepResultOffsets));
-      scheduled->setAttr("scheduled.step_result_counts", rewriter.getDenseI64ArrayAttr(stepResultCounts));
-      scheduled->setAttr("scheduled.source_lane_starts", rewriter.getDenseI64ArrayAttr(sourceLaneStarts));
-      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 {
@@ -928,36 +771,8 @@ materializeScheduledCompute(func::FuncOp funcOp,
                                                          rewriter.getI32IntegerAttr(static_cast<int32_t>(peftClassPlan.cpus.size())),
                                                          peftClassPlan.weights,
                                                          peftClassPlan.inputs);
+      scheduled->setAttr("scheduled.realized", rewriter.getBoolAttr(true));
       scheduled->setAttr(kCoreIdsAttrName, rewriter.getDenseI32ArrayAttr(toI32Array(peftClassPlan.cpus)));
-      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;
-      SmallVector<int64_t> stepSourceIds;
-      for (const ScheduledStepPlan &stepPlan : record.stepPlans) {
-        stepSources.push_back(rewriter.getStringAttr(getInstanceName(stepPlan.stepTuple.instances.front())));
-        stepSourceIds.push_back(graphIds.lookup(stepPlan.stepTuple.instances.front().op));
-        sourceLaneSelectors.push_back(rewriter.getStringAttr(usesAffineSourceLaneMapping(stepPlan.stepTuple) ? "affine" : "table"));
-        resultOffsets.push_back(static_cast<int64_t>(stepPlan.resultOffset));
-        resultCounts.push_back(static_cast<int64_t>(stepPlan.resultCount));
-        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_source_ids", rewriter.getDenseI64ArrayAttr(stepSourceIds));
-      scheduled->setAttr("scheduled.step_result_offsets", rewriter.getDenseI64ArrayAttr(resultOffsets));
-      scheduled->setAttr("scheduled.step_result_counts", rewriter.getDenseI64ArrayAttr(resultCounts));
-      scheduled->setAttr("scheduled.source_lane_starts", DenseElementsAttr::get(sourceLaneTableType, ArrayRef<int64_t>(sourceLaneStarts)));
-      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;
     }

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

@@ -5,12 +5,29 @@
 namespace onnx_mlir {
 namespace spatial {
 
+struct BatchFragmentSpec {
+  RankedTensorType resultType;
+  RankedTensorType sourceSliceType;
+};
+
 struct ScheduledComputeMaterializationResult {
   llvm::MapVector<size_t, PeftClassPlan> peftClassPlans;
   std::vector<ScheduledMaterializationRecord> materializedSchedules;
   DenseMap<GraphComputeBlockKey, Block *> graphComputeToBlockMap;
 };
 
+FailureOr<BatchFragmentSpec>
+getBatchFragmentSpec(SpatComputeBatch batch,
+                     unsigned resultIndex,
+                     uint32_t fragmentLaneCount);
+
+LogicalResult verifyPeftClassPlan(Operation *diagnosticAnchor,
+                                  const PeftClassPlan &peftClassPlan,
+                                  const MergeScheduleResult &schedule);
+
+LogicalResult collectPeftClassOperandsAndResults(
+  PeftClassPlan &peftClassPlan, const MergeScheduleResult &schedule);
+
 FailureOr<ScheduledComputeMaterializationResult>
 materializeScheduledCompute(func::FuncOp funcOp,
                             const MergeScheduleResult &schedule,

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

@@ -3,7 +3,6 @@
 #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"
@@ -20,6 +19,7 @@
 #include <vector>
 
 #include "Scheduling/ComputeInstanceUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/StaticIntSequence.hpp"
 #include "Scheduling/MergeSchedulingAnalysis.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
@@ -38,6 +38,16 @@ struct ProducerValueKey {
   }
 };
 
+struct ScheduledPublication {
+  ProducerValueKey producer;
+  unsigned scheduledResultIndex = 0;
+
+  bool operator==(const ScheduledPublication &other) const {
+    return producer == other.producer
+        && scheduledResultIndex == other.scheduledResultIndex;
+  }
+};
+
 struct GraphComputeBlockKey {
   Operation *op = nullptr;
   uint32_t laneStart = 0;
@@ -56,6 +66,7 @@ struct PeftClassPlan {
   SmallVector<Value> weights;
   SmallVector<Value> inputs;
   SmallVector<Type> resultTypes;
+  SmallVector<ScheduledPublication> publications;
 };
 
 struct ComputeStepTuple {
@@ -65,10 +76,38 @@ struct ComputeStepTuple {
 struct ScheduledStepPlan {
   ComputeStepTuple stepTuple;
   size_t stepIndex = 0;
-  size_t resultOffset = 0;
   size_t resultCount = 0;
 };
 
+struct ScheduledInstanceRun {
+  Operation *sourceOp = nullptr;
+  size_t firstStep = 0;
+  size_t stepCount = 0;
+  SmallVector<ComputeInstance> instances;
+  StaticIntSequence laneStarts;
+  StaticIntSequence laneCounts;
+  bool scalarSource = false;
+  bool resultfulBatchSource = false;
+};
+
+struct MaterializedProducerRef {
+  Value payload;
+  int64_t payloadLaneStart = 0;
+  int64_t payloadLaneCount = 1;
+};
+
+struct MaterializedStepValue {
+  ComputeInstance instance;
+  unsigned resultIndex = 0;
+  Value payload;
+  int64_t graphId = -1;
+  int64_t laneStart = 0;
+  int64_t laneCount = 1;
+  int64_t payloadLaneStart = 0;
+  int64_t payloadLaneCount = 1;
+  Value published;
+};
+
 struct SourceLaneAffineMapping {
   uint32_t baseLaneStart = 0;
   uint32_t laneCount = 1;
@@ -86,14 +125,12 @@ struct ScheduledMaterializationRecord {
   size_t canonicalPeftClassId = 0;
   SmallVector<size_t> cpus;
   SmallVector<ScheduledStepPlan> stepPlans;
+  SmallVector<ScheduledInstanceRun, 0> runs;
   SmallVector<GraphComputeBlockKey> computeKeys;
   SmallVector<Block *> blocks;
-};
-
-struct ScheduledComputePrintContext {
-  mlir::AsmState asmState;
-  explicit ScheduledComputePrintContext(ModuleOp module, const OpPrintingFlags &flags = OpPrintingFlags())
-    : asmState(module.getOperation(), flags) {}
+  SmallVector<Operation *> stepAnchors;
+  SmallVector<MaterializedStepValue> stepValues;
+  SmallVector<ScheduledPublication> publications;
 };
 
 inline GraphComputeBlockKey getGraphComputeBlockKey(const ComputeInstance &instance) {
@@ -133,22 +170,6 @@ inline size_t getScheduledCpuForComputeInstance(const ComputeInstance &instance,
   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());
@@ -209,16 +230,49 @@ inline size_t getComputeInstanceResultValueCount(const ComputeInstance &instance
 
 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;
+    stepPlans.push_back(ScheduledStepPlan {stepTuple, stepIndex, resultCount});
   }
   return stepPlans;
 }
 
+inline SmallVector<ScheduledInstanceRun, 0> buildScheduledInstanceRuns(
+    ArrayRef<ComputeInstance> instances) {
+  SmallVector<ScheduledInstanceRun, 0> runs;
+  for (auto [step, instance] : llvm::enumerate(instances)) {
+    auto batch = dyn_cast<SpatComputeBatch>(instance.op);
+    bool resultfulBatch = batch && batch.getNumResults() != 0;
+    bool append = resultfulBatch && !runs.empty()
+      && runs.back().resultfulBatchSource
+      && runs.back().sourceOp == instance.op
+      && runs.back().instances.back().laneStart
+           + runs.back().instances.back().laneCount == instance.laneStart;
+    if (!append) {
+      ScheduledInstanceRun run;
+      run.sourceOp = instance.op;
+      run.firstStep = step;
+      run.scalarSource = isa<SpatCompute>(instance.op);
+      run.resultfulBatchSource = resultfulBatch;
+      runs.push_back(std::move(run));
+    }
+    ScheduledInstanceRun &run = runs.back();
+    run.instances.push_back(instance);
+    run.stepCount = run.instances.size();
+  }
+  for (ScheduledInstanceRun &run : runs) {
+    SmallVector<int64_t> starts, counts;
+    for (const ComputeInstance &instance : run.instances) {
+      starts.push_back(instance.laneStart);
+      counts.push_back(instance.laneCount);
+    }
+    run.laneStarts = StaticIntSequence::fromValues(starts);
+    run.laneCounts = StaticIntSequence::fromValues(counts);
+  }
+  return runs;
+}
+
 inline bool valueTransitivelyDependsOn(Value value, Value dependency) {
   SmallVector<Value> worklist {value};
   DenseSet<Value> visited;
@@ -255,10 +309,6 @@ inline std::optional<SourceLaneAffineMapping> getSourceLaneAffineMapping(const C
   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());

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

@@ -0,0 +1,233 @@
+#include "ScheduledComputeMaterialization.hpp"
+#include "DeferredCommunicationPlanning.hpp"
+
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+
+namespace onnx_mlir::spatial {
+using namespace mlir;
+namespace {
+
+void appendUnique(SmallVectorImpl<Value> &values, Value value) {
+  if (!llvm::is_contained(values, value))
+    values.push_back(value);
+}
+
+bool requiresScheduledPublication(Value value, DenseSet<Value> &visited) {
+  if (!visited.insert(value).second)
+    return false;
+  return llvm::any_of(value.getUses(), [&](OpOperand &use) {
+    Operation *user = use.getOwner();
+    if (isa<SpatGraphCompute, SpatGraphComputeBatch,
+            SpatDeferredCommunicationOp>(user))
+      return false;
+    auto blueprint = dyn_cast<SpatBlueprintOp>(user);
+    return !blueprint || blueprint.getMode() != "fragment_assembly"
+      || requiresScheduledPublication(blueprint.getOutput(), visited);
+  });
+}
+
+bool requiresScheduledPublication(Value value) {
+  DenseSet<Value> visited;
+  return requiresScheduledPublication(value, visited);
+}
+
+} // namespace
+
+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() != sourceType.getRank() + 1 || destType.getDimSize(0) != batch.getLaneCount()
+        || destType.getElementType() != sourceType.getElementType())
+      return batch.emitOpError("graph_compute_batch result must be a leading physical-slot dimension followed by its fragment");
+    if (!llvm::equal(destType.getShape().drop_front(), sourceType.getShape()))
+      return batch.emitOpError("graph_compute_batch result trailing shape must match its published fragment");
+    if (!insert.hasUnitStride())
+      return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+    auto offsets = insert.getMixedOffsets();
+    auto sizes = insert.getMixedSizes();
+    auto strides = insert.getMixedStrides();
+    if (offsets.size() != static_cast<size_t>(destType.getRank()) || sizes.size() != static_cast<size_t>(destType.getRank())
+        || strides.size() != static_cast<size_t>(destType.getRank()))
+      return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+    if (!isa<Value>(offsets.front()) || !valueTransitivelyDependsOn(cast<Value>(offsets.front()), *batch.getLaneArgument()))
+      return batch.emitOpError("graph_compute_batch publication must select its physical slot in dimension zero");
+    for (unsigned dim = 1; dim < offsets.size(); ++dim) {
+      auto offset = dyn_cast<Attribute>(offsets[dim]);
+      auto integer = dyn_cast_or_null<IntegerAttr>(offset);
+      if (!integer || integer.getInt() != 0)
+        return batch.emitOpError("graph_compute_batch publication must have zero trailing offsets");
+    }
+    auto staticIndex = [](OpFoldResult value) -> std::optional<int64_t> {
+      auto attr = dyn_cast<Attribute>(value);
+      auto integer = dyn_cast_or_null<IntegerAttr>(attr);
+      return integer ? std::optional<int64_t>(integer.getInt()) : std::nullopt;
+    };
+    if (staticIndex(sizes.front()) != 1)
+      return batch.emitOpError("graph_compute_batch publication sizes must be [1] plus the fragment shape");
+    for (auto [size, dim] : llvm::zip_equal(ArrayRef<OpFoldResult>(sizes).drop_front(), sourceType.getShape()))
+      if (staticIndex(size) != dim)
+        return batch.emitOpError("graph_compute_batch publication sizes must be [1] plus the fragment shape");
+    return BatchFragmentSpec {spatial::getGraphBatchPhysicalResultType(fragmentLaneCount, sourceType), sourceType};
+  }
+
+  return batch.emitOpError("scheduled materialization only supports regular leading-lane graph_compute_batch fragments");
+}
+
+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();
+}
+
+LogicalResult collectPeftClassOperandsAndResults(
+  PeftClassPlan &peftClassPlan, const MergeScheduleResult &schedule) {
+  peftClassPlan.weights.clear();
+  peftClassPlan.inputs.clear();
+  peftClassPlan.resultTypes.clear();
+  peftClassPlan.publications.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)) {
+        for (auto [resultIndex, result] : llvm::enumerate(compute.getResults()))
+          if (requiresScheduledPublication(result)) {
+            peftClassPlan.publications.push_back(
+              {{instance, resultIndex}, static_cast<unsigned>(peftClassPlan.resultTypes.size())});
+            peftClassPlan.resultTypes.push_back(result.getType());
+          }
+      } else {
+        auto batch = cast<SpatComputeBatch>(instance.op);
+        for (unsigned resultIndex = 0; resultIndex < batch.getNumResults(); ++resultIndex) {
+          if (!requiresScheduledPublication(batch.getResult(resultIndex)))
+            continue;
+          auto spec = getBatchFragmentSpec(batch, resultIndex, instance.laneCount);
+          if (failed(spec))
+            return failure();
+          peftClassPlan.publications.push_back(
+            {{instance, resultIndex}, static_cast<unsigned>(peftClassPlan.resultTypes.size())});
+          peftClassPlan.resultTypes.push_back(spec->resultType);
+        }
+      }
+
+      for (Value weight : getComputeInstanceWeights(instance))
+        appendUnique(peftClassPlan.weights, weight);
+      for (Value input : getComputeInstanceInputs(instance))
+        if (!getProducerValueRef(input, &instance) && !isDeferredFragmentAssemblyInput(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 (auto [resultIndex, result] : llvm::enumerate(compute.getResults())) {
+        if (!requiresScheduledPublication(result))
+          continue;
+        auto tensorType = dyn_cast<RankedTensorType>(result.getType());
+        if (!tensorType || !tensorType.hasStaticShape())
+          return compute.emitOpError("scheduled materialization only supports static ranked tensor scalar results");
+        SmallVector<int64_t> shape {static_cast<int64_t>(peftClassPlan.cpus.size())};
+        llvm::append_range(shape, tensorType.getShape());
+        unsigned scheduledResultIndex = peftClassPlan.resultTypes.size();
+        peftClassPlan.resultTypes.push_back(RankedTensorType::get(shape, tensorType.getElementType()));
+        for (const ComputeInstance &instance : stepTuple.instances)
+          peftClassPlan.publications.push_back({{instance, resultIndex}, scheduledResultIndex});
+      }
+    } 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) {
+        if (!requiresScheduledPublication(batch.getResult(resultIndex)))
+          continue;
+        auto spec = getBatchFragmentSpec(batch, resultIndex, totalLanes);
+        if (failed(spec))
+          return failure();
+        unsigned scheduledResultIndex = peftClassPlan.resultTypes.size();
+        peftClassPlan.resultTypes.push_back(spec->resultType);
+        for (const ComputeInstance &instance : stepTuple.instances)
+          peftClassPlan.publications.push_back({{instance, resultIndex}, scheduledResultIndex});
+      }
+    }
+
+    for (const ComputeInstance &instance : stepTuple.instances) {
+      for (Value weight : getComputeInstanceWeights(instance))
+        appendUnique(peftClassPlan.weights, weight);
+      for (Value input : getComputeInstanceInputs(instance))
+        if (!getProducerValueRef(input, &instance) && !isDeferredFragmentAssemblyInput(input, instance))
+          appendUnique(peftClassPlan.inputs, input);
+    }
+  }
+
+  return success();
+}
+
+} // namespace onnx_mlir::spatial

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

@@ -1,304 +1,161 @@
 #include "ScheduledComputeReport.hpp"
-
 #include "llvm/Support/raw_os_ostream.h"
-
+#include "mlir/IR/AsmState.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();
+  op->getResult(0).printAsOperand(os, asmState);
+  if (op->getNumResults() == 1)
+    return storage;
+  if (StringRef(storage).ends_with("#0")) storage.resize(storage.size() - 2);
+  return llvm::formatv("{0}:{1}", storage, op->getNumResults()).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();
-  });
+  for (size_t begin = 0; begin < values.size();) {
+    size_t end = begin + 1;
+    if (end < values.size()) {
+      int64_t step = static_cast<int64_t>(values[end]) - static_cast<int64_t>(values[begin]);
+      while (end + 1 < values.size()
+             && static_cast<int64_t>(values[end + 1]) - static_cast<int64_t>(values[end]) == step)
+        ++end;
+    }
+    bool compressed = end - begin + 1 >= 3;
+    if (begin) os << ", ";
+    os << begin << ":" << values[begin];
+    if (compressed) os << ".." << end << ":" << values[end];
+    begin = (compressed ? end : begin) + 1;
+  }
   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 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++;
+static void printSingleCpuSteps(raw_ostream &os, ArrayRef<ComputeInstance> instances,
+                                const ScheduledMaterializationRecord &record,
+                                AsmState &asmState) {
+  for (size_t begin = 0; begin < instances.size();) {
+    const ComputeInstance &first = instances[begin];
+    size_t resultCount = getComputeInstanceResultValueCount(first);
+    size_t end = begin + 1;
+    uint32_t laneEnd = first.laneStart + first.laneCount;
+    if (!isa<SpatCompute>(first.op)) {
+      while (end < instances.size()) {
+        const ComputeInstance &next = instances[end];
+        if (isa<SpatCompute>(next.op) || next.op != first.op || next.laneStart != laneEnd
+            || next.laneCount != first.laneCount || getComputeInstanceResultValueCount(next) != resultCount)
+          break;
+        laneEnd += next.laneCount;
+        ++end;
+      }
     }
+    size_t runResults = (end - begin) * resultCount;
+    size_t publications = llvm::count_if(record.publications, [&](const ScheduledPublication &publication) {
+      return llvm::is_contained(instances.slice(begin, end - begin), publication.producer.instance);
+    });
+    os << "  " << (end - begin == 1 ? "step " : "steps [") << begin;
+    if (end - begin != 1) os << ":" << end << "]";
+    os << " payloads=" << runResults << " publications=" << publications
+       << " " << formatOperationLabel(first.op, asmState);
+    if (!isa<SpatCompute>(first.op)) os << " lanes [" << first.laneStart << ":" << laneEnd << "]";
+    if (end - begin != 1) os << " each(results=" << resultCount << ", lanes=" << first.laneCount << ")";
+    os << "\n";
+    begin = end;
   }
-  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 (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;
-
+void dumpScheduledComputeReport(ModuleOp moduleOp, func::FuncOp funcOp, const MergeScheduleResult &schedule,
+                                const llvm::MapVector<size_t, PeftClassPlan> &peftClassPlans,
+                                ArrayRef<ScheduledMaterializationRecord> records) {
+  std::fstream file = openDialectDumpFileWithExtension("spatial3_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 << "  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);
+  AsmState asmState(moduleOp.getOperation(), flags);
 
-  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();
+  size_t scalarGraph = 0, batchGraph = 0;
+  for (Operation &op : funcOp.getOps())
+    isa<SpatGraphCompute>(op) ? ++scalarGraph : batchGraph += isa<SpatGraphComputeBatch>(op);
+  size_t deferred = 0;
+  funcOp.walk([&](SpatDeferredCommunicationOp) { ++deferred; });
+  size_t scalarizedBatchInstances = 0, materializedRuns = 0;
+  size_t instancesCompacted = 0, largestRun = 0;
+  for (const ScheduledMaterializationRecord &record : records)
+    for (const ScheduledInstanceRun &run : record.runs) {
+      if (!run.resultfulBatchSource)
+        continue;
+      scalarizedBatchInstances += run.instances.size();
+      if (run.instances.size() == 1)
+        continue;
+      ++materializedRuns;
+      instancesCompacted += run.instances.size() - 1;
+      largestRun = std::max(largestRun, run.instances.size());
+    }
+  os << "Summary\n"
+     << "  graph computes: " << scalarGraph + batchGraph << " (scalar=" << scalarGraph << ", batch=" << batchGraph << ")\n"
+     << "  compute instances: " << schedule.dominanceOrderCompute.size() << "\n"
+     << "  PEFT classes: " << peftClassPlans.size() << "\n"
+     << "  scheduled ops: " << records.size() << "\n"
+     << "  deferred communications: " << deferred << "\n"
+     << "  original scalarized batch instances: " << scalarizedBatchInstances << "\n"
+     << "  materialized homogeneous runs: " << materializedRuns << "\n"
+     << "  largest run: " << largestRun << "\n"
+     << "  instances compacted: " << instancesCompacted << "\n"
+     << "  compatible runs rejected: 0\n\n"
+     << "Materialized scheduled ops\n";
+
+  for (const ScheduledMaterializationRecord &record : records) {
+    bool batch = isa<SpatScheduledComputeBatch>(record.scheduledOp);
+    const PeftClassPlan &plan = peftClassPlans.lookup(record.canonicalPeftClassId);
+    size_t steps = batch ? record.stepPlans.size() : plan.instancesByCpu.lookup(record.cpus.front()).size();
+    os << "C" << record.canonicalPeftClassId << " -> " << formatOperationLabel(record.scheduledOp, asmState) << " "
+       << record.scheduledOp->getName() << " " << (batch ? "cpus=" : "cpu=");
+    if (batch) printIndexedList(os, ArrayRef<size_t>(record.cpus));
+    else os << record.cpus.front();
+    os << " results=" << record.scheduledOp->getNumResults() << " steps=" << steps << "\n";
+    if (!batch) {
+      printSingleCpuSteps(os, plan.instancesByCpu.lookup(record.cpus.front()), record, asmState);
+      continue;
+    }
+    for (auto [stepIndex, step] : llvm::enumerate(record.stepPlans)) {
+      const ComputeInstance &source = step.stepTuple.instances.front();
+      llvm::SmallDenseSet<unsigned, 4> publishedResults;
+      for (const ScheduledPublication &publication : record.publications)
+        if (llvm::is_contained(step.stepTuple.instances, publication.producer.instance))
+          publishedResults.insert(publication.scheduledResultIndex);
+      os << "  step " << stepIndex << " payloads=" << step.resultCount
+         << " publications=" << publishedResults.size() << " "
+         << formatOperationLabel(source.op, asmState) << " lanes=" << source.laneCount << " sources=";
+      printIndexedList(os, ArrayRef<uint32_t>(collectSourceLaneStarts(step.stepTuple)));
+      os << "\n";
+      unsigned inputIndex = 0;
+      Operation *end = stepIndex + 1 < record.stepAnchors.size()
+        ? record.stepAnchors[stepIndex + 1]
+        : nullptr;
+      for (Operation *op = record.stepAnchors[stepIndex]; op && op != end; op = op->getNextNode()) {
+        auto transfer = dyn_cast<SpatDeferredCommunicationOp>(op);
+        if (!transfer) continue;
+        auto sources = transfer->getAttrOfType<DenseI64ArrayAttr>("source_operand_for_scheduled_lane");
+        auto multiSource = transfer->getAttrOfType<BoolAttr>("multi_source_payload");
+        if (multiSource && multiSource.getValue() && sources) {
+          os << "    deferred" << inputIndex << " sources=" << transfer.getSources().size() << " laneSources=";
+          printIndexedList(os, sources.asArrayRef());
+          os << "\n";
+        }
+        ++inputIndex;
+      }
+    }
+  }
 }
-
 } // namespace spatial
 } // namespace onnx_mlir

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

@@ -5,17 +5,11 @@
 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);
+void dumpScheduledComputeReport(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

@@ -84,11 +84,20 @@ LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp) {
   funcOp.walk([&](SpatDeferredCommunicationOp transfer) {
     bool ownershipValid = true;
     for (Value source : transfer.getSources()) {
-      auto result = dyn_cast<OpResult>(source);
-      if (!result || !isa<SpatGraphCompute, SpatGraphComputeBatch>(result.getOwner())) {
+      SmallVector<Value> originalSources;
+      if (auto selection = source.getDefiningOp<SpatDeferredSourceSelectOp>())
+        llvm::append_range(originalSources, selection.getSources());
+      else
+        originalSources.push_back(source);
+      for (Value originalSource : originalSources) {
+        auto result = dyn_cast<OpResult>(originalSource);
+        if (result
+            && isa<SpatGraphCompute, SpatGraphComputeBatch>(
+                 result.getOwner()))
+          continue;
         ownershipValid = false;
         diagnostics.report(transfer.getOperation(), [&](Operation *illegalOp) {
-          illegalOp->emitOpError("phase-check deferred communication source operand must be an original graph SSA result");
+          illegalOp->emitOpError("phase-check deferred communication source must resolve to original graph SSA results");
         });
       }
     }
@@ -101,18 +110,23 @@ LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp) {
     }
     if (!ownershipValid)
       return;
-    if (failed(verifyDeferredProgramContract(transfer))) {
+    if (failed(analyzeDeferredProgramTemplate(transfer))) {
       diagnostics.report(transfer.getOperation(), [&](Operation *) {});
       return;
     }
     for (Value source : transfer.getSources()) {
-      auto result = dyn_cast<OpResult>(source);
-      auto batch = result
-        ? dyn_cast<SpatGraphComputeBatch>(result.getOwner())
-        : SpatGraphComputeBatch();
-      if (batch && failed(getGraphBatchPublicationMap(
-                     batch, result.getResultNumber(), publicationCache)))
-        diagnostics.report(transfer.getOperation(), [&](Operation *) {});
+      SmallVector<Value> originalSources;
+      if (auto selection = source.getDefiningOp<SpatDeferredSourceSelectOp>())
+        llvm::append_range(originalSources, selection.getSources());
+      else
+        originalSources.push_back(source);
+      for (Value originalSource : originalSources) {
+        auto result = cast<OpResult>(originalSource);
+        auto batch = dyn_cast<SpatGraphComputeBatch>(result.getOwner());
+        if (batch && failed(getGraphBatchPublicationMap(
+                       batch, result.getResultNumber(), publicationCache)))
+          diagnostics.report(transfer.getOperation(), [&](Operation *) {});
+      }
     }
   });
 
@@ -120,34 +134,42 @@ LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp) {
   return success(!diagnostics.hasFailure());
 }
 
-LogicalResult verifyMultiCpuStepResultRouting(SpatScheduledComputeBatch scheduled,
-                                                     ArrayRef<ScheduledStepPlan> stepPlans) {
+static LogicalResult verifyStepAnchors(
+    const ScheduledMaterializationRecord &record) {
+  Region &body = record.scheduledOp->getRegion(0);
+  size_t materializedStepCount = record.runs.empty()
+    ? record.stepPlans.size() : record.runs.size();
+  if (!body.hasOneBlock()
+      || record.stepAnchors.size() != materializedStepCount)
+    return record.scheduledOp->emitOpError(
+      "scheduled compute requires one block and one anchor per step");
+  Block &block = body.front();
+  for (auto [index, anchor] : llvm::enumerate(record.stepAnchors)) {
+    if (!anchor || anchor->getBlock() != &block
+        || (index && !record.stepAnchors[index - 1]->isBeforeInBlock(anchor)))
+      return record.scheduledOp->emitOpError()
+        << "scheduled step " << index << " has an invalid anchor";
+  }
+  return success();
+}
+
+LogicalResult verifyMultiCpuStepResultRouting(
+    SpatScheduledComputeBatch scheduled) {
   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");
+  if (!scheduled.getBody().hasOneBlock())
     return failure();
-  }
 
+  Block &block = scheduled.getBody().front();
+  if (scheduled.getNumResults() == 0)
+    return success(isa<SpatYieldOp>(block.getTerminator()));
   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;
-    }
-
+  auto inParallel = dyn_cast<SpatInParallelOp>(block.getTerminator());
+  if (!inParallel)
+    diagnostics.report(scheduled.getOperation(), [&](Operation *illegalOp) {
+      illegalOp->emitOpError("scheduled batch is missing spat.in_parallel");
+    });
+  else
     for (Operation &op : inParallel.getRegion().front()) {
       auto insert = dyn_cast<tensor::ParallelInsertSliceOp>(&op);
       if (!insert)
@@ -155,41 +177,23 @@ LogicalResult verifyMultiCpuStepResultRouting(SpatScheduledComputeBatch schedule
       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";
+          illegalOp->emitOpError(
+            "scheduled batch 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();
+          illegalOp->emitOpError()
+            << "scheduled batch 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) {
@@ -204,27 +208,25 @@ LogicalResult verifyMultiCpuStepResultRouting(SpatScheduledComputeBatch schedule
 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) {
+  if (!scheduled.getBody().hasOneBlock())
+    return scheduled.emitOpError(
+      "scheduled batch local fragment verification requires one block");
+  Block &block = scheduled.getBody().front();
+  if (scheduled.getNumResults() == 0)
+    return success(isa<SpatYieldOp>(block.getTerminator()));
+  Value scheduledLane = block.getArgument(0);
+  auto inParallel = dyn_cast<SpatInParallelOp>(block.getTerminator());
+  if (!inParallel)
+    return scheduled.emitOpError("scheduled batch is missing spat.in_parallel");
+  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) {
+  block.walk([&](Operation *op) {
       if (op == block.getTerminator())
         return;
       if (isFinalScheduledOutputInsert(op)) {
@@ -266,13 +268,11 @@ LogicalResult verifyMultiCpuLocalFragmentOffsets(SpatScheduledComputeBatch sched
         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;
+            << " in the shared scheduled block";
         });
         break;
       }
-    });
-  }
-
+  });
   diagnostics.emitSuppressedSummary(scheduled.getOperation(),
                                     "scheduled batch local fragment offset verification failed");
   return success(!diagnostics.hasFailure());
@@ -281,10 +281,49 @@ LogicalResult verifyMultiCpuLocalFragmentOffsets(SpatScheduledComputeBatch sched
 
 LogicalResult verifyScheduledMaterializationRecords(ArrayRef<ScheduledMaterializationRecord> materializedSchedules) {
   for (const ScheduledMaterializationRecord &record : materializedSchedules) {
+    if (failed(verifyStepAnchors(record)))
+      return failure();
+    size_t expectedValues = 0;
+    for (const ScheduledStepPlan &step : record.stepPlans)
+      expectedValues += step.stepTuple.instances.size() * step.resultCount;
+    if (record.stepValues.size() != expectedValues)
+      return record.scheduledOp->emitOpError(
+        "scheduled materialization record has the wrong step-value count");
+    DenseMap<ProducerValueKey, const MaterializedStepValue *> values;
+    for (const MaterializedStepValue &value : record.stepValues) {
+      if (!value.payload || value.graphId < 0 || value.laneStart < 0
+          || value.laneCount <= 0 || value.payloadLaneStart < 0
+          || value.payloadLaneCount <= 0
+          || !values.try_emplace({value.instance, value.resultIndex}, &value).second)
+        return record.scheduledOp->emitOpError(
+          "scheduled materialization record has an invalid or duplicate step value");
+    }
+    SmallVector<unsigned> publicationCounts(record.scheduledOp->getNumResults());
+    for (const ScheduledPublication &publication : record.publications) {
+      auto value = values.find(publication.producer);
+      if (value == values.end()
+          || publication.scheduledResultIndex >= publicationCounts.size()
+          || value->second->published
+               != record.scheduledOp->getResult(publication.scheduledResultIndex))
+        return record.scheduledOp->emitOpError(
+          "scheduled publication does not map to its recorded step value");
+      ++publicationCounts[publication.scheduledResultIndex];
+    }
+    if (llvm::any_of(publicationCounts,
+                     [](unsigned count) { return count == 0; }))
+      return record.scheduledOp->emitOpError(
+        "scheduled result has no declared publication owner");
+    for (const MaterializedStepValue &value : record.stepValues)
+      if (value.published
+          && !llvm::is_contained(record.publications,
+               ScheduledPublication{{value.instance, value.resultIndex},
+                 cast<OpResult>(value.published).getResultNumber()}))
+        return record.scheduledOp->emitOpError(
+          "scheduled step value has an undeclared publication");
     auto scheduled = dyn_cast<SpatScheduledComputeBatch>(record.scheduledOp);
     if (!scheduled)
       continue;
-    if (failed(verifyMultiCpuStepResultRouting(scheduled, record.stepPlans)))
+    if (failed(verifyMultiCpuStepResultRouting(scheduled)))
       return failure();
     if (failed(verifyMultiCpuLocalFragmentOffsets(scheduled)))
       return failure();
@@ -292,5 +331,22 @@ LogicalResult verifyScheduledMaterializationRecords(ArrayRef<ScheduledMaterializ
   return success();
 }
 
+LogicalResult verifyScheduledResultsLive(
+    ArrayRef<ScheduledMaterializationRecord> materializedSchedules) {
+  pim::CappedDiagnosticReporter diagnostics;
+  for (const ScheduledMaterializationRecord &record : materializedSchedules)
+    for (OpResult result : record.scheduledOp->getResults())
+      if (result.use_empty())
+        diagnostics.report(record.scheduledOp, [&](Operation *op) {
+          op->emitOpError() << "scheduled result " << result.getResultNumber()
+                            << " has no use after communication realization";
+        });
+  if (!materializedSchedules.empty())
+    diagnostics.emitSuppressedSummary(
+      materializedSchedules.front().scheduledOp,
+      "scheduled publication liveness verification failed");
+  return success(!diagnostics.hasFailure());
+}
+
 } // namespace spatial
 } // namespace onnx_mlir

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

@@ -13,10 +13,10 @@ LogicalResult verifyMaterializedScheduleMapping(
   ArrayRef<ScheduledMaterializationRecord> materializedSchedules);
 
 LogicalResult verifyDeferredTransferPhase1Invariants(func::FuncOp funcOp);
-LogicalResult verifyMultiCpuStepResultRouting(SpatScheduledComputeBatch scheduled,
-                                              ArrayRef<ScheduledStepPlan> stepPlans);
+LogicalResult verifyMultiCpuStepResultRouting(SpatScheduledComputeBatch scheduled);
 LogicalResult verifyMultiCpuLocalFragmentOffsets(SpatScheduledComputeBatch scheduled);
 LogicalResult verifyScheduledMaterializationRecords(ArrayRef<ScheduledMaterializationRecord> materializedSchedules);
+LogicalResult verifyScheduledResultsLive(ArrayRef<ScheduledMaterializationRecord> materializedSchedules);
 
 } // namespace spatial
 } // namespace onnx_mlir

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

@@ -16,13 +16,6 @@ namespace spatial {
 
 namespace {
 
-MergeSchedulerKind getSchedulerKind() {
-  switch (pimMergeScheduler.getValue()) {
-  case MergeSchedulerPeft: return MergeSchedulerKind::Peft;
-  }
-  llvm_unreachable("unknown merge scheduler kind");
-}
-
 void verifySchedule(const ComputeGraph& graph,
                     const MergeScheduleResult& result,
                     unsigned long crossbarCapacity,
@@ -107,22 +100,19 @@ MergeScheduleResult MergeSchedulingAnalysis::run() {
   if (!verifyAcyclic(graph))
     llvm::report_fatal_error("merge scheduling: compute graph is cyclic");
 
-  MergeSchedulingOptions options;
-  options.kind = getSchedulerKind();
+  size_t processorCount = 0;
   if (coresCount.getValue() > 0)
-    options.processorCount = static_cast<size_t>(coresCount.getValue());
+    processorCount = static_cast<size_t>(coresCount.getValue());
 
-  MergeScheduleResult schedule;
-  if (options.kind == MergeSchedulerKind::Peft) {
-    schedule = runPeftScheduler(graph,
-                                PeftScheduleOptions {options.processorCount,
-                                                     static_cast<unsigned long>(crossbarCountInCore.getValue()),
-                                                     entryOp->getContext()});
-  }
+  MergeScheduleResult schedule = runPeftScheduler(
+    graph, PeftScheduleOptions {
+      processorCount,
+      static_cast<unsigned long>(crossbarCountInCore.getValue()),
+      entryOp->getContext()});
   verifySchedule(graph,
                  schedule,
                  static_cast<unsigned long>(crossbarCountInCore.getValue()),
-                 options.processorCount);
+                 processorCount);
   return schedule;
 }
 

src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.hpp

@@ -2,22 +2,11 @@
 
 #include "mlir/IR/Operation.h"
 
-#include <cstddef>
-
 #include "MergeSchedule.hpp"
 
 namespace onnx_mlir {
 namespace spatial {
 
-enum class MergeSchedulerKind {
-  Peft,
-};
-
-struct MergeSchedulingOptions {
-  MergeSchedulerKind kind = MergeSchedulerKind::Peft;
-  size_t processorCount = 0;
-};
-
 class MergeSchedulingAnalysis {
 public:
   explicit MergeSchedulingAnalysis(mlir::Operation* op);

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

@@ -418,53 +418,34 @@ resolveProducerSourcesForCsv(const ResolvedProducer& producer,
   return sources;
 }
 
-FailureOr<SmallVector<int64_t>> getIntegerValues(Operation* op, StringRef name) {
-  Attribute attr = op->getAttr(name);
-  if (auto array = dyn_cast_or_null<DenseI64ArrayAttr>(attr))
-    return SmallVector<int64_t>(array.asArrayRef());
-  if (auto elements = dyn_cast_or_null<DenseIntElementsAttr>(attr))
-    return SmallVector<int64_t>(elements.getValues<int64_t>());
-  return op->emitOpError() << "expected " << name << " integer array for Spatial dataflow report";
-}
-
 FailureOr<ScheduledNodeByGraphLane>
 buildScheduledNodesByGraphLane(const DenseMap<Operation*, TopLevelOpInfo>& topLevelInfo,
                                const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
-                               const DenseMap<int64_t, Operation*>& graphOpsById) {
+                               ArrayRef<ScheduledMaterializationRecord> records) {
   ScheduledNodeByGraphLane nodesByGraphLane;
-  for (const auto& entry : topLevelInfo) {
-    Operation* scheduledOp = entry.first;
-    auto sourceIds = getIntegerValues(scheduledOp, "scheduled.step_source_ids");
-    auto sourceStarts = getIntegerValues(scheduledOp, "scheduled.source_lane_starts");
-    auto sourceCounts = getIntegerValues(scheduledOp, "scheduled.source_lane_counts");
-    if (failed(sourceIds) || failed(sourceStarts) || failed(sourceCounts))
-      return failure();
-
-    uint32_t scheduledLaneCount = 1;
-    if (auto batch = dyn_cast<SpatScheduledComputeBatch>(scheduledOp))
-      scheduledLaneCount = static_cast<uint32_t>(batch.getLaneCount());
-    size_t expectedEntries = sourceIds->size() * scheduledLaneCount;
-    if (sourceStarts->size() != expectedEntries || sourceCounts->size() != expectedEntries)
-      return scheduledOp->emitOpError("inconsistent scheduling provenance arrays for Spatial dataflow report");
-
-    for (auto [step, graphId] : llvm::enumerate(*sourceIds)) {
-      auto graphIt = graphOpsById.find(graphId);
-      if (graphIt == graphOpsById.end())
-        return scheduledOp->emitOpError() << "references unknown scheduled graph id " << graphId;
-      bool graphIsBatch = isa<SpatGraphComputeBatch>(graphIt->second);
-      for (uint32_t scheduledLane = 0; scheduledLane < scheduledLaneCount; ++scheduledLane) {
-        auto nodeIt = expandedNodes.find({scheduledOp, scheduledLane});
+  for (const ScheduledMaterializationRecord &record : records) {
+    if (!topLevelInfo.count(record.scheduledOp))
+      return record.scheduledOp->emitOpError(
+        "missing scheduled node for Spatial dataflow report");
+    for (const ScheduledStepPlan &step : record.stepPlans)
+      for (auto [scheduledLane, instance] :
+           llvm::enumerate(step.stepTuple.instances)) {
+        auto nodeIt = expandedNodes.find(
+          {record.scheduledOp, static_cast<uint32_t>(scheduledLane)});
         if (nodeIt == expandedNodes.end())
           continue;
-        size_t index = step * scheduledLaneCount + scheduledLane;
-        int64_t start = graphIsBatch ? (*sourceStarts)[index] : 0;
-        int64_t count = graphIsBatch ? (*sourceCounts)[index] : 1;
-        if (start < 0 || count < 0)
-          return scheduledOp->emitOpError("negative scheduling provenance range for Spatial dataflow report");
-        for (int64_t lane = start; lane < start + count; ++lane)
-          nodesByGraphLane[{graphId, static_cast<uint32_t>(lane)}] = nodeIt->second;
+        auto graphId = instance.op->getAttrOfType<IntegerAttr>(
+          "scheduled.graph_id");
+        if (!graphId)
+          return instance.op->emitOpError(
+            "missing graph identity for Spatial dataflow report");
+        uint32_t laneStart = isa<SpatGraphComputeBatch>(instance.op)
+          ? instance.laneStart : 0;
+        uint32_t laneCount = isa<SpatGraphComputeBatch>(instance.op)
+          ? instance.laneCount : 1;
+        for (uint32_t lane = laneStart; lane < laneStart + laneCount; ++lane)
+          nodesByGraphLane[{graphId.getInt(), lane}] = nodeIt->second;
       }
-    }
   }
   return nodesByGraphLane;
 }
@@ -489,13 +470,10 @@ emitScheduledPlanningEdges(std::fstream& edgesFile,
                            func::FuncOp func,
                            const DenseMap<Operation*, TopLevelOpInfo>& topLevelInfo,
                            const DenseMap<std::pair<Operation*, uint32_t>, ExpandedNodeInfo>& expandedNodes,
+                           ArrayRef<ScheduledMaterializationRecord> records,
                            StringRef stage) {
-  DenseMap<int64_t, Operation*> graphOpsById;
-  for (Operation& op : func.getBody().front())
-    if (auto graphId = op.getAttrOfType<IntegerAttr>("scheduled.graph_id"))
-      graphOpsById[graphId.getInt()] = &op;
-
-  auto nodesByGraphLane = buildScheduledNodesByGraphLane(topLevelInfo, expandedNodes, graphOpsById);
+  auto nodesByGraphLane = buildScheduledNodesByGraphLane(
+    topLevelInfo, expandedNodes, records);
   if (failed(nodesByGraphLane))
     return failure();
 
@@ -734,7 +712,7 @@ LogicalResult emitExplicitChannelEdges(std::fstream& edgesFile,
   return success();
 }
 
-LogicalResult exportGraph(func::FuncOp func, StringRef reportName) {
+LogicalResult exportGraph(func::FuncOp func, StringRef reportName, StringRef stage) {
   std::fstream nodesFile = openDialectDumpFileWithExtension((reportName + ".nodes").str(), "/reports", "csv");
   std::fstream edgesFile = openDialectDumpFileWithExtension((reportName + ".edges").str(), "/reports", "csv");
   if (!nodesFile.is_open() || !edgesFile.is_open())
@@ -772,10 +750,12 @@ LogicalResult exportGraph(func::FuncOp func, StringRef reportName) {
     addBatchNodeRows(nodesFile, expandedNodes, *info, batch, laneCoreIds, &asmState);
   }
 
-  return emitDataEdges<SpatGraphCompute, SpatGraphComputeBatch>(edgesFile, topLevelInfo, expandedNodes, "spatial1");
+  return emitDataEdges<SpatGraphCompute, SpatGraphComputeBatch>(edgesFile, topLevelInfo, expandedNodes, stage);
 }
 
-LogicalResult exportScheduled(func::FuncOp func, StringRef reportName, StringRef stage) {
+LogicalResult exportScheduled(func::FuncOp func,
+                              ArrayRef<ScheduledMaterializationRecord> records,
+                              StringRef reportName, StringRef stage) {
   std::fstream nodesFile = openDialectDumpFileWithExtension((reportName + ".nodes").str(), "/reports", "csv");
   std::fstream edgesFile = openDialectDumpFileWithExtension((reportName + ".edges").str(), "/reports", "csv");
   if (!nodesFile.is_open() || !edgesFile.is_open())
@@ -820,8 +800,9 @@ LogicalResult exportScheduled(func::FuncOp func, StringRef reportName, StringRef
     addBatchNodeRows(nodesFile, expandedNodes, *info, batch, laneCoreIds, &asmState);
   }
 
-  if (stage == "spatial2")
-    return emitScheduledPlanningEdges(edgesFile, func, topLevelInfo, expandedNodes, stage);
+  if (stage == "spatial3")
+    return emitScheduledPlanningEdges(
+      edgesFile, func, topLevelInfo, expandedNodes, records, stage);
   if (failed(
         emitDataEdges<SpatScheduledCompute, SpatScheduledComputeBatch>(edgesFile, topLevelInfo, expandedNodes, stage)))
     return failure();
@@ -870,6 +851,7 @@ SpatialDataflowExportStage getSpatialDataflowExportStage() {
   case SpatialDataflowExportSpatial1: return SpatialDataflowExportStage::Spatial1;
   case SpatialDataflowExportSpatial2: return SpatialDataflowExportStage::Spatial2;
   case SpatialDataflowExportSpatial3: return SpatialDataflowExportStage::Spatial3;
+  case SpatialDataflowExportSpatial4: return SpatialDataflowExportStage::Spatial4;
   case SpatialDataflowExportAll:      return SpatialDataflowExportStage::All;
   }
   llvm_unreachable("unknown spatial dataflow export mode");
@@ -881,17 +863,20 @@ bool shouldExportSpatialDataflowStage(SpatialDataflowExportStage mode, SpatialDa
   case SpatialDataflowExportStage::Spatial1: return stage == SpatialDataflowExportStage::Spatial1;
   case SpatialDataflowExportStage::Spatial2: return stage == SpatialDataflowExportStage::Spatial2;
   case SpatialDataflowExportStage::Spatial3: return stage == SpatialDataflowExportStage::Spatial3;
+  case SpatialDataflowExportStage::Spatial4: return stage == SpatialDataflowExportStage::Spatial4;
   case SpatialDataflowExportStage::All:      return stage != SpatialDataflowExportStage::None;
   }
   return false;
 }
 
-LogicalResult exportSpatialDataflowCsvGraph(func::FuncOp func, StringRef reportName) {
-  return exportGraph(func, reportName);
+LogicalResult exportSpatialDataflowCsvGraph(func::FuncOp func, StringRef reportName, StringRef stage) {
+  return exportGraph(func, reportName, stage);
 }
 
-LogicalResult exportSpatialDataflowCsvScheduled(func::FuncOp func, StringRef reportName, StringRef stage) {
-  return exportScheduled(func, reportName, stage);
+LogicalResult exportSpatialDataflowCsvScheduled(
+    func::FuncOp func, ArrayRef<ScheduledMaterializationRecord> records,
+    StringRef reportName, StringRef stage) {
+  return exportScheduled(func, records, reportName, stage);
 }
 
 } // namespace spatial

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

@@ -5,6 +5,8 @@
 
 #include "llvm/ADT/StringRef.h"
 
+#include "ScheduledComputeMaterialization.hpp"
+
 namespace onnx_mlir {
 namespace spatial {
 
@@ -13,14 +15,20 @@ enum class SpatialDataflowExportStage {
   Spatial1,
   Spatial2,
   Spatial3,
+  Spatial4,
   All,
 };
 
 SpatialDataflowExportStage getSpatialDataflowExportStage();
 
-mlir::LogicalResult exportSpatialDataflowCsvGraph(mlir::func::FuncOp func, llvm::StringRef reportName);
+mlir::LogicalResult exportSpatialDataflowCsvGraph(mlir::func::FuncOp func,
+                                                  llvm::StringRef reportName,
+                                                  llvm::StringRef stage = "spatial1");
 mlir::LogicalResult
-exportSpatialDataflowCsvScheduled(mlir::func::FuncOp func, llvm::StringRef reportName, llvm::StringRef stage);
+exportSpatialDataflowCsvScheduled(mlir::func::FuncOp func,
+                                  llvm::ArrayRef<ScheduledMaterializationRecord> records,
+                                  llvm::StringRef reportName,
+                                  llvm::StringRef stage);
 
 bool shouldExportSpatialDataflowStage(SpatialDataflowExportStage mode, SpatialDataflowExportStage stage);
 

src/PIM/Dialect/Spatial/Transforms/TrivialGraphComputeMergePass.cpp

@@ -0,0 +1,508 @@
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/IRMapping.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SetVector.h"
+#include "llvm/Support/FormatVariadic.h"
+#include "llvm/Support/raw_os_ostream.h"
+
+#include <fstream>
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
+#include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.hpp"
+#include "src/Accelerators/PIM/Pass/PIMPasses.h"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace spatial {
+namespace {
+
+template <typename ComputeOp>
+static bool hasOnlyStructuralAttrs(ComputeOp op) {
+  return llvm::all_of(op->getAttrs(), [&](NamedAttribute attr) {
+    return attr.getName() == op.getOperandSegmentSizesAttrName()
+        || (isa<SpatGraphComputeBatch>(op.getOperation())
+            && attr.getName() == cast<SpatGraphComputeBatch>(op.getOperation()).getLaneCountAttrName());
+  });
+}
+
+static bool hasCapacityFor(Operation* producer, Operation* consumer) {
+  CrossbarUsage producerWeights = collectDistinctCrossbarWeights(producer);
+  CrossbarUsage consumerWeights = collectDistinctCrossbarWeights(consumer);
+  return getCrossbarUnionSize(producerWeights, consumerWeights) <= static_cast<size_t>(crossbarCountInCore.getValue());
+}
+
+struct TrivialGraphMergeStats {
+  size_t scalarBefore = 0;
+  size_t batchBefore = 0;
+  size_t scalarAfter = 0;
+  size_t batchAfter = 0;
+  size_t scalarProducerConsumerMerges = 0;
+  size_t batchProducerConsumerMerges = 0;
+  size_t leadingUnitNormalizationFolds = 0;
+};
+
+static std::pair<size_t, size_t> countGraphComputes(ModuleOp module) {
+  size_t scalar = 0, batch = 0;
+  module.walk([&](Operation *op) {
+    scalar += isa<SpatGraphCompute>(op);
+    batch += isa<SpatGraphComputeBatch>(op);
+  });
+  return {scalar, batch};
+}
+
+static void dumpTrivialMergeReport(const TrivialGraphMergeStats &stats,
+                                   size_t largestBatchLaneCount) {
+  std::fstream file = openDialectDumpFileWithExtension("spatial2_trivial_merged", "/reports", "txt");
+  if (!file.is_open())
+    return;
+  llvm::raw_os_ostream os(file);
+  size_t before = stats.scalarBefore + stats.batchBefore;
+  size_t after = stats.scalarAfter + stats.batchAfter;
+  size_t removed = before - after;
+  double percentage = before ? 100.0 * removed / before : 0.0;
+  os << "Summary\n"
+     << "  graph computes: " << before << " -> " << after << "\n"
+     << "    removed: " << removed << " (" << llvm::formatv("{0:F2}", percentage) << "%)\n"
+     << "  scalar: " << stats.scalarBefore << " -> " << stats.scalarAfter << "\n"
+     << "  batch: " << stats.batchBefore << " -> " << stats.batchAfter << "\n"
+     << "  transformations:\n"
+     << "    scalar producer-consumer merges: " << stats.scalarProducerConsumerMerges << "\n"
+     << "    batch producer-consumer merges: " << stats.batchProducerConsumerMerges << "\n"
+     << "    leading-unit normalization folds: " << stats.leadingUnitNormalizationFolds << "\n\n"
+     << "Resulting graph\n"
+     << "  graph computes: " << after << "\n"
+     << "  scalar: " << stats.scalarAfter << "\n"
+     << "  batch: " << stats.batchAfter << "\n"
+     << "  largest batch lane count: " << largestBatchLaneCount << "\n";
+}
+
+template <typename ProducerOp, typename ConsumerOp>
+static bool isExclusivelyConsumedBy(ProducerOp producer, ConsumerOp consumer) {
+  bool hasDependency = false;
+  for (Value result : producer.getResults()) {
+    for (OpOperand& use : result.getUses()) {
+      if (use.getOwner() != consumer.getOperation() || !llvm::is_contained(consumer.getInputs(), result))
+        return false;
+      hasDependency = true;
+    }
+  }
+  return hasDependency;
+}
+
+static bool hasNoNestedArgumentCaptures(SpatGraphCompute compute) {
+  Block &body = compute.getBody().front();
+  return llvm::all_of(body.getArguments(), [&](BlockArgument argument) {
+    return llvm::all_of(argument.getUsers(), [&](Operation *user) { return user->getBlock() == &body; });
+  });
+}
+
+template <typename ProducerOp, typename ConsumerOp>
+static void collectExternalOperands(ProducerOp producer,
+                                    ConsumerOp consumer,
+                                    llvm::SetVector<Value>& weights,
+                                    llvm::SetVector<Value>& inputs) {
+  weights.insert(producer.getWeights().begin(), producer.getWeights().end());
+  weights.insert(consumer.getWeights().begin(), consumer.getWeights().end());
+  auto appendInput = [&](Value value) {
+    if (value.getDefiningOp() != producer.getOperation() && !weights.contains(value))
+      inputs.insert(value);
+  };
+  llvm::for_each(producer.getInputs(), appendInput);
+  llvm::for_each(consumer.getInputs(), appendInput);
+}
+
+template <typename OldOp, typename NewOp>
+static void mapExternalArguments(OldOp oldOp, NewOp newOp, IRMapping& mapper, bool mapForwardedInputs = true) {
+  for (auto [index, operand] : llvm::enumerate(oldOp.getWeights())) {
+    auto oldArg = oldOp.getWeightArgument(index);
+    auto newOperand = llvm::find(newOp.getWeights(), operand);
+    assert(oldArg && newOperand != newOp.getWeights().end());
+    mapper.map(*oldArg, *newOp.getWeightArgument(std::distance(newOp.getWeights().begin(), newOperand)));
+  }
+  for (auto [index, operand] : llvm::enumerate(oldOp.getInputs())) {
+    auto oldArg = oldOp.getInputArgument(index);
+    assert(oldArg);
+    if (mapper.contains(operand)) {
+      if (mapForwardedInputs)
+        mapper.map(*oldArg, mapper.lookup(operand));
+      continue;
+    }
+    auto newOperand = llvm::find(newOp.getInputs(), operand);
+    assert(oldArg && newOperand != newOp.getInputs().end());
+    mapper.map(*oldArg, *newOp.getInputArgument(std::distance(newOp.getInputs().begin(), newOperand)));
+  }
+}
+
+struct MergeTrivialScalarComputes : OpRewritePattern<SpatGraphCompute> {
+  MergeTrivialScalarComputes(MLIRContext *context, TrivialGraphMergeStats *stats)
+      : OpRewritePattern(context), stats(stats) {}
+
+  LogicalResult matchAndRewrite(SpatGraphCompute consumer, PatternRewriter& rewriter) const override {
+    SpatGraphCompute producer;
+    for (Value input : consumer.getInputs()) {
+      auto candidate = input.getDefiningOp<SpatGraphCompute>();
+      if (candidate && candidate->getBlock() == consumer->getBlock() && hasOnlyStructuralAttrs(candidate)
+          && hasOnlyStructuralAttrs(consumer) && isExclusivelyConsumedBy(candidate, consumer)
+          && hasCapacityFor(candidate, consumer) && hasNoNestedArgumentCaptures(candidate)
+          && hasNoNestedArgumentCaptures(consumer)) {
+        producer = candidate;
+        break;
+      }
+    }
+    if (!producer)
+      return failure();
+
+    llvm::SetVector<Value> weights, inputs;
+    collectExternalOperands(producer, consumer, weights, inputs);
+    rewriter.setInsertionPoint(consumer);
+    SpatGraphCompute merged = createEmptySpatGraphCompute(
+      rewriter, consumer.getLoc(), consumer.getResultTypes(), weights.getArrayRef(), inputs.getArrayRef());
+
+    IRMapping mapper;
+    mapExternalArguments(producer, merged, mapper);
+    for (Operation& op : producer.getBody().front().without_terminator())
+      rewriter.clone(op, mapper);
+    auto producerYield = cast<SpatYieldOp>(producer.getBody().front().getTerminator());
+    for (auto [result, yielded] : llvm::zip(producer.getResults(), producerYield.getOutputs()))
+      mapper.map(result, mapper.lookupOrDefault(yielded));
+
+    mapExternalArguments(consumer, merged, mapper);
+    for (Operation& op : consumer.getBody().front())
+      rewriter.clone(op, mapper);
+
+    rewriter.replaceOp(consumer, merged.getResults());
+    rewriter.eraseOp(producer);
+    ++stats->scalarProducerConsumerMerges;
+    return success();
+  }
+
+private:
+  TrivialGraphMergeStats *stats;
+};
+
+static bool isLaneIndex(Value value, Value lane, int64_t laneCount) {
+  if (!value)
+    return false;
+  if (value == lane)
+    return true;
+  auto apply = value.getDefiningOp<affine::AffineApplyOp>();
+  if (!apply || apply.getMapOperands().size() != 1 || apply.getMapOperands().front() != lane)
+    return false;
+  AffineMap map = apply.getAffineMap();
+  if (map.getNumDims() != 1 || map.getNumSymbols() != 0 || map.getNumResults() != 1)
+    return false;
+  AffineExpr expression = map.getResult(0);
+  if (auto dim = dyn_cast<AffineDimExpr>(expression))
+    return dim.getPosition() == 0;
+  auto modulo = dyn_cast<AffineBinaryOpExpr>(expression);
+  if (!modulo || modulo.getKind() != AffineExprKind::Mod)
+    return false;
+  auto dim = dyn_cast<AffineDimExpr>(modulo.getLHS());
+  auto divisor = dyn_cast<AffineConstantExpr>(modulo.getRHS());
+  return dim && dim.getPosition() == 0
+      && divisor && divisor.getValue() >= laneCount;
+}
+
+static bool isDirectLaneSlot(ArrayRef<OpFoldResult> offsets,
+                             ArrayRef<OpFoldResult> sizes,
+                             ArrayRef<OpFoldResult> strides,
+                             Value lane,
+                             RankedTensorType physicalType) {
+  size_t rank = physicalType.getRank();
+  if (offsets.size() != rank || sizes.size() != rank || strides.size() != rank
+      || !isLaneIndex(dyn_cast<Value>(offsets.front()), lane, physicalType.getDimSize(0)))
+    return false;
+  for (size_t dim = 0; dim < rank; ++dim) {
+    int64_t expectedOffset = 0;
+    int64_t expectedSize = dim == 0 ? 1 : physicalType.getDimSize(dim);
+    if ((dim != 0 && getConstantIntValue(offsets[dim]) != expectedOffset)
+        || getConstantIntValue(sizes[dim]) != expectedSize || getConstantIntValue(strides[dim]) != 1)
+      return false;
+  }
+  return true;
+}
+
+static FailureOr<SmallVector<Value>> collectPublishedFragments(SpatGraphComputeBatch producer) {
+  auto terminator = dyn_cast<SpatInParallelOp>(producer.getBody().front().getTerminator());
+  auto lane = producer.getLaneArgument();
+  if (!terminator || !lane)
+    return failure();
+  SmallVector<Value> fragments(producer.getNumResults());
+  for (Operation& op : terminator.getRegion().front()) {
+    auto insert = dyn_cast<tensor::ParallelInsertSliceOp>(op);
+    auto destination = insert ? dyn_cast<BlockArgument>(insert.getDest()) : BlockArgument();
+    if (!insert || !destination)
+      return failure();
+    unsigned resultIndex = 0;
+    while (resultIndex < producer.getNumResults() && producer.getOutputArgument(resultIndex) != destination)
+      ++resultIndex;
+    if (resultIndex == producer.getNumResults() || fragments[resultIndex])
+      return failure();
+    auto physicalType = dyn_cast<RankedTensorType>(producer.getResult(resultIndex).getType());
+    auto fragmentType = dyn_cast<RankedTensorType>(insert.getSource().getType());
+    auto expectedFragmentType = getGraphBatchFragmentType(physicalType, producer.getLaneCount());
+    if (!physicalType || !fragmentType || failed(expectedFragmentType) || *expectedFragmentType != fragmentType
+        || !isDirectLaneSlot(
+          insert.getMixedOffsets(), insert.getMixedSizes(), insert.getMixedStrides(), *lane, physicalType))
+      return failure();
+    fragments[resultIndex] = insert.getSource();
+  }
+  if (!llvm::all_of(fragments, [](Value value) { return value; }))
+    return failure();
+  return fragments;
+}
+
+static bool matchLeadingUnitNormalization(SpatGraphComputeBatch producer, SpatGraphCompute consumer) {
+  if (producer.getNumResults() != 1 || !consumer.getWeights().empty()
+      || consumer.getInputs().size() != 1 || consumer.getInputs().front() != producer.getResult(0)
+      || consumer.getNumResults() != 1)
+    return false;
+  auto yield = dyn_cast<SpatYieldOp>(consumer.getBody().front().getTerminator());
+  auto loopResult = yield && yield.getOutputs().size() == 1
+    ? dyn_cast<OpResult>(yield.getOutputs().front())
+    : OpResult();
+  auto loop = loopResult ? dyn_cast<scf::ForOp>(loopResult.getOwner()) : scf::ForOp();
+  auto targetType = dyn_cast<RankedTensorType>(consumer.getResult(0).getType());
+  if (!loop || llvm::range_size(consumer.getBody().front().without_terminator()) != 2
+      || loop->getBlock() != &consumer.getBody().front() || loopResult.getResultNumber() != 0
+      || loop.getNumResults() != 1 || loop.getResult(0).getType() != targetType
+      || !targetType || !targetType.hasStaticShape()
+      || getConstantIntValue(loop.getLowerBound()) != 0
+      || getConstantIntValue(loop.getUpperBound()) != producer.getLaneCount()
+      || getConstantIntValue(loop.getStep()) != 1
+      || llvm::range_size(loop.getBody()->without_terminator()) != 2)
+    return false;
+  auto empty = loop.getInitArgs().front().getDefiningOp<tensor::EmptyOp>();
+  auto bodyIt = loop.getBody()->begin();
+  auto extract = dyn_cast<tensor::ExtractSliceOp>(&*bodyIt++);
+  auto insert = dyn_cast<tensor::InsertSliceOp>(&*bodyIt);
+  auto loopYield = dyn_cast<scf::YieldOp>(loop.getBody()->getTerminator());
+  auto sourceType = dyn_cast<RankedTensorType>(producer.getResult(0).getType());
+  auto sourceFragmentType = sourceType
+    ? getGraphBatchFragmentType(sourceType, producer.getLaneCount())
+    : FailureOr<RankedTensorType>(failure());
+  auto targetFragmentType = getGraphBatchFragmentType(targetType, producer.getLaneCount());
+  return empty && empty->getBlock() == &consumer.getBody().front() && empty.getType() == targetType
+      && extract && insert && loopYield && loopYield.getResults().size() == 1
+      && extract.getSource() == consumer.getInputArgument(0)
+      && insert.getSource() == extract.getResult() && insert.getDest() == loop.getRegionIterArgs().front()
+      && loopYield.getResults().front() == insert.getResult()
+      && succeeded(sourceFragmentType) && succeeded(targetFragmentType)
+      && extract.getType() == *sourceFragmentType
+      && sourceFragmentType->getRank() == targetFragmentType->getRank() + 1
+      && sourceFragmentType->getDimSize(0) == 1
+      && sourceFragmentType->getShape().drop_front() == targetFragmentType->getShape()
+      && sourceFragmentType->getElementType() == targetFragmentType->getElementType()
+      && isDirectLaneSlot(extract.getMixedOffsets(), extract.getMixedSizes(), extract.getMixedStrides(),
+                          loop.getInductionVar(), sourceType)
+      && isDirectLaneSlot(insert.getMixedOffsets(), insert.getMixedSizes(), insert.getMixedStrides(),
+                          loop.getInductionVar(), targetType);
+}
+
+struct FoldBatchLeadingUnitNormalization : OpRewritePattern<SpatGraphCompute> {
+  FoldBatchLeadingUnitNormalization(MLIRContext *context, TrivialGraphMergeStats *stats)
+      : OpRewritePattern(context), stats(stats) {}
+
+  LogicalResult matchAndRewrite(SpatGraphCompute consumer, PatternRewriter& rewriter) const override {
+    auto producer = consumer.getInputs().empty()
+      ? SpatGraphComputeBatch()
+      : consumer.getInputs().front().getDefiningOp<SpatGraphComputeBatch>();
+    if (!producer || producer->getBlock() != consumer->getBlock() || !hasOnlyStructuralAttrs(producer)
+        || !hasOnlyStructuralAttrs(consumer) || !isExclusivelyConsumedBy(producer, consumer))
+      return failure();
+    auto fragments = collectPublishedFragments(producer);
+    if (!matchLeadingUnitNormalization(producer, consumer) || failed(fragments))
+      return failure();
+
+    rewriter.setInsertionPoint(consumer);
+    auto folded = createEmptySpatGraphComputeBatch(rewriter, consumer.getLoc(), consumer.getResultTypes(),
+      producer.getLaneCount(), producer.getWeights(), producer.getInputs());
+    if (failed(folded))
+      return failure();
+    IRMapping mapper;
+    mapper.map(*producer.getLaneArgument(), *folded->getLaneArgument());
+    mapExternalArguments(producer, *folded, mapper);
+    for (Operation& op : producer.getBody().front().without_terminator())
+      rewriter.clone(op, mapper);
+
+    auto outputType = cast<RankedTensorType>(consumer.getResult(0).getType());
+    auto fragmentType = *getGraphBatchFragmentType(outputType, producer.getLaneCount());
+    auto fragment = removeLeadingUnitTensorDimension(
+      rewriter, consumer.getLoc(), mapper.lookup(fragments->front()), fragmentType);
+    assert(succeeded(fragment) && "normalization fragment types were prechecked");
+    publishGraphBatchPhysicalFragment(rewriter, consumer.getLoc(), *fragment,
+      *folded->getOutputArgument(0), *folded->getLaneArgument());
+    rewriter.replaceOp(consumer, folded->getResults());
+    rewriter.eraseOp(producer);
+    ++stats->leadingUnitNormalizationFolds;
+    return success();
+  }
+
+private:
+  TrivialGraphMergeStats *stats;
+};
+
+static bool hasDirectLaneConsumers(SpatGraphComputeBatch producer, SpatGraphComputeBatch consumer) {
+  auto lane = consumer.getLaneArgument();
+  if (!lane)
+    return false;
+  for (auto [inputIndex, input] : llvm::enumerate(consumer.getInputs())) {
+    if (input.getDefiningOp() != producer.getOperation())
+      continue;
+    auto inputArg = consumer.getInputArgument(inputIndex);
+    auto physicalType = dyn_cast<RankedTensorType>(input.getType());
+    if (!inputArg || !physicalType || inputArg->use_empty())
+      return false;
+    for (Operation* user : inputArg->getUsers()) {
+      auto extract = dyn_cast<tensor::ExtractSliceOp>(user);
+      if (!extract || extract.getSource() != *inputArg
+          || !isDirectLaneSlot(
+            extract.getMixedOffsets(), extract.getMixedSizes(), extract.getMixedStrides(), *lane, physicalType))
+        return false;
+      auto resultType = dyn_cast<RankedTensorType>(extract.getType());
+      auto fragmentType = getGraphBatchFragmentType(physicalType, producer.getLaneCount());
+      if (failed(fragmentType) || !resultType
+          || (resultType != *fragmentType
+              && (resultType.getRank() != fragmentType->getRank() + 1 || resultType.getDimSize(0) != 1
+                  || resultType.getShape().drop_front() != fragmentType->getShape()
+                  || resultType.getElementType() != fragmentType->getElementType())))
+        return false;
+    }
+  }
+  return true;
+}
+
+struct MergeTrivialBatchComputes : OpRewritePattern<SpatGraphComputeBatch> {
+  MergeTrivialBatchComputes(MLIRContext *context, TrivialGraphMergeStats *stats)
+      : OpRewritePattern(context), stats(stats) {}
+
+  LogicalResult matchAndRewrite(SpatGraphComputeBatch consumer, PatternRewriter& rewriter) const override {
+    SpatGraphComputeBatch producer;
+    FailureOr<SmallVector<Value>> fragments = failure();
+    for (Value input : consumer.getInputs()) {
+      auto candidate = input.getDefiningOp<SpatGraphComputeBatch>();
+      if (candidate && candidate->getBlock() == consumer->getBlock()
+          && candidate.getLaneCount() == consumer.getLaneCount() && hasOnlyStructuralAttrs(candidate)
+          && hasOnlyStructuralAttrs(consumer) && isExclusivelyConsumedBy(candidate, consumer)
+          && hasCapacityFor(candidate, consumer) && hasDirectLaneConsumers(candidate, consumer)
+          && succeeded(fragments = collectPublishedFragments(candidate))) {
+        producer = candidate;
+        break;
+      }
+    }
+    if (!producer)
+      return failure();
+
+    llvm::SetVector<Value> weights, inputs;
+    collectExternalOperands(producer, consumer, weights, inputs);
+    rewriter.setInsertionPoint(consumer);
+    SpatGraphComputeBatch merged = *createEmptySpatGraphComputeBatch(rewriter, consumer.getLoc(), consumer.getResultTypes(),
+      consumer.getLaneCount(), weights.getArrayRef(), inputs.getArrayRef());
+    IRMapping mapper;
+    mapper.map(*producer.getLaneArgument(), *merged.getLaneArgument());
+    mapExternalArguments(producer, merged, mapper);
+    for (Operation& op : producer.getBody().front().without_terminator())
+      rewriter.clone(op, mapper);
+    for (auto [result, fragment] : llvm::zip(producer.getResults(), *fragments))
+      mapper.map(result, mapper.lookupOrDefault(fragment));
+
+    mapper.map(*consumer.getLaneArgument(), *merged.getLaneArgument());
+    mapExternalArguments(consumer, merged, mapper, /*mapForwardedInputs=*/false);
+    for (auto [index, output] : llvm::enumerate(consumer.getResults()))
+      mapper.map(*consumer.getOutputArgument(index), *merged.getOutputArgument(index));
+    for (Operation& op : consumer.getBody().front().without_terminator()) {
+      auto extract = dyn_cast<tensor::ExtractSliceOp>(op);
+      auto sourceArg = extract ? dyn_cast<BlockArgument>(extract.getSource()) : BlockArgument();
+      unsigned firstInputArg = 1 + consumer.getWeights().size();
+      if (!sourceArg || sourceArg.getOwner() != &consumer.getBody().front() || sourceArg.getArgNumber() < firstInputArg
+          || sourceArg.getArgNumber() >= firstInputArg + consumer.getInputs().size()) {
+        rewriter.clone(op, mapper);
+        continue;
+      }
+      unsigned inputIndex = sourceArg.getArgNumber() - firstInputArg;
+      Value input = consumer.getInputs()[inputIndex];
+      if (input.getDefiningOp() != producer.getOperation()) {
+        rewriter.clone(op, mapper);
+        continue;
+      }
+      Value fragment = mapper.lookup(input);
+      if (fragment.getType() != extract.getType()) {
+        auto expanded = addLeadingUnitTensorDimension(rewriter, extract.getLoc(), fragment);
+        assert(succeeded(expanded) && expanded->getType() == extract.getType()
+               && "prechecked physical fragment type must be forwardable");
+        fragment = *expanded;
+      }
+      mapper.map(extract.getResult(), fragment);
+    }
+    rewriter.clone(*consumer.getBody().front().getTerminator(), mapper);
+
+    rewriter.replaceOp(consumer, merged.getResults());
+    rewriter.eraseOp(producer);
+    ++stats->batchProducerConsumerMerges;
+    return success();
+  }
+
+private:
+  TrivialGraphMergeStats *stats;
+};
+
+struct TrivialGraphComputeMergePass final : PassWrapper<TrivialGraphComputeMergePass, OperationPass<ModuleOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(TrivialGraphComputeMergePass)
+
+  StringRef getArgument() const override { return "pim-trivial-graph-compute-merge"; }
+  StringRef getDescription() const override {
+    return "Inline graph computes with exclusive direct dependencies before PEFT scheduling.";
+  }
+
+  void runOnOperation() override {
+    ModuleOp module = getOperation();
+    TrivialGraphMergeStats stats;
+    std::tie(stats.scalarBefore, stats.batchBefore) = countGraphComputes(module);
+    RewritePatternSet patterns(&getContext());
+    patterns.add<MergeTrivialScalarComputes, MergeTrivialBatchComputes>(&getContext(), &stats);
+    if (failed(applyPatternsGreedily(module, std::move(patterns)))) {
+      signalPassFailure();
+      return;
+    }
+    RewritePatternSet normalization(&getContext());
+    normalization.add<FoldBatchLeadingUnitNormalization>(&getContext(), &stats);
+    if (failed(applyPatternsGreedily(module, std::move(normalization)))) {
+      signalPassFailure();
+      return;
+    }
+    std::tie(stats.scalarAfter, stats.batchAfter) = countGraphComputes(module);
+    size_t largestBatchLaneCount = 0;
+    module.walk([&](SpatGraphComputeBatch batch) {
+      largestBatchLaneCount = std::max(largestBatchLaneCount, static_cast<size_t>(batch.getLaneCount()));
+    });
+    dumpTrivialMergeReport(stats, largestBatchLaneCount);
+    dumpModule(module, "spatial2_trivial_merged");
+    SpatialDataflowExportStage exportMode = getSpatialDataflowExportStage();
+    if (shouldExportSpatialDataflowStage(exportMode, SpatialDataflowExportStage::Spatial2)) {
+      auto entryFunc = getPimEntryFunc(module);
+      if (failed(entryFunc)
+          || failed(exportSpatialDataflowCsvGraph(*entryFunc, "spatial2_trivial_merged", "spatial2")))
+        signalPassFailure();
+    }
+  }
+};
+
+} // namespace
+} // namespace spatial
+
+std::unique_ptr<Pass> createTrivialGraphComputeMergePass() {
+  return std::make_unique<spatial::TrivialGraphComputeMergePass>();
+}
+
+} // namespace onnx_mlir

src/PIM/Pass/PIMPasses.h

@@ -19,6 +19,8 @@ std::unique_ptr<mlir::Pass> createPimMemoryCoalescingPass();
 
 std::unique_ptr<mlir::Pass> createMergeComputeNodesPass();
 
+std::unique_ptr<mlir::Pass> createTrivialGraphComputeMergePass();
+
 std::unique_ptr<mlir::Pass> createPimHostConstantFoldingPass();
 
 std::unique_ptr<mlir::Pass> createPimVerificationPass();

src/PIM/PimAccelerator.cpp

@@ -77,6 +77,7 @@ void PimAccelerator::registerPasses(int optLevel) const {
   registerPass(createSpatialToPimPass);
   registerPass(createPimBufferizationPass);
   registerPass(createPimMemoryCoalescingPass);
+  registerPass(createTrivialGraphComputeMergePass);
   registerPass(createMergeComputeNodesPass);
   registerPass(createPimHostConstantFoldingPass);
   registerPass(createPimVerificationPass);

tools/raptor_graph_explorer/README.md

@@ -70,7 +70,7 @@ The server defaults to `127.0.0.1:8765`. It has no authentication and is intende
 
 ## Supported CSV reports
 
-Files are paired using the complete prefix before `.nodes.csv` and `.edges.csv`. That prefix is the stable `report_id`; `spatial2_scheduled_no_comm.nodes.csv` therefore pairs only with `spatial2_scheduled_no_comm.edges.csv`.
+Files are paired using the complete prefix before `.nodes.csv` and `.edges.csv`. That prefix is the stable `report_id`; `spatial3_scheduled_no_comm.nodes.csv` therefore pairs only with `spatial3_scheduled_no_comm.edges.csv`.
 
 Node columns:
 
@@ -116,7 +116,7 @@ Nodes group by `(report_id, node_kind)`. This deliberately small diagnostic view
 
 ### Raw instances
 
-Raw instances displays every `raw_nodes` row exactly once and every valid `raw_edges` row exactly once. Nodes are queried independently from edges, so degree-zero nodes remain visible. A spatial3 report starts in this view automatically; Operation, Core, and Node kind remain selectable.
+Raw instances displays every `raw_nodes` row exactly once and every valid `raw_edges` row exactly once. Nodes are queried independently from edges, so degree-zero nodes remain visible. A spatial4 report starts in this view automatically; Operation, Core, and Node kind remain selectable.
 
 Spatial3 CSVs record explicit communication only and may legitimately contain many isolated nodes. The explorer never fabricates dependencies from semicolon-separated SSA names or other node text.
 
@@ -155,7 +155,7 @@ Motifs are not inferred from rendered geometry. For each operation graph the too
 
 ## Viewer and API
 
-Except for spatial3, the viewer initially fetches an aggregate operation graph unless the browser has a saved view choice. Sigma renders the graph with WebGL. Controls cover report/view/metric selection, text and tensor search, mapping filters, self edges, relayout, fitting, and motif selection. Raw nodes and edges have their own selection details and never call aggregate-only detail or mapping endpoints. Mapping panels remain available for operation aggregate edges.
+Except for spatial4, the viewer initially fetches an aggregate operation graph unless the browser has a saved view choice. Sigma renders the graph with WebGL. Controls cover report/view/metric selection, text and tensor search, mapping filters, self edges, relayout, fitting, and motif selection. Raw nodes and edges have their own selection details and never call aggregate-only detail or mapping endpoints. Mapping panels remain available for operation aggregate edges.
 
 Operation expansion is a deterministic projection of the source graph. Expand selected, Expand all operations, Collapse selected operation, and Collapse all rebuild the complete display from the current expanded-operation set. An expanded operation's raw nodes replace its aggregate node, including isolated raw nodes. An aggregate edge is retained only when both endpoint operations are collapsed; otherwise its raw edges replace it with endpoints calculated from the complete expansion set. Expansion order therefore cannot leave stale endpoints or simultaneous aggregate/raw representations.
 

tools/raptor_graph_explorer/raptor_graph_explorer.egg-info/PKG-INFO

@@ -84,7 +84,7 @@ The server defaults to `127.0.0.1:8765`. It has no authentication and is intende
 
 ## Supported CSV reports
 
-Files are paired using the complete prefix before `.nodes.csv` and `.edges.csv`. That prefix is the stable `report_id`; `spatial2_scheduled_no_comm.nodes.csv` therefore pairs only with `spatial2_scheduled_no_comm.edges.csv`.
+Files are paired using the complete prefix before `.nodes.csv` and `.edges.csv`. That prefix is the stable `report_id`; `spatial3_scheduled_no_comm.nodes.csv` therefore pairs only with `spatial3_scheduled_no_comm.edges.csv`.
 
 Node columns:
 
@@ -130,7 +130,7 @@ Nodes group by `(report_id, node_kind)`. This deliberately small diagnostic view
 
 ### Raw instances
 
-Raw instances displays every `raw_nodes` row exactly once and every valid `raw_edges` row exactly once. Nodes are queried independently from edges, so degree-zero nodes remain visible. A spatial3 report starts in this view automatically; Operation, Core, and Node kind remain selectable.
+Raw instances displays every `raw_nodes` row exactly once and every valid `raw_edges` row exactly once. Nodes are queried independently from edges, so degree-zero nodes remain visible. A spatial4 report starts in this view automatically; Operation, Core, and Node kind remain selectable.
 
 Spatial3 CSVs record explicit communication only and may legitimately contain many isolated nodes. The explorer never fabricates dependencies from semicolon-separated SSA names or other node text.
 
@@ -169,7 +169,7 @@ Motifs are not inferred from rendered geometry. For each operation graph the too
 
 ## Viewer and API
 
-Except for spatial3, the viewer initially fetches an aggregate operation graph unless the browser has a saved view choice. Sigma renders the graph with WebGL. Controls cover report/view/metric selection, text and tensor search, mapping filters, self edges, relayout, fitting, and motif selection. Raw nodes and edges have their own selection details and never call aggregate-only detail or mapping endpoints. Mapping panels remain available for operation aggregate edges.
+Except for spatial4, the viewer initially fetches an aggregate operation graph unless the browser has a saved view choice. Sigma renders the graph with WebGL. Controls cover report/view/metric selection, text and tensor search, mapping filters, self edges, relayout, fitting, and motif selection. Raw nodes and edges have their own selection details and never call aggregate-only detail or mapping endpoints. Mapping panels remain available for operation aggregate edges.
 
 Operation expansion is a deterministic projection of the source graph. Expand selected, Expand all operations, Collapse selected operation, and Collapse all rebuild the complete display from the current expanded-operation set. An expanded operation's raw nodes replace its aggregate node, including isolated raw nodes. An aggregate edge is retained only when both endpoint operations are collapsed; otherwise its raw edges replace it with endpoints calculated from the complete expansion set. Expansion order therefore cannot leave stale endpoints or simultaneous aggregate/raw representations.
 

tools/raptor_graph_explorer/raptor_graph_explorer.egg-info/SOURCES.txt

@@ -1,7 +1,6 @@
 README.md
 pyproject.toml
 raptor_graph_explorer/__init__.py
-raptor_graph_explorer/__main__.py
 raptor_graph_explorer/aggregate.py
 raptor_graph_explorer/api.py
 raptor_graph_explorer/cli.py

tools/raptor_graph_explorer/raptor_graph_explorer/__init__.py

@@ -0,0 +1 @@
+__version__ = "0.1.0"

tools/raptor_graph_explorer/raptor_graph_explorer/static/app.js

@@ -18,7 +18,7 @@ function definition(object, keys) {
 function status(message, error = false) { $("status").textContent = message; $("status").className = error ? "error" : ""; }
 function selectedReport() { return state.reports.find(report => report.report_id === $("report").value); }
 function defaultViewForReport(report) {
-  if (report?.stage === "spatial3") return "raw";
+  if (report?.stage === "spatial4") return "raw";
   const saved = localStorage.getItem("raptor-graph-view");
   return ["operation", "raw", "core", "node_kind"].includes(saved) ? saved : "operation";
 }

tools/raptor_graph_explorer/tests/fixtures/regressions/spatial4_scheduled.edges.csv

@@ -1,2 +1,2 @@
 Source,Target,Weight,Type,stage,source_lane,target_lane,channel_id
-scb:1:5,sc:0,512,tensor<1x8x1x16xf32>,spatial3,5,,0
+scb:1:5,sc:0,512,tensor<1x8x1x16xf32>,spatial4,5,,0

tools/raptor_graph_explorer/tests/test_projection.py

@@ -79,8 +79,8 @@ def test_collapse_reconstructs_every_boundary(regression_output: Path):
     assert all(edge["representation"] == "raw" for edge in graph["edges"])
 
 
-def test_spatial3_includes_isolated_nodes_and_has_distinct_positions(regression_output: Path):
-    graph = project(regression_output, "spatial3_scheduled", view="raw")
+def test_spatial4_includes_isolated_nodes_and_has_distinct_positions(regression_output: Path):
+    graph = project(regression_output, "spatial4_scheduled", view="raw")
     assert graph["counts"] == {
         "nodes": 16, "edges": 1, "isolated_nodes": 14,
         "raw_nodes": 16, "aggregate_nodes": 0, "raw_edges": 1, "aggregate_edges": 0,
@@ -100,20 +100,20 @@ def test_header_only_edges_preserve_all_nodes(regression_output: Path):
 
 def test_projection_cap_is_actionable_and_never_partial(regression_output: Path):
     with pytest.raises(DisplayGraphTooLarge, match="Use an aggregate view or increase"):
-        project(regression_output, "spatial3_scheduled", view="raw", cap=16)
+        project(regression_output, "spatial4_scheduled", view="raw", cap=16)
     client = TestClient(create_app(regression_output, expansion_cap=16))
-    response = client.get("/api/reports/spatial3_scheduled/display-graph?view=raw")
+    response = client.get("/api/reports/spatial4_scheduled/display-graph?view=raw")
     assert response.status_code == 400
     assert "safety cap" in response.json()["detail"] and "aggregate view" in response.json()["detail"]
 
 
 def test_display_api_and_static_raw_interactions(regression_output: Path):
     client = TestClient(create_app(regression_output))
-    assert client.get("/api/reports/spatial3_scheduled/display-graph?view=raw").json()["counts"]["nodes"] == 16
+    assert client.get("/api/reports/spatial4_scheduled/display-graph?view=raw").json()["counts"]["nodes"] == 16
     index = client.get("/").text
     script = client.get("/app.js").text
     assert '<option value="raw">Raw instances</option>' in index
-    assert 'report?.stage === "spatial3"' in script
+    assert 'report?.stage === "spatial4"' in script
     assert 'local.representation === "raw"' in script
     reducer = script[script.index("function reduceEdge"):script.index("function applyFilters")]
     assert "state.graph.source(id)" not in reducer

validation/raptor.py

@@ -40,7 +40,7 @@ def _format_command(cmd):
 
 
 def compile_with_raptor(network_path, raptor_onnx_path: Path, output_base: Path,
-                        crossbar_size, crossbar_count, core_count=None, pim_merge_scheduler="peft",
+                        crossbar_size, crossbar_count, core_count=None,
                         pim_memory_report="none", raptor_extra_args=None, cwd=None, verbose=False,
                         reporter=None, timeout_sec=None):
     # Define the arguments, with the possibility to set crossbar size and count
@@ -52,7 +52,6 @@ def compile_with_raptor(network_path, raptor_onnx_path: Path, output_base: Path,
         "--EmitPimCodegen",
         f"--crossbar-size={crossbar_size}",
         f"--crossbar-count={crossbar_count}",
-        f"--pim-merge-scheduler={pim_merge_scheduler}",
     ]
     if core_count is not None:
         args.append(f"--core-count={core_count}")

validation/validate.py

@@ -76,8 +76,6 @@ def main():
     ap.add_argument("--crossbar-count", type=int, default=8)
     ap.add_argument("--core-count", type=int, default=None,
                     help="Core count to pass to Raptor. Required for PIM validation.")
-    ap.add_argument("--pim-merge-scheduler", choices=("peft"), default="peft",
-                    help="Scheduler used by the Spatial merge-compute-nodes pass.")
     ap.add_argument("--pim-memory-report", choices=("none", "summary", "full"), default="none",
                     help="Emit a human-readable PIM memory planning report during codegen.")
     ap.add_argument("--raptor-extra-arg", action="append", default=[],
@@ -149,7 +147,7 @@ def main():
             result = validate_network(
                 onnx_path, a.raptor_path, a.onnx_include_dir, simulator_dir,
                 crossbar_size=a.crossbar_size, crossbar_count=a.crossbar_count, core_count=a.core_count,
-                pim_merge_scheduler=a.pim_merge_scheduler, pim_memory_report=a.pim_memory_report,
+                pim_memory_report=a.pim_memory_report,
                 raptor_extra_args=a.raptor_extra_arg,
                 command_timeout_seconds=a.command_timeout_seconds,
                 threshold=a.threshold,

validation/validate_one.py

@@ -1,4 +1,5 @@
 import json
+import os
 import numpy as np
 import shutil
 import sys
@@ -61,32 +62,37 @@ class ProgressReporter:
         self.passed_models = 0
         self.failed_models = 0
         self.current_label = ""
-        self.enabled = sys.stdout.isatty() if enabled is None else enabled
+        self.enabled = (
+            sys.stdout.isatty() and "CODEX_CI" not in os.environ
+            if enabled is None else enabled
+        )
         self.verbose = verbose
-        self.columns = shutil.get_terminal_size((100, 20)).columns
+        self.columns = max(1, shutil.get_terminal_size((100, 20)).columns)
         self.suspended = False
         self.rendered_width = 0
+        self.rendered_rows = 0
 
     def _clear(self):
-        if self.enabled:
-            sys.stdout.write("\r" + (" " * self.rendered_width) + "\r")
+        if self.enabled and self.rendered_rows:
+            columns = max(1, shutil.get_terminal_size((100, 20)).columns)
+            rows = max(self.rendered_rows, (self.rendered_width + columns - 1) // columns)
+            sys.stdout.write("\r\033[2K")
+            for _ in range(rows - 1):
+                sys.stdout.write("\033[1A\r\033[2K")
             sys.stdout.flush()
+            self.rendered_width = 0
+            self.rendered_rows = 0
 
     def _render(self):
         if not self.enabled or self.suspended:
             return
-        bar_width = 24
+        self.columns = max(1, shutil.get_terminal_size((100, 20)).columns)
+        bar_width = min(24, max(4, self.columns - 24))
         filled = int(bar_width * self.completed_steps / self.total_steps)
         counts_text = f"P:{self.passed_models} F:{self.failed_models}"
         prefix_text = f"[{'#' * filled}{'-' * (bar_width - filled)}] {self.completed_steps}/{self.total_steps}"
-        if len(prefix_text) > self.columns:
-            prefix_text = f"{self.completed_steps}/{self.total_steps}"
-
-        if prefix_text.startswith("["):
-            bar = Fore.GREEN + ("#" * filled) + Fore.CYAN + ("-" * (bar_width - filled))
-            prefix = Fore.CYAN + f"[{bar}{Fore.CYAN}] {self.completed_steps}/{self.total_steps}" + Style.RESET_ALL
-        else:
-            prefix = Fore.CYAN + prefix_text + Style.RESET_ALL
+        bar = Fore.GREEN + ("#" * filled) + Fore.CYAN + ("-" * (bar_width - filled))
+        prefix = Fore.CYAN + f"[{bar}{Fore.CYAN}] {self.completed_steps}/{self.total_steps}" + Style.RESET_ALL
 
         counts = (
                 " "
@@ -109,14 +115,28 @@ class ProgressReporter:
         elif self.current_label:
             label = f" {self.current_label}"
 
-        available_label_width = max(0, self.columns - len(prefix_text) - len(model_counter) - len(counts_text) - 3)
+        fixed_width = len(prefix_text) + len(model_counter) + len(counts_text) + 2
+        if fixed_width > self.columns:
+            model_counter = ""
+            fixed_width = len(prefix_text) + len(counts_text) + 2
+        if fixed_width > self.columns:
+            prefix_text = f"{self.completed_steps}/{self.total_steps}"
+            prefix = Fore.CYAN + prefix_text + Style.RESET_ALL
+            fixed_width = len(prefix_text) + len(counts_text) + 2
+        if fixed_width > self.columns:
+            counts = ""
+            counts_text = ""
+            fixed_width = len(prefix_text) + 1
+        available_label_width = max(0, self.columns - fixed_width)
         label = label[:available_label_width]
-        plain_line = prefix_text + model_counter + f" P:{self.passed_models} F:{self.failed_models}" + label
+        plain_counts = f" {counts_text}" if counts_text else ""
+        plain_line = prefix_text + model_counter + plain_counts + label
         rendered_line = prefix + model_counter + counts + label + Style.RESET_ALL
-        padded_width = max(self.rendered_width, len(plain_line))
-        sys.stdout.write("\r" + rendered_line + (" " * max(0, padded_width - len(plain_line))))
+        self._clear()
+        sys.stdout.write(rendered_line)
         sys.stdout.flush()
         self.rendered_width = len(plain_line)
+        self.rendered_rows = max(1, (self.rendered_width + self.columns - 1) // self.columns)
 
     def log(self, message="", color=None):
         if not self.verbose:
@@ -158,7 +178,6 @@ class ProgressReporter:
         if self.enabled:
             self.suspended = True
             self._clear()
-            self.rendered_width = 0
 
 
 def run_command(cmd, cwd=None, reporter=None, timeout_sec=None):
@@ -293,7 +312,7 @@ def validate_outputs(sim_arrays, runner_out_dir, outputs_descriptor, threshold=1
 
 def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
                      simulator_dir, crossbar_size=64, crossbar_count=8, core_count=None,
-                     pim_merge_scheduler="peft", pim_memory_report="none", raptor_extra_args=None,
+                     pim_memory_report="none", raptor_extra_args=None,
                      threshold=1e-3, rtol=1e-5,
                      seed=0, reporter=None, model_index=1, model_total=1, verbose=False,
                      command_timeout_seconds=60.0, mode=MODE_FULL):
@@ -346,8 +365,8 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
             if mode == MODE_COMPILE_ONLY:
                 print_stage(reporter, model_index, model_total, network_onnx_path.name, "Compile PIM")
                 pim_pass_timings = compile_with_raptor(
-                    network_mlir_path, raptor_path, pim_output_base, crossbar_size, crossbar_count,
-                    core_count=core_count, pim_merge_scheduler=pim_merge_scheduler,
+                    network_onnx_path, raptor_path, pim_output_base, crossbar_size, crossbar_count,
+                    core_count=core_count,
                     pim_memory_report=pim_memory_report, raptor_extra_args=raptor_extra_args,
                     cwd=raptor_dir, verbose=verbose, reporter=reporter, timeout_sec=command_timeout_seconds)
                 print_info(reporter, f"PIM artifacts saved to {raptor_dir / 'pim'}")
@@ -386,8 +405,8 @@ def validate_network(network_onnx_path, raptor_path, onnx_include_dir,
         if mode != MODE_RUN_ONLY:
             print_stage(reporter, model_index, model_total, network_onnx_path.name, "Compile PIM")
             pim_pass_timings = compile_with_raptor(
-                network_mlir_path, raptor_path, pim_output_base, crossbar_size, crossbar_count,
-                core_count=core_count, pim_merge_scheduler=pim_merge_scheduler,
+                network_onnx_path, raptor_path, pim_output_base, crossbar_size, crossbar_count,
+                core_count=core_count,
                 pim_memory_report=pim_memory_report, raptor_extra_args=raptor_extra_args,
                 cwd=raptor_dir, verbose=verbose, reporter=reporter, timeout_sec=command_timeout_seconds)
             print_info(reporter, f"PIM artifacts saved to {raptor_dir / 'pim'}")