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

AGENTS.md

@@ -1,4 +1,5 @@
 * Always read the full README.md before doing anything
+* Always read the full invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md before modifying Spatial graph IR, Blueprint handling, or MergeComputeNodes.
 * Build commands:
   * `cmake --build ./build_release`
   * `cmake --build ./build_debug`

invariants/GRAPH_COMPUTE_BATCH_INVARIANT.md

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

src/PIM/CMakeLists.txt

@@ -117,7 +117,6 @@ add_pim_library(OMPIMAccel
   SpatialOps
   PimOps
   OMONNXToSpatial
-  OMSpatialToGraphviz
   OMSpatialToPim
   OMPimCommon
   OMPimBufferization

src/PIM/Common/IR/CoreBlockUtils.cpp

@@ -74,6 +74,21 @@ walkPimCoreBlock(mlir::Block& block,
       continue;
     }
 
+    if (auto ifOp = mlir::dyn_cast<mlir::scf::IfOp>(op)) {
+      auto condition = resolveIndexValue(ifOp.getCondition(), knowledge);
+      if (failed(condition)) {
+        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen");
+        hasFailure = true;
+        continue;
+      }
+
+      mlir::Region& selectedRegion = *condition != 0 ? ifOp.getThenRegion() : ifOp.getElseRegion();
+      if (!selectedRegion.empty())
+        if (failed(walkPimCoreBlock(selectedRegion.front(), knowledge, callback)))
+          hasFailure = true;
+      continue;
+    }
+
     if (failed(callback(op, knowledge)))
       hasFailure = true;
   }
@@ -128,6 +143,22 @@ mlir::LogicalResult walkPimCoreBlockStructurally(
       continue;
     }
 
+    if (auto ifOp = mlir::dyn_cast<mlir::scf::IfOp>(op)) {
+      if (failed(resolveIndexValue(ifOp.getCondition(), knowledge))) {
+        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM verification");
+        hasFailure = true;
+        continue;
+      }
+
+      if (!ifOp.getThenRegion().empty())
+        if (failed(walkPimCoreBlockStructurally(ifOp.getThenRegion().front(), knowledge, callback)))
+          hasFailure = true;
+      if (!ifOp.getElseRegion().empty())
+        if (failed(walkPimCoreBlockStructurally(ifOp.getElseRegion().front(), knowledge, callback)))
+          hasFailure = true;
+      continue;
+    }
+
     if (failed(callback(op, knowledge)))
       hasFailure = true;
   }

src/PIM/Common/Support/DebugDump.cpp

@@ -7,18 +7,26 @@
 
 namespace onnx_mlir {
 
-void dumpModule(mlir::ModuleOp moduleOp, const std::string& name) {
+std::fstream openDialectDumpFileWithExtension(const std::string& name, llvm::StringRef destination,  llvm::StringRef extension) {
   std::string outputDir = getOutputDir();
   if (outputDir.empty())
+    return {};
+
+  std::string dialectsDir = (outputDir + destination).str();
+  createDirectory(dialectsDir);
+  return std::fstream(dialectsDir + "/" + name + "." + extension.str(), std::ios::out);
+}
+
+void dumpModule(mlir::ModuleOp moduleOp, const std::string& name, bool assumeVerified) {
+  std::fstream file = openDialectDumpFileWithExtension(name, "/dialects", "mlir");
+  if (!file.is_open())
     return;
 
-  std::string dialectsDir = outputDir + "/dialects";
-  createDirectory(dialectsDir);
-
-  std::fstream file(dialectsDir + "/" + name + ".mlir", std::ios::out);
   llvm::raw_os_ostream os(file);
   mlir::OpPrintingFlags flags;
-  flags.elideLargeElementsAttrs().enableDebugInfo(true, false);
+  flags.elideLargeElementsAttrs().enableDebugInfo(false, false);
+  if (assumeVerified)
+    flags.assumeVerified();
   moduleOp.print(os, flags);
   os.flush();
   file.close();

src/PIM/Common/Support/DebugDump.hpp

@@ -1,13 +1,18 @@
 #pragma once
 
 #include "mlir/IR/BuiltinOps.h"
+#include "llvm/ADT/StringRef.h"
 
+#include <fstream>
 #include <string>
 
 namespace onnx_mlir {
 
 /// Emits a MLIR snapshot under the current compiler output
 /// directory for pass-level debugging.
-void dumpModule(mlir::ModuleOp moduleOp, const std::string& name);
+void dumpModule(mlir::ModuleOp moduleOp, const std::string& name, bool assumeVerified = false);
+
+/// Opens a file under the same dialect dump directory used by dumpModule.
+std::fstream openDialectDumpFileWithExtension(const std::string& name,llvm::StringRef destination = "/dialects",  llvm::StringRef extension = "mlir");
 
 } // namespace onnx_mlir

src/PIM/Compiler/PimCodeGen.cpp

@@ -414,31 +414,35 @@ size_t PimAcceleratorMemory::getValueAddress(mlir::Value value,
                                              const StaticValueKnowledge& knowledge,
                                              std::optional<unsigned> lane) const {
   value = resolveCachedAlias(value, knowledge);
-  auto compiledIt = compiledAddressExprs.find(value);
-  if (compiledIt == compiledAddressExprs.end()) {
-    auto compiledExpr = compileContiguousAddressExpr(value);
-    if (failed(compiledExpr)) {
-      errs() << "Failed to compile contiguous address for value: ";
+
+  FailureOr<ResolvedContiguousAddress> resolvedAddress = resolveContiguousAddress(value, knowledge);
+  if (failed(resolvedAddress)) {
+    auto compiledIt = compiledAddressExprs.find(value);
+    if (compiledIt == compiledAddressExprs.end()) {
+      auto compiledExpr = compileContiguousAddressExpr(value);
+      if (failed(compiledExpr)) {
+        errs() << "Failed to compile contiguous address for value: ";
+        value.print(errs());
+        errs() << " : " << value.getType();
+        errs() << "\n";
+        llvm_unreachable("Failed to compile contiguous address");
+      }
+      compiledIt = compiledAddressExprs.try_emplace(value, *compiledExpr).first;
+    }
+
+    resolvedAddress = compiledIt->second.evaluate(knowledge, lane);
+    if (failed(resolvedAddress)) {
+      errs() << "Failed to evaluate contiguous address for value: ";
       value.print(errs());
       errs() << " : " << value.getType();
       errs() << "\n";
-      llvm_unreachable("Failed to compile contiguous address");
+      if (auto* definingOp = value.getDefiningOp()) {
+        errs() << "Defining op:\n";
+        definingOp->print(errs());
+        errs() << "\n";
+      }
+      llvm_unreachable("Failed to resolve contiguous address");
     }
-    compiledIt = compiledAddressExprs.try_emplace(value, *compiledExpr).first;
-  }
-
-  auto resolvedAddress = compiledIt->second.evaluate(knowledge, lane);
-  if (failed(resolvedAddress)) {
-    errs() << "Failed to evaluate contiguous address for value: ";
-    value.print(errs());
-    errs() << " : " << value.getType();
-    errs() << "\n";
-    if (auto* definingOp = value.getDefiningOp()) {
-      errs() << "Defining op:\n";
-      definingOp->print(errs());
-      errs() << "\n";
-    }
-    llvm_unreachable("Failed to resolve contiguous address");
   }
 
   MemoryValueKey key = getMemoryValueKey(resolvedAddress->base, lane);
@@ -1114,7 +1118,8 @@ enum class CompiledCoreOpKind : uint8_t {
 struct CompiledCoreNode {
   enum class Kind : uint8_t {
     Op,
-    Loop
+    Loop,
+    If
   };
 
   Kind kind = Kind::Op;
@@ -1123,7 +1128,10 @@ struct CompiledCoreNode {
   CompiledIndexExpr lowerBound;
   CompiledIndexExpr upperBound;
   CompiledIndexExpr step;
+  CompiledIndexExpr condition;
   std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> loopBody;
+  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> thenBody;
+  std::unique_ptr<llvm::SmallVector<CompiledCoreNode, 8>> elseBody;
 };
 
 static FailureOr<CompiledCoreOpKind> classifyCompiledCoreOpKind(Operation& op) {
@@ -1201,6 +1209,28 @@ compileCoreEmissionPlan(Block& block, Operation* weightOwner, llvm::SmallVectorI
       continue;
     }
 
+    if (auto ifOp = dyn_cast<mlir::scf::IfOp>(op)) {
+      auto condition = compileIndexExpr(ifOp.getCondition());
+      if (failed(condition)) {
+        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen");
+        return failure();
+      }
+
+      CompiledCoreNode ifNode;
+      ifNode.kind = CompiledCoreNode::Kind::If;
+      ifNode.op = ifOp.getOperation();
+      ifNode.condition = *condition;
+      ifNode.thenBody = std::make_unique<llvm::SmallVector<CompiledCoreNode, 8>>();
+      if (failed(compileCoreEmissionPlan(ifOp.getThenRegion().front(), weightOwner, *ifNode.thenBody)))
+        return failure();
+      ifNode.elseBody = std::make_unique<llvm::SmallVector<CompiledCoreNode, 8>>();
+      if (!ifOp.getElseRegion().empty())
+        if (failed(compileCoreEmissionPlan(ifOp.getElseRegion().front(), weightOwner, *ifNode.elseBody)))
+          return failure();
+      plan.push_back(std::move(ifNode));
+      continue;
+    }
+
     auto opKind = classifyCompiledCoreOpKind(op);
     if (failed(opKind)) {
       InFlightDiagnostic diag = op.emitError() << "unsupported codegen for op '" << op.getName().getStringRef() << "'";
@@ -1263,6 +1293,26 @@ static LogicalResult executeCompiledCorePlan(
       continue;
     }
 
+    if (node.kind == CompiledCoreNode::Kind::If) {
+      auto condition = node.condition.evaluate(knowledge);
+      auto ifOp = cast<mlir::scf::IfOp>(node.op);
+      if (failed(condition)) {
+        ifOp.emitOpError("requires statically evaluable scf.if condition for PIM codegen");
+        return failure();
+      }
+
+      const auto& selectedBody = *condition != 0 ? node.thenBody : node.elseBody;
+      if (selectedBody && failed(executeCompiledCorePlan(*selectedBody,
+                                                         coreCodeGen,
+                                                         knowledge,
+                                                         resolveWeightSlot,
+                                                         processedOperations,
+                                                         batchLane,
+                                                         batchLaneCount)))
+        return failure();
+      continue;
+    }
+
     switch (node.opKind) {
     case CompiledCoreOpKind::Load:
       coreCodeGen.codeGenLoadOp(cast<pim::PimMemCopyHostToDevOp>(node.op), knowledge);

src/PIM/Compiler/PimCompilerOptions.cpp

@@ -57,6 +57,18 @@ llvm::cl::opt<PimConvLoweringType> pimConvLowering(
   llvm::cl::init(PimConvLoweringAuto),
   llvm::cl::cat(OnnxMlirOptions));
 
+llvm::cl::opt<PimSpatialDataflowExportType> pimExportSpatialDataflow(
+  "pim-export-spatial-dataflow",
+  llvm::cl::desc("Emit Gephi-importable CSV dataflow reports around MergeComputeNodes materialization"),
+  llvm::cl::values(clEnumValN(SpatialDataflowExportNone, "none", "Do not emit Spatial dataflow CSV reports")),
+  llvm::cl::values(clEnumValN(SpatialDataflowExportPre, "pre", "Emit pre-materialization Spatial dataflow CSV reports")),
+  llvm::cl::values(
+    clEnumValN(SpatialDataflowExportPost, "post", "Emit post-materialization Spatial dataflow CSV reports")),
+  llvm::cl::values(
+    clEnumValN(SpatialDataflowExportBoth, "both", "Emit both pre- and post-materialization Spatial dataflow CSV reports")),
+  llvm::cl::init(SpatialDataflowExportNone),
+  llvm::cl::cat(OnnxMlirOptions));
+
 llvm::cl::opt<bool>
   pimOnlyCodegen("pim-only-codegen",
                  llvm::cl::desc("Only generate code for PIM (assume input is already in bufferized PIM IR)"),

src/PIM/Compiler/PimCompilerOptions.hpp

@@ -42,11 +42,19 @@ typedef enum {
   PimConvLoweringTiled2D = 8,
 } PimConvLoweringType;
 
+typedef enum {
+  SpatialDataflowExportNone = 0,
+  SpatialDataflowExportPre = 1,
+  SpatialDataflowExportPost = 2,
+  SpatialDataflowExportBoth = 3,
+} 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;
 
 extern llvm::cl::opt<bool> pimOnlyCodegen;
 extern llvm::cl::opt<bool> pimDisableMemoryCoalescing;

src/PIM/Conversion/CMakeLists.txt

@@ -1,3 +1,2 @@
 add_subdirectory(ONNXToSpatial)
-add_subdirectory(SpatialToGraphviz)
-add_subdirectory(SpatialToPim)
+add_subdirectory(SpatialToPim)

src/PIM/Conversion/ONNXToSpatial/CMakeLists.txt

@@ -20,6 +20,7 @@ add_pim_library(OMONNXToSpatial
   Patterns/NN/Sigmoid.cpp
   Patterns/NN/Softmax.cpp
   Patterns/Tensor/Concat.cpp
+  Patterns/Tensor/Flatten.cpp
   Patterns/Tensor/Gather.cpp
   Patterns/Tensor/Resize.cpp
   Patterns/Tensor/Reshape.cpp
@@ -30,8 +31,10 @@ add_pim_library(OMONNXToSpatial
   SpatialLayoutPlanningPass.cpp
   LowerSpatialPlansPass.cpp
   Common/AttributeUtils.cpp
+  Common/BiasAddUtils.cpp
   Common/ComputeRegionBuilder.cpp
   Common/MatrixProductLowering.cpp
+  Common/RowStripLayoutUtils.cpp
   Common/ShapeTilingUtils.cpp
   Common/WeightMaterialization.cpp
 

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

@@ -0,0 +1,112 @@
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+LogicalResult isSupportedBiasAddShape(RankedTensorType biasType, RankedTensorType resultType) {
+  if (!biasType || !resultType || !biasType.hasStaticShape() || !resultType.hasStaticShape())
+    return failure();
+  if (resultType.getRank() != 4)
+    return failure();
+  if (biasType.getElementType() != resultType.getElementType())
+    return failure();
+
+  const int64_t channels = resultType.getDimSize(1);
+  ArrayRef<int64_t> shape = biasType.getShape();
+  if (shape.empty())
+    return success();
+  if (shape.size() == 1)
+    return success(shape[0] == channels);
+  if (shape.size() == 2)
+    return success(shape[0] == 1 && shape[1] == channels);
+  if (shape.size() == 4)
+    return success(shape[0] == 1 && shape[1] == channels && shape[2] == 1 && shape[3] == 1);
+  return failure();
+}
+
+FailureOr<SmallVector<Attribute>> getBiasChannelValues(DenseElementsAttr denseAttr, RankedTensorType resultType) {
+  auto biasType = dyn_cast<RankedTensorType>(denseAttr.getType());
+  if (!biasType || failed(isSupportedBiasAddShape(biasType, resultType)))
+    return failure();
+
+  const int64_t channels = resultType.getDimSize(1);
+  if (denseAttr.isSplat()) {
+    return SmallVector<Attribute>(channels, denseAttr.getSplatValue<Attribute>());
+  }
+
+  SmallVector<Attribute> flattened(denseAttr.getValues<Attribute>());
+  if (biasType.getRank() == 1)
+    return flattened;
+  if (biasType.getRank() == 2)
+    return flattened;
+
+  SmallVector<Attribute> channelValues;
+  channelValues.reserve(channels);
+  const int64_t channelStride = biasType.getDimSize(2) * biasType.getDimSize(3);
+  for (int64_t channel = 0; channel < channels; ++channel)
+    channelValues.push_back(flattened[channel * channelStride]);
+  return channelValues;
+}
+
+bool isSupportedBiasAddValue(Value bias, RankedTensorType resultType, DenseElementsAttr* denseAttr) {
+  auto attr = getHostConstDenseElementsAttr(bias);
+  if (!attr)
+    return false;
+  auto biasType = dyn_cast<RankedTensorType>(attr.getType());
+  if (!biasType || failed(isSupportedBiasAddShape(biasType, resultType)))
+    return false;
+  if (failed(getBiasChannelValues(attr, resultType)))
+    return false;
+  if (denseAttr)
+    *denseAttr = attr;
+  return true;
+}
+
+FailureOr<BiasAddPlanCandidate> classifyBiasAddPlanCandidate(Value lhs, Value rhs, RankedTensorType resultType) {
+  auto lhsType = dyn_cast<RankedTensorType>(lhs.getType());
+  auto rhsType = dyn_cast<RankedTensorType>(rhs.getType());
+  if (!lhsType || !rhsType)
+    return failure();
+  if (lhsType == resultType && isSupportedBiasAddValue(rhs, resultType))
+    return BiasAddPlanCandidate {lhs, rhs};
+  if (rhsType == resultType && isSupportedBiasAddValue(lhs, resultType))
+    return BiasAddPlanCandidate {rhs, lhs};
+  return failure();
+}
+
+FailureOr<Value>
+materializeDenseBiasAddTensor(Value bias, RankedTensorType resultType, RewriterBase& rewriter, Location loc) {
+  DenseElementsAttr denseAttr;
+  if (!isSupportedBiasAddValue(bias, resultType, &denseAttr))
+    return failure();
+
+  FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, resultType);
+  if (failed(channelValues))
+    return failure();
+
+  SmallVector<Attribute> resultValues;
+  resultValues.reserve(resultType.getNumElements());
+  const int64_t batches = resultType.getDimSize(0);
+  const int64_t channels = resultType.getDimSize(1);
+  const int64_t height = resultType.getDimSize(2);
+  const int64_t width = resultType.getDimSize(3);
+  for (int64_t n = 0; n < batches; ++n)
+    for (int64_t c = 0; c < channels; ++c)
+      for (int64_t h = 0; h < height; ++h)
+        for (int64_t w = 0; w < width; ++w)
+          resultValues.push_back((*channelValues)[c]);
+
+  auto resultAttr = DenseElementsAttr::get(resultType, resultValues);
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), resultAttr, resultType);
+}
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,30 @@
+#pragma once
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/Value.h"
+#include "mlir/Support/LogicalResult.h"
+
+namespace onnx_mlir {
+
+struct BiasAddPlanCandidate {
+  mlir::Value data;
+  mlir::Value bias;
+};
+
+mlir::LogicalResult isSupportedBiasAddShape(mlir::RankedTensorType biasType, mlir::RankedTensorType resultType);
+bool isSupportedBiasAddValue(mlir::Value bias,
+                             mlir::RankedTensorType resultType,
+                             mlir::DenseElementsAttr* denseAttr = nullptr);
+mlir::FailureOr<llvm::SmallVector<mlir::Attribute>>
+getBiasChannelValues(mlir::DenseElementsAttr denseAttr, mlir::RankedTensorType resultType);
+mlir::FailureOr<BiasAddPlanCandidate> classifyBiasAddPlanCandidate(mlir::Value lhs,
+                                                                   mlir::Value rhs,
+                                                                   mlir::RankedTensorType resultType);
+mlir::FailureOr<mlir::Value> materializeDenseBiasAddTensor(mlir::Value bias,
+                                                           mlir::RankedTensorType resultType,
+                                                           mlir::RewriterBase& rewriter,
+                                                           mlir::Location loc);
+
+} // namespace onnx_mlir

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

@@ -60,6 +60,56 @@ struct SpatComputeBatchBodyArgs {
   mlir::ValueRange outputs;
 };
 
+inline mlir::SmallVector<mlir::Type> getGraphComputeBlockArgTypes(mlir::ValueRange weights, mlir::ValueRange inputs) {
+  mlir::SmallVector<mlir::Type> blockArgTypes;
+  blockArgTypes.reserve(weights.size() + inputs.size());
+  for (mlir::Value weight : weights)
+    blockArgTypes.push_back(weight.getType());
+  for (mlir::Value input : inputs)
+    blockArgTypes.push_back(input.getType());
+  return blockArgTypes;
+}
+
+inline mlir::SmallVector<mlir::Location> getGraphComputeBlockArgLocs(mlir::Location defaultLoc,
+                                                                     mlir::ValueRange weights,
+                                                                     mlir::ValueRange inputs) {
+  mlir::SmallVector<mlir::Location> blockArgLocs;
+  blockArgLocs.reserve(weights.size() + inputs.size());
+  for (mlir::Value weight : weights)
+    blockArgLocs.push_back(weight.getLoc());
+  for (mlir::Value input : inputs)
+    blockArgLocs.push_back(input.getLoc());
+  return blockArgLocs;
+}
+
+inline mlir::SmallVector<mlir::Type> getGraphComputeBatchBlockArgTypes(mlir::OpBuilder& builder,
+                                                                       mlir::TypeRange resultTypes,
+                                                                       mlir::ValueRange weights,
+                                                                       mlir::ValueRange inputs) {
+  mlir::SmallVector<mlir::Type> blockArgTypes {builder.getIndexType()};
+  blockArgTypes.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
+  for (mlir::Value weight : weights)
+    blockArgTypes.push_back(weight.getType());
+  for (mlir::Value input : inputs)
+    blockArgTypes.push_back(input.getType());
+  llvm::append_range(blockArgTypes, resultTypes);
+  return blockArgTypes;
+}
+
+inline mlir::SmallVector<mlir::Location> getGraphComputeBatchBlockArgLocs(mlir::Location defaultLoc,
+                                                                          mlir::TypeRange resultTypes,
+                                                                          mlir::ValueRange weights,
+                                                                          mlir::ValueRange inputs) {
+  mlir::SmallVector<mlir::Location> blockArgLocs {defaultLoc};
+  blockArgLocs.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
+  for (mlir::Value weight : weights)
+    blockArgLocs.push_back(weight.getLoc());
+  for (mlir::Value input : inputs)
+    blockArgLocs.push_back(input.getLoc());
+  blockArgLocs.append(resultTypes.size(), defaultLoc);
+  return blockArgLocs;
+}
+
 } // namespace detail
 
 template <typename RewriterT>
@@ -87,6 +137,31 @@ inline mlir::Value createSpatConcat(RewriterT& rewriter, mlir::Location loc, int
   return spatial::SpatConcatOp::create(rewriter, loc, outputType, rewriter.getI64IntegerAttr(axis), inputs).getOutput();
 }
 
+template <typename RewriterT>
+spatial::SpatGraphCompute createEmptySpatGraphCompute(RewriterT& rewriter,
+                                                      mlir::Location loc,
+                                                      mlir::TypeRange resultTypes,
+                                                      mlir::ValueRange weights,
+                                                      mlir::ValueRange inputs,
+                                                      mlir::TypeRange blockArgTypes,
+                                                      llvm::ArrayRef<mlir::Location> blockArgLocs) {
+  auto computeOp = spatial::SpatGraphCompute::create(rewriter, loc, resultTypes, weights, inputs);
+  rewriter.createBlock(&computeOp.getBody(), computeOp.getBody().end(), blockArgTypes, blockArgLocs);
+  rewriter.setInsertionPointToStart(&computeOp.getBody().front());
+  return computeOp;
+}
+
+template <typename RewriterT>
+spatial::SpatGraphCompute createEmptySpatGraphCompute(RewriterT& rewriter,
+                                                      mlir::Location loc,
+                                                      mlir::TypeRange resultTypes,
+                                                      mlir::ValueRange weights,
+                                                      mlir::ValueRange inputs) {
+  auto blockArgTypes = detail::getGraphComputeBlockArgTypes(weights, inputs);
+  auto blockArgLocs = detail::getGraphComputeBlockArgLocs(loc, weights, inputs);
+  return createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs, blockArgTypes, blockArgLocs);
+}
+
 /// Builds a `spat.graph_compute` with a fixed number of SSA inputs and erases it if
 /// the body callback reports failure.
 template <size_t NumInputs, typename RewriterT, typename BodyFn>
@@ -97,16 +172,8 @@ auto createSpatGraphCompute(RewriterT& rewriter,
                             mlir::ValueRange inputs,
                             BodyFn&& body) {
   assert(inputs.size() == NumInputs && "NumInputs must match the number of input values");
-  auto computeOp = spatial::SpatGraphCompute::create(rewriter, loc, resultTypes, weights, inputs);
-
-  auto* block = new mlir::Block();
-  for (mlir::Value weight : weights)
-    block->addArgument(weight.getType(), loc);
-  for (mlir::Value input : inputs)
-    block->addArgument(input.getType(), loc);
-
-  computeOp.getBody().push_back(block);
-  rewriter.setInsertionPointToStart(block);
+  auto computeOp = createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs);
+  auto* block = &computeOp.getBody().front();
 
   using BodyResult = detail::InvokeWithBlockArgsResultT<std::decay_t<BodyFn>, std::make_index_sequence<NumInputs>>;
   if constexpr (std::is_same_v<BodyResult, void>) {
@@ -140,16 +207,8 @@ auto createSpatGraphCompute(RewriterT& rewriter,
                             mlir::ValueRange weights,
                             mlir::ValueRange inputs,
                             BodyFn&& body) {
-  auto computeOp = spatial::SpatGraphCompute::create(rewriter, loc, resultTypes, weights, inputs);
-
-  auto* block = new mlir::Block();
-  for (mlir::Value weight : weights)
-    block->addArgument(weight.getType(), loc);
-  for (mlir::Value input : inputs)
-    block->addArgument(input.getType(), loc);
-
-  computeOp.getBody().push_back(block);
-  rewriter.setInsertionPointToStart(block);
+  auto computeOp = createEmptySpatGraphCompute(rewriter, loc, resultTypes, weights, inputs);
+  auto* block = &computeOp.getBody().front();
 
   using BodyResult = detail::InvokeWithValueRangeResultT<std::decay_t<BodyFn>>;
   if constexpr (std::is_same_v<BodyResult, void>) {
@@ -170,14 +229,15 @@ auto createSpatGraphCompute(RewriterT& rewriter,
   }
 }
 
-template <typename RewriterT, typename BodyFn>
-auto createSpatGraphComputeBatch(RewriterT& rewriter,
-                                 mlir::Location loc,
-                                 mlir::TypeRange resultTypes,
-                                 int64_t laneCount,
-                                 mlir::ValueRange weights,
-                                 mlir::ValueRange inputs,
-                                 BodyFn&& body) {
+template <typename RewriterT>
+auto createEmptySpatGraphComputeBatch(RewriterT& rewriter,
+                                      mlir::Location loc,
+                                      mlir::TypeRange resultTypes,
+                                      int64_t laneCount,
+                                      mlir::ValueRange weights,
+                                      mlir::ValueRange inputs,
+                                      mlir::TypeRange blockArgTypes,
+                                      llvm::ArrayRef<mlir::Location> blockArgLocs) {
   if (laneCount <= 0 || laneCount > std::numeric_limits<int32_t>::max())
     return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
 
@@ -186,27 +246,36 @@ auto createSpatGraphComputeBatch(RewriterT& rewriter,
     return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
 
   auto batchOp = spatial::SpatGraphComputeBatch::create(rewriter, loc, resultTypes, *laneCountAttr, weights, inputs);
+  rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), blockArgTypes, blockArgLocs);
+  rewriter.setInsertionPointToStart(&batchOp.getBody().front());
+  return mlir::FailureOr<spatial::SpatGraphComputeBatch>(batchOp);
+}
 
-  mlir::SmallVector<mlir::Type> blockArgTypes {rewriter.getIndexType()};
-  mlir::SmallVector<mlir::Location> blockArgLocs {loc};
-  blockArgTypes.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
-  blockArgLocs.reserve(1 + weights.size() + inputs.size() + resultTypes.size());
-  for (mlir::Value weight : weights) {
-    blockArgTypes.push_back(weight.getType());
-    blockArgLocs.push_back(weight.getLoc());
-  }
-  for (mlir::Value input : inputs) {
-    blockArgTypes.push_back(input.getType());
-    blockArgLocs.push_back(input.getLoc());
-  }
-  for (mlir::Type resultType : resultTypes) {
-    blockArgTypes.push_back(resultType);
-    blockArgLocs.push_back(loc);
-  }
+template <typename RewriterT>
+auto createEmptySpatGraphComputeBatch(RewriterT& rewriter,
+                                      mlir::Location loc,
+                                      mlir::TypeRange resultTypes,
+                                      int64_t laneCount,
+                                      mlir::ValueRange weights,
+                                      mlir::ValueRange inputs) {
+  auto blockArgTypes = detail::getGraphComputeBatchBlockArgTypes(rewriter, resultTypes, weights, inputs);
+  auto blockArgLocs = detail::getGraphComputeBatchBlockArgLocs(loc, resultTypes, weights, inputs);
+  return createEmptySpatGraphComputeBatch(
+    rewriter, loc, resultTypes, laneCount, weights, inputs, blockArgTypes, blockArgLocs);
+}
 
-  auto* block =
-    rewriter.createBlock(&batchOp.getBody(), batchOp.getBody().end(), mlir::TypeRange(blockArgTypes), blockArgLocs);
-  rewriter.setInsertionPointToStart(block);
+template <typename RewriterT, typename BodyFn>
+auto createSpatGraphComputeBatch(RewriterT& rewriter,
+                                 mlir::Location loc,
+                                 mlir::TypeRange resultTypes,
+                                 int64_t laneCount,
+                                 mlir::ValueRange weights,
+                                 mlir::ValueRange inputs,
+                                 BodyFn&& body) {
+  auto batchOp = createEmptySpatGraphComputeBatch(rewriter, loc, resultTypes, laneCount, weights, inputs);
+  if (failed(batchOp))
+    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
+  auto* block = &(*batchOp).getBody().front();
 
   detail::SpatComputeBatchBodyArgs args {
     block->getArgument(0),
@@ -217,18 +286,18 @@ auto createSpatGraphComputeBatch(RewriterT& rewriter,
   using BodyResult = std::invoke_result_t<BodyFn, detail::SpatComputeBatchBodyArgs>;
   if constexpr (std::is_same_v<BodyResult, void>) {
     std::forward<BodyFn>(body)(args);
-    rewriter.setInsertionPointAfter(batchOp);
-    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(batchOp);
+    rewriter.setInsertionPointAfter(*batchOp);
+    return batchOp;
   }
   else {
     auto bodyResult = std::forward<BodyFn>(body)(args);
     if (mlir::failed(bodyResult)) {
-      rewriter.setInsertionPointAfter(batchOp);
-      rewriter.eraseOp(batchOp);
+      rewriter.setInsertionPointAfter(*batchOp);
+      rewriter.eraseOp(*batchOp);
       return mlir::FailureOr<spatial::SpatGraphComputeBatch>(mlir::failure());
     }
-    rewriter.setInsertionPointAfter(batchOp);
-    return mlir::FailureOr<spatial::SpatGraphComputeBatch>(batchOp);
+    rewriter.setInsertionPointAfter(*batchOp);
+    return batchOp;
   }
 }
 
@@ -277,6 +346,52 @@ inline void createParallelInsertSliceIntoBatchOutput(mlir::PatternRewriter& rewr
   mlir::tensor::ParallelInsertSliceOp::create(rewriter, loc, source, dest, offsets, sizes, strides);
 }
 
+inline void publishGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
+                                              mlir::Location loc,
+                                              mlir::Value fragment,
+                                              mlir::Value output,
+                                              mlir::Value physicalSlot) {
+  auto fragmentType = mlir::cast<mlir::RankedTensorType>(fragment.getType());
+  mlir::SmallVector<mlir::OpFoldResult> offsets {physicalSlot};
+  mlir::SmallVector<mlir::OpFoldResult> sizes {rewriter.getIndexAttr(1)};
+  mlir::SmallVector<mlir::OpFoldResult> strides {rewriter.getIndexAttr(1)};
+  for (int64_t dim : fragmentType.getShape()) {
+    offsets.push_back(rewriter.getIndexAttr(0));
+    sizes.push_back(rewriter.getIndexAttr(dim));
+    strides.push_back(rewriter.getIndexAttr(1));
+  }
+  createParallelInsertSliceIntoBatchOutput(rewriter, loc, fragment, output, offsets, sizes, strides);
+}
+
+inline mlir::FailureOr<mlir::Value>
+extractGraphBatchPhysicalFragment(mlir::PatternRewriter& rewriter,
+                                  mlir::Location loc,
+                                  mlir::Value physicalBatch,
+                                  mlir::OpFoldResult slot,
+                                  mlir::RankedTensorType fragmentType) {
+  if (fragmentType.getRank() == 0)
+    return mlir::failure();
+  auto physicalType = mlir::dyn_cast<mlir::RankedTensorType>(physicalBatch.getType());
+  if (!physicalType || physicalType.getRank() != fragmentType.getRank() + 1)
+    return mlir::failure();
+  mlir::SmallVector<int64_t> selectedShape {1};
+  llvm::append_range(selectedShape, fragmentType.getShape());
+  auto selectedType = mlir::RankedTensorType::get(selectedShape, fragmentType.getElementType(), fragmentType.getEncoding());
+  mlir::SmallVector<mlir::OpFoldResult> offsets {slot};
+  mlir::SmallVector<mlir::OpFoldResult> sizes {rewriter.getIndexAttr(1)};
+  mlir::SmallVector<mlir::OpFoldResult> strides {rewriter.getIndexAttr(1)};
+  for (int64_t dim : fragmentType.getShape()) {
+    offsets.push_back(rewriter.getIndexAttr(0));
+    sizes.push_back(rewriter.getIndexAttr(dim));
+    strides.push_back(rewriter.getIndexAttr(1));
+  }
+  mlir::Value selected = mlir::tensor::ExtractSliceOp::create(rewriter, loc, selectedType, physicalBatch, offsets, sizes, strides);
+  mlir::SmallVector<mlir::ReassociationIndices> reassociation {{0, 1}};
+  for (int64_t dim = 2; dim <= fragmentType.getRank(); ++dim)
+    reassociation.push_back({dim});
+  return mlir::tensor::CollapseShapeOp::create(rewriter, loc, fragmentType, selected, reassociation).getResult();
+}
+
 template <typename BodyFn>
 mlir::Value materializeOrComputeUnary(mlir::Value input,
                                       mlir::RankedTensorType resultType,

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

@@ -0,0 +1,211 @@
+#include "llvm/ADT/SmallVector.h"
+
+#include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+
+RankedTensorType getRowStripFragmentType(RankedTensorType logicalType) {
+  return RankedTensorType::get({logicalType.getDimSize(0), logicalType.getDimSize(1), 1, logicalType.getDimSize(3)},
+                               logicalType.getElementType(),
+                               logicalType.getEncoding());
+}
+
+RankedTensorType getRowStripStorageType(RankedTensorType logicalType) {
+  return spatial::getGraphBatchPhysicalResultType(logicalType.getDimSize(2), getRowStripFragmentType(logicalType));
+}
+
+std::pair<SmallVector<int64_t>, SmallVector<int64_t>> buildRowStripMetadata(RankedTensorType type) {
+  SmallVector<int64_t> offsets;
+  SmallVector<int64_t> sizes;
+  const int64_t channels = type.getDimSize(1);
+  const int64_t height = type.getDimSize(2);
+  const int64_t width = type.getDimSize(3);
+  offsets.reserve(height * 4);
+  sizes.reserve(height * 4);
+  for (int64_t row = 0; row < height; ++row) {
+    offsets.append({0, 0, row, 0});
+    sizes.append({1, channels, 1, width});
+  }
+  return {offsets, sizes};
+}
+
+Value extractRowStripFragment(Value storage,
+                              RankedTensorType logicalType,
+                              OpFoldResult row,
+                              PatternRewriter& rewriter,
+                              Location loc) {
+  return *extractGraphBatchPhysicalFragment(rewriter, loc, storage, row, getRowStripFragmentType(logicalType));
+}
+
+void insertRowStripFragment(Value fragment,
+                            Value output,
+                            RankedTensorType logicalType,
+                            OpFoldResult row,
+                            PatternRewriter& rewriter,
+                            Location loc) {
+  assert(fragment.getType() == getRowStripFragmentType(logicalType));
+  assert(output.getType() == getRowStripStorageType(logicalType));
+  auto slot = dyn_cast<Value>(row);
+  assert(slot && "row-strip graph publication requires a dynamic physical slot");
+  publishGraphBatchPhysicalFragment(rewriter, loc, fragment, output, slot);
+}
+
+FailureOr<Value> createPerChannelConstantFragment(DenseElementsAttr denseAttr,
+                                                  RankedTensorType fragmentType,
+                                                  PatternRewriter& rewriter) {
+  FailureOr<SmallVector<Attribute>> channelValues = getBiasChannelValues(denseAttr, fragmentType);
+  if (failed(channelValues))
+    return failure();
+
+  SmallVector<Attribute> values;
+  values.reserve(fragmentType.getNumElements());
+  for (int64_t n = 0; n < fragmentType.getDimSize(0); ++n)
+    for (int64_t channel = 0; channel < fragmentType.getDimSize(1); ++channel)
+      for (int64_t h = 0; h < fragmentType.getDimSize(2); ++h)
+        for (int64_t w = 0; w < fragmentType.getDimSize(3); ++w)
+          values.push_back((*channelValues)[channel]);
+
+  auto attr = DenseElementsAttr::get(fragmentType, values);
+  return getOrCreateConstant(rewriter, rewriter.getInsertionBlock()->getParentOp(), attr, fragmentType);
+}
+
+FailureOr<Value> createRowStripStorageFromRows(Value rows,
+                                               RankedTensorType logicalType,
+                                               PatternRewriter& rewriter,
+                                               Location loc) {
+  auto rowsType = dyn_cast<RankedTensorType>(rows.getType());
+  if (!rowsType || !rowsType.hasStaticShape() || rowsType.getRank() != 2)
+    return failure();
+  if (!logicalType || !logicalType.hasStaticShape() || logicalType.getRank() != 4)
+    return failure();
+  if (logicalType.getDimSize(0) != 1)
+    return failure();
+  if (rowsType.getElementType() != logicalType.getElementType())
+    return failure();
+
+  const int64_t channels = logicalType.getDimSize(1);
+  const int64_t height = logicalType.getDimSize(2);
+  const int64_t width = logicalType.getDimSize(3);
+  if (rowsType.getDimSize(0) != height * width)
+    return failure();
+  if (rowsType.getDimSize(1) != channels)
+    return failure();
+
+  auto rowSliceType = RankedTensorType::get({width, channels}, logicalType.getElementType(), rowsType.getEncoding());
+  auto channelWidthType = RankedTensorType::get({channels, width}, logicalType.getElementType(), rowsType.getEncoding());
+  auto fragmentType = getRowStripFragmentType(logicalType);
+  auto storageType = getRowStripStorageType(logicalType);
+  auto batchOp = createSpatComputeBatch(
+    rewriter, loc, TypeRange {storageType}, height, {}, ValueRange {rows}, [&](detail::SpatComputeBatchBodyArgs args) {
+      Operation* anchorOp = rewriter.getInsertionBlock()->getParentOp();
+      Value rowStart = affineMulConst(rewriter, loc, args.lane, width, anchorOp);
+      SmallVector<OpFoldResult> rowOffsets {rowStart, rewriter.getIndexAttr(0)};
+      SmallVector<OpFoldResult> rowSizes {rewriter.getIndexAttr(width), rewriter.getIndexAttr(channels)};
+      Value rowSlice = tensor::ExtractSliceOp::create(
+        rewriter, loc, rowSliceType, args.inputs.front(), rowOffsets, rowSizes, getUnitStrides(rewriter, 2));
+      Value channelWidth = ONNXTransposeOp::create(
+        rewriter, loc, channelWidthType, rowSlice, rewriter.getI64ArrayAttr({1, 0})).getResult();
+      Value fragment = tensor::ExpandShapeOp::create(
+        rewriter, loc, fragmentType, channelWidth, SmallVector<ReassociationIndices> {{0, 1}, {2, 3}});
+      insertRowStripFragment(fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
+      return success();
+    });
+  if (failed(batchOp))
+    return failure();
+  return batchOp->getResult(0);
+}
+
+FailureOr<Value>
+createRowStripAssemblyBlueprint(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
+  auto storageType = dyn_cast<RankedTensorType>(storage.getType());
+  if (!storageType || storageType != getRowStripStorageType(logicalType))
+    return failure();
+
+  auto [offsets, sizes] = buildRowStripMetadata(logicalType);
+  int64_t height = logicalType.getDimSize(2);
+  SmallVector<int64_t> operandIndices(height, 0), sourceSlots, sourceOffsets(height, 0), strides(height * 4, 1);
+  for (int64_t row = 0; row < height; ++row)
+    sourceSlots.push_back(row);
+  return spatial::SpatBlueprintOp::create(rewriter, loc, logicalType, storage, ValueRange {},
+    rewriter.getStringAttr("nchw"), rewriter.getStringAttr("nchw_row_strip"),
+    rewriter.getDenseI64ArrayAttr(offsets), rewriter.getDenseI64ArrayAttr(sizes),
+    rewriter.getStringAttr("nchw_row_strip_fragments"), rewriter.getStringAttr("fragment_assembly"),
+    rewriter.getDenseI64ArrayAttr(operandIndices), rewriter.getDenseI64ArrayAttr(sourceSlots),
+    rewriter.getDenseI64ArrayAttr(sourceOffsets), rewriter.getDenseI64ArrayAttr(strides),
+    rewriter.getStringAttr("disjoint"), rewriter.getStringAttr("complete")).getOutput();
+}
+
+FailureOr<Value>
+applyRowStripRelu(Value storage, RankedTensorType logicalType, PatternRewriter& rewriter, Location loc) {
+  auto fragmentType = getRowStripFragmentType(logicalType);
+  auto storageType = getRowStripStorageType(logicalType);
+  auto batchOp = createSpatComputeBatch(rewriter,
+                                        loc,
+                                        TypeRange {storageType},
+                                        logicalType.getDimSize(2),
+                                        {},
+                                        ValueRange {storage},
+                                        [&](detail::SpatComputeBatchBodyArgs args) {
+                                          Value fragment =
+                                            extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
+                                          fragment = spatial::SpatReluOp::create(rewriter, loc, fragmentType, fragment).getResult();
+                                          insertRowStripFragment(
+                                            fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
+                                          return success();
+                                        });
+  if (failed(batchOp))
+    return failure();
+  return batchOp->getResult(0);
+}
+
+FailureOr<Value>
+applyRowStripBiasAdd(Value storage, RankedTensorType logicalType, Value bias, PatternRewriter& rewriter, Location loc) {
+  DenseElementsAttr denseAttr;
+  if (!isSupportedBiasAddValue(bias, logicalType, &denseAttr))
+    return failure();
+  auto fragmentType = getRowStripFragmentType(logicalType);
+  auto storageType = getRowStripStorageType(logicalType);
+  auto batchOp = createSpatComputeBatch(rewriter,
+                                        loc,
+                                        TypeRange {storageType},
+                                        logicalType.getDimSize(2),
+                                        {},
+                                        ValueRange {storage},
+                                        [&](detail::SpatComputeBatchBodyArgs args) {
+                                          Value fragment =
+                                            extractRowStripFragment(args.inputs.front(), logicalType, args.lane, rewriter, loc);
+                                          Value constant;
+                                          if (denseAttr.isSplat()) {
+                                            constant = getOrCreateConstant(
+                                              rewriter,
+                                              rewriter.getInsertionBlock()->getParentOp(),
+                                              DenseElementsAttr::get(fragmentType, denseAttr.getSplatValue<Attribute>()),
+                                              fragmentType);
+                                          }
+                                          else {
+                                            FailureOr<Value> perChannel =
+                                              createPerChannelConstantFragment(denseAttr, fragmentType, rewriter);
+                                            if (failed(perChannel))
+                                              return failure();
+                                            constant = *perChannel;
+                                          }
+                                          fragment =
+                                            spatial::SpatVAddOp::create(rewriter, loc, fragmentType, fragment, constant).getResult();
+                                          insertRowStripFragment(
+                                            fragment, args.outputs.front(), logicalType, args.lane, rewriter, loc);
+                                          return success();
+                                        });
+  if (failed(batchOp))
+    return failure();
+  return batchOp->getResult(0);
+}
+
+} // namespace onnx_mlir

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

@@ -0,0 +1,69 @@
+#pragma once
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/IR/PatternMatch.h"
+
+namespace onnx_mlir {
+
+inline constexpr llvm::StringLiteral kRowStripIndexMap = "nchw_row_strip_fragments";
+
+struct RowStripPhysicalValue {
+  mlir::Value storage;
+  mlir::RankedTensorType logicalType;
+  llvm::SmallVector<int64_t, 16> fragmentOffsets;
+  llvm::SmallVector<int64_t, 16> fragmentSizes;
+};
+
+std::pair<llvm::SmallVector<int64_t>, llvm::SmallVector<int64_t>>
+buildRowStripMetadata(mlir::RankedTensorType type);
+
+mlir::RankedTensorType getRowStripFragmentType(mlir::RankedTensorType logicalType);
+
+mlir::RankedTensorType getRowStripStorageType(mlir::RankedTensorType logicalType);
+
+llvm::SmallVector<mlir::OpFoldResult> buildRowStripFragmentOffsets(mlir::PatternRewriter& rewriter,
+                                                                   mlir::OpFoldResult row);
+
+llvm::SmallVector<mlir::OpFoldResult> buildRowStripFragmentSizes(mlir::PatternRewriter& rewriter,
+                                                                 mlir::RankedTensorType logicalType);
+
+mlir::Value extractRowStripFragment(mlir::Value storage,
+                                    mlir::RankedTensorType logicalType,
+                                    mlir::OpFoldResult row,
+                                    mlir::PatternRewriter& rewriter,
+                                    mlir::Location loc);
+
+void insertRowStripFragment(mlir::Value fragment,
+                            mlir::Value output,
+                            mlir::RankedTensorType logicalType,
+                            mlir::OpFoldResult row,
+                            mlir::PatternRewriter& rewriter,
+                            mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> createPerChannelConstantFragment(mlir::DenseElementsAttr denseAttr,
+                                                              mlir::RankedTensorType fragmentType,
+                                                              mlir::PatternRewriter& rewriter);
+
+mlir::FailureOr<mlir::Value> createRowStripStorageFromRows(mlir::Value rows,
+                                                           mlir::RankedTensorType logicalType,
+                                                           mlir::PatternRewriter& rewriter,
+                                                           mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> createRowStripAssemblyBlueprint(mlir::Value storage,
+                                                               mlir::RankedTensorType logicalType,
+                                                               mlir::PatternRewriter& rewriter,
+                                                               mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> applyRowStripRelu(mlir::Value storage,
+                                               mlir::RankedTensorType logicalType,
+                                               mlir::PatternRewriter& rewriter,
+                                               mlir::Location loc);
+
+mlir::FailureOr<mlir::Value> applyRowStripBiasAdd(mlir::Value storage,
+                                                  mlir::RankedTensorType logicalType,
+                                                  mlir::Value bias,
+                                                  mlir::PatternRewriter& rewriter,
+                                                  mlir::Location loc);
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/LowerSpatialPlansPass.cpp

@@ -11,9 +11,12 @@
 #include "llvm/ADT/SmallPtrSet.h"
 
 #include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
+#include "mlir/Transforms/Passes.h"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -28,14 +31,6 @@ namespace {
 static constexpr StringLiteral kDenseLayout = "dense_nchw";
 static constexpr StringLiteral kRowStripLayout = "nchw_row_strip";
 
-struct RowStripPhysicalValue {
-  Value physicalValue;
-  RankedTensorType logicalType;
-  SmallVector<int64_t, 16> fragmentOffsets;
-  SmallVector<int64_t, 16> fragmentSizes;
-  std::string indexMap;
-};
-
 static FailureOr<RowStripPhysicalValue> getRowStripValue(llvm::DenseMap<Value, RowStripPhysicalValue>& rowStripValues,
                                                          Value value) {
   auto it = rowStripValues.find(value);
@@ -45,112 +40,42 @@ static FailureOr<RowStripPhysicalValue> getRowStripValue(llvm::DenseMap<Value, R
 }
 
 static FailureOr<RowStripPhysicalValue> buildRowStripValue(spatial::SpatBlueprintOp blueprint,
-                                                           Value physicalValue) {
+                                                           Value storage) {
   auto logicalType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
   if (!logicalType)
     return blueprint.emitOpError("requires ranked logical output type"), failure();
   RowStripPhysicalValue value;
-  value.physicalValue = physicalValue;
+  value.storage = storage;
   value.logicalType = logicalType;
   value.fragmentOffsets.append(blueprint.getFragmentOffsets().begin(), blueprint.getFragmentOffsets().end());
   value.fragmentSizes.append(blueprint.getFragmentSizes().begin(), blueprint.getFragmentSizes().end());
-  value.indexMap = blueprint.getIndexMap().str();
+  if (blueprint.getIndexMap() != kRowStripIndexMap)
+    return blueprint.emitOpError("requires the canonical row-strip index map"), failure();
+  auto storageType = dyn_cast<RankedTensorType>(storage.getType());
+  if (!storageType || storageType != getRowStripStorageType(logicalType))
+    return blueprint.emitOpError("requires physical row-strip fragment storage"), failure();
   return value;
 }
 
 static FailureOr<Value>
 lowerRowStripRelu(const RowStripPhysicalValue& input, spatial::SpatReluPlanOp planOp, PatternRewriter& rewriter) {
-  auto packedType = cast<RankedTensorType>(input.physicalValue.getType());
-  auto computeOp =
-    createSpatCompute<1>(rewriter, planOp.getLoc(), TypeRange {packedType}, {}, input.physicalValue, [&](Value x) {
-      auto relu = spatial::SpatReluOp::create(rewriter, planOp.getLoc(), packedType, x);
-      spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), relu.getResult());
-    });
-  return computeOp.getResult(0);
+  return applyRowStripRelu(input.storage, input.logicalType, rewriter, planOp.getLoc());
+}
+
+static FailureOr<Value> lowerRowStripBiasAdd(const RowStripPhysicalValue& input,
+                                             spatial::SpatBiasAddPlanOp planOp,
+                                             PatternRewriter& rewriter) {
+  return applyRowStripBiasAdd(input.storage, input.logicalType, planOp.getBias(), rewriter, planOp.getLoc());
 }
 
 static FailureOr<Value>
 materializeRowStripToDense(const RowStripPhysicalValue& rowStripValue, Location loc, PatternRewriter& rewriter) {
-  auto packedType = dyn_cast<RankedTensorType>(rowStripValue.physicalValue.getType());
-  if (!packedType || packedType.getRank() != 3 || !packedType.hasStaticShape())
-    return failure();
   if (rowStripValue.logicalType.getRank() != 4 || !rowStripValue.logicalType.hasStaticShape())
     return failure();
-  if (rowStripValue.indexMap != "packed_hwc_rows_to_nchw")
+  auto [expectedOffsets, expectedSizes] = buildRowStripMetadata(rowStripValue.logicalType);
+  if (!llvm::equal(rowStripValue.fragmentOffsets, expectedOffsets) || !llvm::equal(rowStripValue.fragmentSizes, expectedSizes))
     return failure();
-
-  const int64_t rank = rowStripValue.logicalType.getRank();
-  const int64_t fragmentCount = rowStripValue.fragmentOffsets.size() / rank;
-  const int64_t packedWidth = packedType.getDimSize(1);
-  const int64_t packedChannels = packedType.getDimSize(2);
-  if (fragmentCount != packedType.getDimSize(0))
-    return failure();
-  for (int64_t fragmentIndex = 0; fragmentIndex < fragmentCount; ++fragmentIndex) {
-    if (rowStripValue.fragmentOffsets[fragmentIndex * rank + 0] != 0
-        || rowStripValue.fragmentOffsets[fragmentIndex * rank + 1] != 0
-        || rowStripValue.fragmentOffsets[fragmentIndex * rank + 2] != fragmentIndex
-        || rowStripValue.fragmentOffsets[fragmentIndex * rank + 3] != 0)
-      return failure();
-    if (rowStripValue.fragmentSizes[fragmentIndex * rank + 0] != 1
-        || rowStripValue.fragmentSizes[fragmentIndex * rank + 1] != packedChannels
-        || rowStripValue.fragmentSizes[fragmentIndex * rank + 2] != 1
-        || rowStripValue.fragmentSizes[fragmentIndex * rank + 3] != packedWidth)
-      return failure();
-  }
-
-  auto packedSliceType =
-    RankedTensorType::get({1, packedWidth, packedChannels}, packedType.getElementType(), packedType.getEncoding());
-  auto expandedType =
-    RankedTensorType::get({1, 1, packedWidth, packedChannels}, packedType.getElementType(), packedType.getEncoding());
-  auto logicalFragmentType =
-    RankedTensorType::get({1, packedChannels, 1, packedWidth}, packedType.getElementType(), packedType.getEncoding());
-  auto batchOp = createSpatComputeBatch(
-    rewriter,
-    loc,
-    TypeRange {rowStripValue.logicalType},
-    fragmentCount,
-    {},
-    ValueRange {rowStripValue.physicalValue},
-    [&](detail::SpatComputeBatchBodyArgs args) {
-      SmallVector<OpFoldResult> packedOffsets {args.lane, rewriter.getIndexAttr(0), rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> packedSizes {
-        rewriter.getIndexAttr(1), rewriter.getIndexAttr(packedWidth), rewriter.getIndexAttr(packedChannels)};
-      Value packedSlice = tensor::ExtractSliceOp::create(
-        rewriter, loc, packedSliceType, args.inputs.front(), packedOffsets, packedSizes, getUnitStrides(rewriter, 3));
-
-      Value expanded = tensor::ExpandShapeOp::create(rewriter,
-                                                     loc,
-                                                     expandedType,
-                                                     packedSlice,
-                                                     SmallVector<ReassociationIndices> {
-                                                       {0, 1},
-                                                       {2},
-                                                       {3}
-      });
-      Value transposeInit =
-        tensor::EmptyOp::create(rewriter, loc, logicalFragmentType.getShape(), logicalFragmentType.getElementType());
-      Value logicalFragment =
-        linalg::TransposeOp::create(rewriter, loc, expanded, transposeInit, SmallVector<int64_t> {0, 3, 1, 2})
-          .getResult()[0];
-
-      SmallVector<OpFoldResult> logicalOffsets {
-        rewriter.getIndexAttr(0), rewriter.getIndexAttr(0), args.lane, rewriter.getIndexAttr(0)};
-      SmallVector<OpFoldResult> logicalSizes {rewriter.getIndexAttr(1),
-                                              rewriter.getIndexAttr(packedChannels),
-                                              rewriter.getIndexAttr(1),
-                                              rewriter.getIndexAttr(packedWidth)};
-      createParallelInsertSliceIntoBatchOutput(rewriter,
-                                               loc,
-                                               logicalFragment,
-                                               args.outputs.front(),
-                                               logicalOffsets,
-                                               logicalSizes,
-                                               getUnitStrides(rewriter, 4));
-      return success();
-    });
-  if (failed(batchOp))
-    return failure();
-  return batchOp->getResult(0);
+  return createRowStripAssemblyBlueprint(rowStripValue.storage, rowStripValue.logicalType, rewriter, loc);
 }
 
 struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, OperationPass<ModuleOp>> {
@@ -193,7 +118,7 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
           rewriter.setInsertionPoint(planOp);
           FailureOr<Value> lowered = lowerSelectedConv2DPlan(
             planOp,
-            succeeded(rowStripInput) ? std::optional<Value> {rowStripInput->physicalValue} : std::nullopt,
+            succeeded(rowStripInput) ? std::optional<Value> {rowStripInput->storage} : std::nullopt,
             /*emitRowStripLayout=*/true,
             rewriter);
           if (failed(lowered)) {
@@ -265,6 +190,64 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
         rewriter.replaceOp(planOp, computeOp.getResults());
         continue;
       }
+      if (auto planOp = dyn_cast<spatial::SpatBiasAddPlanOp>(&op)) {
+        if (succeeded(getRowStripValue(rowStripValues, planOp.getInput()))) {
+          auto outputBlueprint = llvm::find_if(planOp.getResult().getUsers(), [](Operation* user) {
+            auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(user);
+            return blueprint && blueprint.getPhysicalLayout() == kRowStripLayout;
+          });
+          if (outputBlueprint == planOp.getResult().getUsers().end()) {
+            planOp.emitOpError("row-strip bias_add plan requires a row-strip blueprint result");
+            signalPassFailure();
+            return;
+          }
+
+          FailureOr<RowStripPhysicalValue> input = getRowStripValue(rowStripValues, planOp.getInput());
+          rewriter.setInsertionPoint(planOp);
+          FailureOr<Value> lowered = lowerRowStripBiasAdd(*input, planOp, rewriter);
+          if (failed(lowered)) {
+            planOp.emitOpError("failed to lower selected row-strip Spatial bias_add plan");
+            signalPassFailure();
+            return;
+          }
+          auto blueprint = cast<spatial::SpatBlueprintOp>(*outputBlueprint);
+          FailureOr<RowStripPhysicalValue> output = buildRowStripValue(blueprint, *lowered);
+          if (failed(output)) {
+            signalPassFailure();
+            return;
+          }
+          rowStripValues[blueprint.getResult()] = *output;
+          eraseAfterLowering.insert(planOp);
+          eraseAfterLowering.insert(blueprint);
+          continue;
+        }
+
+        auto resultType = dyn_cast<RankedTensorType>(planOp.getOutput().getType());
+        if (!resultType) {
+          planOp.emitOpError("requires ranked output type");
+          signalPassFailure();
+          return;
+        }
+        rewriter.setInsertionPoint(planOp);
+        FailureOr<Value> denseBias = materializeDenseBiasAddTensor(planOp.getBias(), resultType, rewriter, planOp.getLoc());
+        if (failed(denseBias)) {
+          planOp.emitOpError("failed to materialize dense Conv-style bias");
+          signalPassFailure();
+          return;
+        }
+        auto computeOp = createSpatCompute<2>(rewriter,
+                                              planOp.getLoc(),
+                                              planOp.getOutput().getType(),
+                                              {},
+                                              ValueRange {planOp.getInput(), *denseBias},
+                                              [&](Value x, Value y) {
+                                                auto added = spatial::SpatVAddOp::create(
+                                                  rewriter, planOp.getLoc(), planOp.getOutput().getType(), x, y);
+                                                spatial::SpatYieldOp::create(rewriter, planOp.getLoc(), added.getResult());
+                                              });
+        rewriter.replaceOp(planOp, computeOp.getResults());
+        continue;
+      }
       if (auto materializeOp = dyn_cast<spatial::SpatMaterializeLayoutOp>(&op)) {
         if (materializeOp.getSourcePhysicalLayout() == kDenseLayout
             && materializeOp.getTargetPhysicalLayout() == kDenseLayout) {
@@ -294,6 +277,8 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
         continue;
       }
       if (auto blueprintOp = dyn_cast<spatial::SpatBlueprintOp>(&op)) {
+        if (std::optional<StringRef> mode = blueprintOp.getMode(); mode && *mode == "fragment_assembly")
+          continue;
         if (blueprintOp.getPhysicalLayout() == kDenseLayout) {
           rewriter.replaceOp(blueprintOp, blueprintOp.getInput());
           continue;
@@ -383,19 +368,31 @@ struct LowerSpatialPlansPass final : PassWrapper<LowerSpatialPlansPass, Operatio
     moduleOp.walk([&](Operation* op) {
       if (isa<ONNXEntryPointOp>(op))
         return;
-      if (isa<spatial::SpatConv2DPlanOp,
-              spatial::SpatReluPlanOp,
-              spatial::SpatBlueprintOp,
-              spatial::SpatMaterializeLayoutOp>(op)
+      if (auto blueprint = dyn_cast<spatial::SpatBlueprintOp>(op)) {
+        if (std::optional<StringRef> mode = blueprint.getMode(); mode && *mode == "fragment_assembly")
+          return;
+        op->emitOpError("planning blueprint must not remain after LowerSpatialPlans");
+        hasIllegalOps = true;
+      } else if (isa<spatial::SpatConv2DPlanOp,
+                     spatial::SpatBiasAddPlanOp,
+                     spatial::SpatReluPlanOp,
+                     spatial::SpatMaterializeLayoutOp>(op)
           || op->getDialect()->getNamespace() == "onnx") {
         op->emitOpError("operation must not remain after LowerSpatialPlans");
         hasIllegalOps = true;
       }
     });
-    if (hasIllegalOps)
+
+    PassManager canonicalizationPM(ctx);
+    canonicalizationPM.addPass(createCanonicalizerPass());
+    if (failed(canonicalizationPM.run(moduleOp)))
+      moduleOp.emitWarning("failed to run LowerSpatialPlansPass canonicalization; continuing");
+
+    if (hasIllegalOps) {
       signalPassFailure();
-    else
-      dumpModule(moduleOp, "spatial1_premerge");
+    } else {
+      dumpModule(moduleOp, "spatial1_graph");
+    }
 
     if (!verifyLogicalPhase("at the end of LowerSpatialPlans"))
       return;

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialPass.cpp

@@ -13,6 +13,7 @@
 
 #include "Common/Common.hpp"
 #include "Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ComputeRegionBuilder.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/CompileTime.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
@@ -45,11 +46,12 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
   SmallVector<spatial::SpatGraphCompute> computes(funcOp.getOps<spatial::SpatGraphCompute>());
   SmallVector<spatial::SpatGraphComputeBatch> computeBatches(funcOp.getOps<spatial::SpatGraphComputeBatch>());
   SmallVector<spatial::SpatConv2DPlanOp> convPlans(funcOp.getOps<spatial::SpatConv2DPlanOp>());
+  SmallVector<spatial::SpatBiasAddPlanOp> biasAddPlans(funcOp.getOps<spatial::SpatBiasAddPlanOp>());
   SmallVector<spatial::SpatReluPlanOp> reluPlans(funcOp.getOps<spatial::SpatReluPlanOp>());
   SmallVector<spatial::SpatBlueprintOp> blueprints(funcOp.getOps<spatial::SpatBlueprintOp>());
   SmallVector<spatial::SpatMaterializeLayoutOp> materializers(funcOp.getOps<spatial::SpatMaterializeLayoutOp>());
-  if (!computes.empty() || !computeBatches.empty() || !convPlans.empty() || !reluPlans.empty() || !blueprints.empty()
-      || !materializers.empty()) {
+  if (!computes.empty() || !computeBatches.empty() || !convPlans.empty() || !biasAddPlans.empty() || !reluPlans.empty()
+      || !blueprints.empty() || !materializers.empty()) {
     return;
   }
 
@@ -65,9 +67,9 @@ static void populateEmptyFunction(func::FuncOp funcOp) {
     sourceLocs.push_back(source.getLoc());
   }
 
-  auto newCompute = spatial::SpatGraphCompute::create(
-    rewriter, returnOp.getLoc(), returnOp.getOperandTypes(), funcOp.getArguments(), {}, {});
-  auto* newBlock = rewriter.createBlock(&newCompute.getBody(), newCompute.getBody().end(), sourceTypes, sourceLocs);
+  auto newCompute = createEmptySpatGraphCompute(
+    rewriter, returnOp.getLoc(), returnOp.getOperandTypes(), {}, funcOp.getArguments(), sourceTypes, sourceLocs);
+  auto* newBlock = &newCompute.getBody().front();
   for (auto [blockArg, computeArg] : llvm::zip(newBlock->getArguments(), newCompute.getOperands()))
     mapper.map(computeArg, blockArg);
   newCompute.getProperties().setOperandSegmentSizes({0, static_cast<int>(sourceTypes.size())});
@@ -103,7 +105,7 @@ void ONNXToSpatialPass::runOnOperation() {
                             affine::AffineDialect,
                             arith::ArithDialect,
                             scf::SCFDialect>();
-  preTarget.addIllegalOp<ONNXConstantOp, ONNXFlattenOp>();
+  preTarget.addIllegalOp<ONNXConstantOp>();
 
   RewritePatternSet prePatterns(ctx);
   populatePrePatterns(prePatterns, ctx);
@@ -142,6 +144,7 @@ void ONNXToSpatialPass::runOnOperation() {
   target.addIllegalOp<ONNXSigmoidOp>();
   target.addIllegalOp<ONNXSoftmaxOp>();
   target.addIllegalOp<ONNXConcatOp>();
+  target.addIllegalOp<ONNXFlattenOp>();
   target.addIllegalOp<ONNXGatherOp>();
   target.addIllegalOp<ONNXReshapeOp>();
   target.addIllegalOp<ONNXResizeOp>();
@@ -173,11 +176,6 @@ void ONNXToSpatialPass::runOnOperation() {
                                   arith::ArithDialect,
                                   scf::SCFDialect>();
 
-  PassManager cleanupPM(ctx);
-  cleanupPM.addPass(createCanonicalizerPass());
-  if (failed(cleanupPM.run(moduleOp)))
-    moduleOp.emitWarning("failed to run ONNX-to-Spatial canonicalization cleanup; continuing");
-
   annotateWeightsConstants(*entryFunc);
 
     if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
@@ -213,13 +211,18 @@ void ONNXToSpatialPass::runOnOperation() {
 
   populateEmptyFunction(*entryFunc);
 
-    if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
+  PassManager canonicalizationPM(ctx);
+  canonicalizationPM.addPass(createCanonicalizerPass());
+  if (failed(canonicalizationPM.run(moduleOp)))
+    moduleOp.emitWarning("failed to run ONNXToSpatial canonicalization; continuing");
+
+  dumpModule(moduleOp, "spatial0");
+
+  if (failed(verifyLogicalSpatialGraphInvariants(*entryFunc))) {
       moduleOp.emitError("logical Spatial graph verification failed after ONNX-to-Spatial");
       signalPassFailure();
       return;
     }
-  dumpModule(moduleOp, "spatial0");
-
   if (failed(verifyONNXToSpatial(*entryFunc))) {
     moduleOp.emitError("ONNX-to-Spatial host legality verification failed");
     signalPassFailure();

src/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.cpp

@@ -56,13 +56,18 @@ bool isLegalExternalCapture(Value value, Region& region) {
   return definingOp && definingOp->hasTrait<OpTrait::ConstantLike>();
 }
 
+bool isRecordedDeferredCommunicationSource(Operation* op, Value value) {
+  auto transfer = dyn_cast<spatial::SpatDeferredCommunicationOp>(op);
+  return transfer && llvm::is_contained(transfer.getSources(), value);
+}
+
 template <typename ComputeOpTy>
 void verifyComputeBodyCaptures(ComputeOpTy compute, StringRef kind, pim::CappedDiagnosticReporter& diagnostics) {
   Region& body = compute.getBody();
   body.walk([&](Operation* nestedOp) {
     for (OpOperand& operand : nestedOp->getOpOperands()) {
       Value value = operand.get();
-      if (isLegalExternalCapture(value, body))
+      if (isLegalExternalCapture(value, body) || isRecordedDeferredCommunicationSource(nestedOp, value))
         continue;
 
       Operation* definingOp = value.getDefiningOp();
@@ -90,21 +95,29 @@ bool isLegalHostBackedValue(Value value) {
   return definingOp->getDialect()->getNamespace() != "spat";
 }
 
+bool isScheduledPhase1Value(Value value) {
+  Operation* definingOp = value.getDefiningOp();
+  return isa_and_nonnull<spatial::SpatScheduledCompute, spatial::SpatScheduledComputeBatch>(definingOp);
+}
+
 template <typename ComputeOpTy>
 void verifyScheduledInputs(ComputeOpTy compute,
                            bool allowChannelReceiveInputs,
                            StringRef kind,
                            pim::CappedDiagnosticReporter& diagnostics) {
   for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs())) {
+    size_t currentInputIndex = inputIndex;
     Operation* definingOp = input.getDefiningOp();
     if (allowChannelReceiveInputs && isa_and_nonnull<spatial::SpatChannelReceiveOp>(definingOp))
       continue;
+    if (isScheduledPhase1Value(input))
+      continue;
     if (isLegalHostBackedValue(input))
       continue;
 
     diagnostics.report(compute.getOperation(), [&](Operation* illegalOp) {
       InFlightDiagnostic diag = illegalOp->emitOpError()
-                                << kPhaseMarker << " " << kind << " input #" << inputIndex
+                                << kPhaseMarker << " " << kind << " input #" << currentInputIndex
                                 << (allowChannelReceiveInputs ? " must come from the host or explicit spat.channel_receive"
                                                               : " must come from the host");
       if (definingOp)
@@ -132,6 +145,7 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
             spatial::SpatGraphCompute,
             spatial::SpatGraphComputeBatch,
             spatial::SpatConv2DPlanOp,
+            spatial::SpatBiasAddPlanOp,
             spatial::SpatReluPlanOp,
             spatial::SpatBlueprintOp,
             spatial::SpatMaterializeLayoutOp>(&op)) {
@@ -162,9 +176,9 @@ void verifyLogicalTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter
 
 void verifyScheduledTopLevelOps(func::FuncOp funcOp, pim::CappedDiagnosticReporter& diagnostics) {
   for (Operation& op : funcOp.getOps()) {
-    if (isa<spatial::SpatGraphCompute, spatial::SpatGraphComputeBatch>(&op)) {
+    if (isa<spatial::SpatChannelSendOp, spatial::SpatChannelReceiveOp>(&op)) {
       diagnostics.report(&op, [&](Operation* illegalOp) {
-        illegalOp->emitOpError() << kPhaseMarker << " graph Spatial compute op remained after merge materialization";
+        illegalOp->emitOpError() << kPhaseMarker << " real channel communication is not allowed in scheduled phase 1";
       });
     }
   }

src/PIM/Conversion/ONNXToSpatial/Patterns.cpp

@@ -19,6 +19,7 @@ void populateConversionPatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
   populateSigmoidPatterns(patterns, ctx);
   populateSoftmaxPatterns(patterns, ctx);
   populateConcatPatterns(patterns, ctx);
+  populateFlattenPatterns(patterns, ctx);
   populateGatherPatterns(patterns, ctx);
   populateResizePatterns(patterns, ctx);
   populateReshapePatterns(patterns, ctx);

src/PIM/Conversion/ONNXToSpatial/Patterns.hpp

@@ -26,6 +26,7 @@ void populateReluPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext*
 void populateSigmoidPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateSoftmaxPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateConcatPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
+void populateFlattenPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateGatherPatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateResizePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);
 void populateReshapePatterns(mlir::RewritePatternSet& patterns, mlir::MLIRContext* ctx);

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

@@ -5,7 +5,7 @@
 
 #include "llvm/ADT/SmallVector.h"
 
-#include "src/Accelerators/PIM/Common/IR/ShapeUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
@@ -47,38 +47,28 @@ static FailureOr<Value> materializeBroadcastedConstantTensor(Value value,
     return failure();
 
   const int64_t rankOffset = static_cast<int64_t>(resultShape.size() - sourceShape.size());
-  for (int64_t i = 0; i < static_cast<int64_t>(resultShape.size()); ++i) {
-    const int64_t sourceIndex = i - rankOffset;
-    const int64_t sourceDim = sourceIndex < 0 ? 1 : sourceShape[sourceIndex];
-    const int64_t resultDim = resultShape[i];
-    if (sourceDim != 1 && sourceDim != resultDim)
-      return failure();
-  }
-
-  SmallVector<Attribute> sourceValues(denseAttr.getValues<Attribute>());
   SmallVector<int64_t> sourceStrides = computeRowMajorStrides(sourceShape);
   SmallVector<int64_t> resultStrides = computeRowMajorStrides(resultShape);
-
+  SmallVector<Attribute> sourceValues(denseAttr.getValues<Attribute>());
   SmallVector<Attribute> resultValues;
   resultValues.reserve(resultType.getNumElements());
 
   for (int64_t flatIndex = 0; flatIndex < resultType.getNumElements(); ++flatIndex) {
     int64_t remaining = flatIndex;
     int64_t sourceFlatIndex = 0;
-
     for (int64_t i = 0; i < static_cast<int64_t>(resultShape.size()); ++i) {
       const int64_t resultIndex = resultStrides.empty() ? 0 : remaining / resultStrides[i];
       remaining = resultStrides.empty() ? 0 : remaining % resultStrides[i];
-
       const int64_t sourceIndex = i - rankOffset;
       if (sourceIndex < 0)
         continue;
-
       const int64_t sourceDim = sourceShape[sourceIndex];
+      const int64_t resultDim = resultShape[i];
+      if (sourceDim != 1 && sourceDim != resultDim)
+        return failure();
       const int64_t mappedIndex = sourceDim == 1 ? 0 : resultIndex;
       sourceFlatIndex += mappedIndex * sourceStrides[sourceIndex];
     }
-
     resultValues.push_back(sourceValues[sourceFlatIndex]);
   }
 
@@ -106,7 +96,7 @@ static FailureOr<Value> materializeReciprocalTensor(Value value,
   if (failed(broadcastedValue))
     return failure();
 
-  auto denseAttr = dyn_cast<DenseFPElementsAttr>(getDenseConstantAttr(*broadcastedValue));
+  auto denseAttr = dyn_cast<DenseFPElementsAttr>(getHostConstDenseElementsAttr(*broadcastedValue));
   if (!denseAttr)
     return failure();
 
@@ -185,10 +175,45 @@ struct DivToSpatialCompute : OpConversionPattern<ONNXDivOp> {
   }
 };
 
+struct AddToSpatialCompute : OpConversionPattern<ONNXAddOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult
+  matchAndRewrite(ONNXAddOp op, ONNXAddOpAdaptor adaptor, ConversionPatternRewriter& rewriter) const override {
+    auto resultType = dyn_cast<RankedTensorType>(op.getResult().getType());
+    if (!resultType || !resultType.hasStaticShape())
+      return failure();
+
+    FailureOr<BiasAddPlanCandidate> candidate =
+      classifyBiasAddPlanCandidate(adaptor.getA(), adaptor.getB(), resultType);
+    if (succeeded(candidate)) {
+      auto plan = spatial::SpatBiasAddPlanOp::create(
+        rewriter, op.getLoc(), resultType, candidate->data, candidate->bias, rewriter.getStringAttr("nchw"));
+      rewriter.replaceOp(op, plan.getResult());
+      return success();
+    }
+
+    auto lhs = prepareElementwiseOperand(adaptor.getA(), resultType, rewriter, op.getLoc());
+    if (failed(lhs))
+      return failure();
+    auto rhs = prepareElementwiseOperand(adaptor.getB(), resultType, rewriter, op.getLoc());
+    if (failed(rhs))
+      return failure();
+
+    auto computeOp =
+      createSpatCompute<2>(rewriter, op.getLoc(), resultType, {}, ValueRange {*lhs, *rhs}, [&](Value x, Value y) {
+        auto loweredOp = spatial::SpatVAddOp::create(rewriter, op.getLoc(), resultType, x, y);
+        spatial::SpatYieldOp::create(rewriter, op.getLoc(), loweredOp.getResult());
+      });
+    rewriter.replaceOp(op, computeOp);
+    return success();
+  }
+};
+
 } // namespace
 
 void populateElementwisePatterns(RewritePatternSet& patterns, MLIRContext* ctx) {
-  patterns.add<BinaryElementwiseToSpatialCompute<ONNXAddOp, spatial::SpatVAddOp>>(ctx);
+  patterns.add<AddToSpatialCompute>(ctx);
   patterns.add<BinaryElementwiseToSpatialCompute<ONNXSubOp, spatial::SpatVSubOp>>(ctx);
   patterns.add<BinaryElementwiseToSpatialCompute<ONNXMulOp, spatial::SpatVMulOp>>(ctx);
   patterns.add<DivToSpatialCompute>(ctx);

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

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

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

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

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

@@ -139,9 +139,7 @@ static RankedTensorType getReducedSliceType(RankedTensorType inputType, ArrayRef
 }
 
 static RankedTensorType getLanePackedKeepdimsType(int64_t laneCount, RankedTensorType leafType) {
-  SmallVector<int64_t> shape(leafType.getShape().begin(), leafType.getShape().end());
-  shape.front() = laneCount;
-  return RankedTensorType::get(shape, leafType.getElementType(), leafType.getEncoding());
+  return spatial::getGraphBatchPhysicalResultType(laneCount, leafType);
 }
 
 static SmallVector<int64_t> getKeptAxes(ArrayRef<bool> reducedAxes) {
@@ -191,12 +189,9 @@ static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
 
   SmallVector<OpFoldResult> sliceOffsets;
   SmallVector<OpFoldResult> sliceSizes;
-  SmallVector<OpFoldResult> insertOffsets;
-  SmallVector<OpFoldResult> insertSizes(inputType.getRank(), rewriter.getIndexAttr(1));
   SmallVector<OpFoldResult> unitStrides = getUnitStrides(rewriter, inputType.getRank());
   sliceOffsets.reserve(inputType.getRank());
   sliceSizes.reserve(inputType.getRank());
-  insertOffsets.reserve(inputType.getRank());
 
   auto batchOp =
     createSpatComputeBatch(rewriter,
@@ -209,7 +204,6 @@ static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
                              size_t keptAxisIndex = 0;
                              sliceOffsets.clear();
                              sliceSizes.clear();
-                             insertOffsets.clear();
                              for (auto [axis, isReduced] : llvm::enumerate(reducedAxes)) {
                                if (isReduced) {
                                  sliceOffsets.push_back(rewriter.getIndexAttr(0));
@@ -224,14 +218,10 @@ static FailureOr<Value> buildReduceMeanKeepdimsBatch(Value input,
                                sliceSizes.push_back(rewriter.getIndexAttr(1));
                              }
 
-                             insertOffsets.push_back(args.lane);
-                             insertOffsets.append(inputType.getRank() - 1, rewriter.getIndexAttr(0));
-
                              Value slice = tensor::ExtractSliceOp::create(
                                rewriter, loc, sliceType, args.inputs.front(), sliceOffsets, sliceSizes, unitStrides);
                              Value reduced = spatial::SpatVAvgOp::create(rewriter, loc, leafType, slice).getResult();
-                             createParallelInsertSliceIntoBatchOutput(
-                               rewriter, loc, reduced, args.outputs.front(), insertOffsets, insertSizes, unitStrides);
+                             publishGraphBatchPhysicalFragment(rewriter, loc, reduced, args.outputs.front(), args.lane);
                            });
   if (failed(batchOp))
     return failure();
@@ -276,10 +266,11 @@ static Value buildKeepdimsFromLanePackedBatch(Value batchValue,
   if (!pendingLeadingReducedAxes.empty())
     expandCompactToKeepdims.back().append(pendingLeadingReducedAxes.begin(), pendingLeadingReducedAxes.end());
 
+  if (batchType.getNumElements() != batchType.getDimSize(0))
+    return {};
   auto reshapeCompute =
     createSpatCompute<1>(rewriter, loc, TypeRange {keepdimsType}, {}, ValueRange {batchValue}, [&](Value input) {
-      auto flatType =
-        RankedTensorType::get({batchType.getDimSize(0)}, batchType.getElementType(), batchType.getEncoding());
+      auto flatType = RankedTensorType::get({batchType.getNumElements()}, batchType.getElementType(), batchType.getEncoding());
       Value flat = tensor::CollapseShapeOp::create(rewriter, loc, flatType, input, collapseToFlat);
       Value compact = flat;
       if (compactKeptType != flatType)

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

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

src/PIM/Conversion/ONNXToSpatial/Patterns/Tensor/Flatten.cpp

@@ -0,0 +1,112 @@
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Transforms/DialectConversion.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/Common.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Patterns.hpp"
+#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Dialect/ONNX/ONNXOps.hpp"
+
+using namespace mlir;
+
+namespace onnx_mlir {
+namespace {
+
+static FailureOr<int64_t> normalizeFlattenAxis(int64_t axis, int64_t rank) {
+  int64_t normalizedAxis = axis < 0 ? rank + axis : axis;
+  if (normalizedAxis < 0 || normalizedAxis > rank)
+    return failure();
+  return normalizedAxis;
+}
+
+static int64_t product(ArrayRef<int64_t> values) {
+  int64_t result = 1;
+  for (int64_t value : values)
+    result *= value;
+  return result;
+}
+
+static SmallVector<ReassociationIndices> getCollapseTo1DReassociation(int64_t rank) {
+  SmallVector<ReassociationIndices> reassociation(1);
+  reassociation.front().reserve(rank);
+  for (int64_t dim = 0; dim < rank; ++dim)
+    reassociation.front().push_back(dim);
+  return reassociation;
+}
+
+static SmallVector<ReassociationIndices> getExpandFrom1DReassociation(int64_t rank) {
+  SmallVector<ReassociationIndices> reassociation(1);
+  reassociation.front().reserve(rank);
+  for (int64_t dim = 0; dim < rank; ++dim)
+    reassociation.front().push_back(dim);
+  return reassociation;
+}
+
+static Value buildFlatten(Value input,
+                          RankedTensorType sourceType,
+                          RankedTensorType resultType,
+                          int64_t axis,
+                          ConversionPatternRewriter& rewriter,
+                          Location loc) {
+  if (sourceType == resultType)
+    return input;
+
+  if (axis > 0 && axis < sourceType.getRank()) {
+    SmallVector<ReassociationIndices> reassociation(2);
+    for (int64_t dim = 0; dim < axis; ++dim)
+      reassociation[0].push_back(dim);
+    for (int64_t dim = axis; dim < sourceType.getRank(); ++dim)
+      reassociation[1].push_back(dim);
+    return tensor::CollapseShapeOp::create(rewriter, loc, resultType, input, reassociation);
+  }
+
+  Value flattened = input;
+  if (sourceType.getRank() != 1) {
+    auto flatType = RankedTensorType::get({sourceType.getNumElements()}, sourceType.getElementType());
+    flattened = tensor::CollapseShapeOp::create(
+      rewriter, loc, flatType, flattened, getCollapseTo1DReassociation(sourceType.getRank()));
+  }
+  return tensor::ExpandShapeOp::create(
+    rewriter, loc, resultType, flattened, getExpandFrom1DReassociation(resultType.getRank()));
+}
+
+struct Flatten : OpConversionPattern<ONNXFlattenOp> {
+  using OpConversionPattern::OpConversionPattern;
+
+  LogicalResult matchAndRewrite(ONNXFlattenOp flattenOp,
+                                ONNXFlattenOpAdaptor adaptor,
+                                ConversionPatternRewriter& rewriter) const override {
+    auto sourceType = dyn_cast<RankedTensorType>(adaptor.getInput().getType());
+    auto resultType = dyn_cast<RankedTensorType>(flattenOp.getOperation()->getResult(0).getType());
+    if (!sourceType || !resultType || !sourceType.hasStaticShape() || !resultType.hasStaticShape())
+      return failure();
+    if (!hasStaticPositiveShape(sourceType) || !hasStaticPositiveShape(resultType) || resultType.getRank() != 2)
+      return failure();
+
+    auto axis = normalizeFlattenAxis(flattenOp.getAxis(), sourceType.getRank());
+    if (failed(axis))
+      return failure();
+
+    int64_t outerDim = product(sourceType.getShape().take_front(*axis));
+    int64_t innerDim = product(sourceType.getShape().drop_front(*axis));
+    if (resultType.getShape()[0] != outerDim || resultType.getShape()[1] != innerDim)
+      return failure();
+
+    auto replaceWithFlatten = [&](auto build) -> LogicalResult {
+      Value flattened = materializeOrComputeUnary(adaptor.getInput(), resultType, rewriter, flattenOp.getLoc(), build);
+      rewriter.replaceOp(flattenOp, flattened);
+      return success();
+    };
+
+    return replaceWithFlatten([&](Value input) {
+      return buildFlatten(input, sourceType, resultType, *axis, rewriter, flattenOp.getLoc());
+    });
+  }
+};
+
+} // namespace
+
+void populateFlattenPatterns(RewritePatternSet& patterns, MLIRContext* ctx) { patterns.add<Flatten>(ctx); }
+
+} // namespace onnx_mlir

src/PIM/Conversion/ONNXToSpatial/SpatialLayoutPlanningPass.cpp

@@ -6,10 +6,11 @@
 
 #include "Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/BiasAddUtils.hpp"
+#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/RowStripLayoutUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/PlanLowering.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 #include "src/Accelerators/PIM/Pass/PIMPasses.h"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
 
 using namespace mlir;
 
@@ -19,7 +20,6 @@ namespace {
 static constexpr StringLiteral kLogicalLayout = "nchw";
 static constexpr StringLiteral kDenseLayout = "dense_nchw";
 static constexpr StringLiteral kRowStripLayout = "nchw_row_strip";
-static constexpr StringLiteral kRowStripIndexMap = "packed_hwc_rows_to_nchw";
 
 enum class SelectedLayout {
   DenseNchw,
@@ -34,6 +34,8 @@ static SelectedLayout getSelectedLayout(llvm::DenseMap<Value, SelectedLayout>& l
 static bool usesSelectedRowStrip(Operation* user, llvm::DenseMap<Value, SelectedLayout>& layouts) {
   if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(user))
     return getSelectedLayout(layouts, reluPlan.getResult()) == SelectedLayout::NchwRowStrip;
+  if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user))
+    return getSelectedLayout(layouts, biasAddPlan.getResult()) == SelectedLayout::NchwRowStrip;
   if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
     return getSelectedLayout(layouts, convPlan.getResult()) == SelectedLayout::NchwRowStrip;
   return false;
@@ -49,21 +51,26 @@ static bool allUsersCanHandleRowStrip(Value value, llvm::DenseMap<Value, Selecte
   return true;
 }
 
-static std::pair<SmallVector<int64_t>, SmallVector<int64_t>> buildRowStripMetadata(RankedTensorType type) {
-  SmallVector<int64_t> offsets;
-  SmallVector<int64_t> sizes;
-  const int64_t channels = type.getDimSize(1);
-  const int64_t height = type.getDimSize(2);
-  const int64_t width = type.getDimSize(3);
-  offsets.reserve(height * 4);
-  sizes.reserve(height * 4);
-  for (int64_t row = 0; row < height; ++row) {
-    offsets.append({0, 0, row, 0});
-    sizes.append({1, channels, 1, width});
+static bool canConsumeRowStripAsUser(Operation* user) {
+  if (isa<spatial::SpatReluPlanOp>(user))
+    return true;
+  if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(user)) {
+    auto resultType = dyn_cast<RankedTensorType>(biasAddPlan.getOutput().getType());
+    return resultType && isSupportedBiasAddValue(biasAddPlan.getBias(), resultType);
   }
-  return {offsets, sizes};
+  if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(user))
+    return succeeded(canConsumeAndProduceRowStrip(convPlan));
+  return false;
 }
 
+static bool hasRowStripConsumer(Value value) {
+  for (Operation* user : value.getUsers())
+    if (canConsumeRowStripAsUser(user))
+      return true;
+  return false;
+}
+
+
 static bool canSelectConvRowStrip(spatial::SpatConv2DPlanOp convPlan,
                                   llvm::DenseMap<Value, SelectedLayout>& layouts) {
   SelectedLayout inputLayout = getSelectedLayout(layouts, convPlan.getInput());
@@ -76,6 +83,9 @@ static SelectedLayout chooseConvLayout(spatial::SpatConv2DPlanOp convPlan,
                                        llvm::DenseMap<Value, SelectedLayout>& layouts) {
   if (!canSelectConvRowStrip(convPlan, layouts))
     return SelectedLayout::DenseNchw;
+  if (getSelectedLayout(layouts, convPlan.getInput()) != SelectedLayout::NchwRowStrip
+      && !hasRowStripConsumer(convPlan.getResult()))
+    return SelectedLayout::DenseNchw;
   if (!allUsersCanHandleRowStrip(convPlan.getResult(), layouts))
     return SelectedLayout::DenseNchw;
   return SelectedLayout::NchwRowStrip;
@@ -85,11 +95,27 @@ static SelectedLayout chooseReluLayout(spatial::SpatReluPlanOp reluPlan,
                                        llvm::DenseMap<Value, SelectedLayout>& layouts) {
   if (getSelectedLayout(layouts, reluPlan.getInput()) != SelectedLayout::NchwRowStrip)
     return SelectedLayout::DenseNchw;
+  if (!hasRowStripConsumer(reluPlan.getResult()))
+    return SelectedLayout::DenseNchw;
   if (!allUsersCanHandleRowStrip(reluPlan.getResult(), layouts))
     return SelectedLayout::DenseNchw;
   return SelectedLayout::NchwRowStrip;
 }
 
+static SelectedLayout chooseBiasAddLayout(spatial::SpatBiasAddPlanOp biasAddPlan,
+                                          llvm::DenseMap<Value, SelectedLayout>& layouts) {
+  if (getSelectedLayout(layouts, biasAddPlan.getInput()) != SelectedLayout::NchwRowStrip)
+    return SelectedLayout::DenseNchw;
+  auto resultType = dyn_cast<RankedTensorType>(biasAddPlan.getOutput().getType());
+  if (!resultType || !isSupportedBiasAddValue(biasAddPlan.getBias(), resultType))
+    return SelectedLayout::DenseNchw;
+  if (!hasRowStripConsumer(biasAddPlan.getResult()))
+    return SelectedLayout::DenseNchw;
+  if (!allUsersCanHandleRowStrip(biasAddPlan.getResult(), layouts))
+    return SelectedLayout::DenseNchw;
+  return SelectedLayout::NchwRowStrip;
+}
+
 static spatial::SpatBlueprintOp insertRowStripBlueprint(IRRewriter& rewriter, Value value) {
   auto outputType = cast<RankedTensorType>(value.getType());
   auto [offsets, sizes] = buildRowStripMetadata(outputType);
@@ -108,6 +134,7 @@ static spatial::SpatBlueprintOp insertRowStripBlueprint(IRRewriter& rewriter, Va
                                               nullptr,
                                               nullptr,
                                               nullptr,
+                                              nullptr,
                                               nullptr);
 }
 
@@ -173,6 +200,14 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
           }
           continue;
         }
+        if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(&op)) {
+          SelectedLayout selected = chooseBiasAddLayout(biasAddPlan, layouts);
+          if (layouts[biasAddPlan.getResult()] != selected) {
+            layouts[biasAddPlan.getResult()] = selected;
+            changed = true;
+          }
+          continue;
+        }
       }
     }
 
@@ -180,6 +215,8 @@ struct SpatialLayoutPlanningPass final : PassWrapper<SpatialLayoutPlanningPass,
       Value producedValue;
       if (auto convPlan = dyn_cast<spatial::SpatConv2DPlanOp>(&op))
         producedValue = convPlan.getResult();
+      else if (auto biasAddPlan = dyn_cast<spatial::SpatBiasAddPlanOp>(&op))
+        producedValue = biasAddPlan.getResult();
       else if (auto reluPlan = dyn_cast<spatial::SpatReluPlanOp>(&op))
         producedValue = reluPlan.getResult();
       else

src/PIM/Conversion/SpatialToGraphviz/CMakeLists.txt

@@ -1,17 +0,0 @@
-add_onnx_mlir_rewriter(SpatialToGraphviz)
-
-add_pim_library(OMSpatialToGraphviz
-  SpatialToGraphviz.cpp
-
-  EXCLUDE_FROM_OM_LIBS
-
-  LINK_LIBS PUBLIC
-  MLIRTosaDialect
-  OMCompilerOptions
-  OMPimCommon
-  OMONNXOps
-  SpatialOps
-
-  ACCEL_INCLUDE_DIRS PRIVATE
-  ${PIM_GENERATED_INCLUDE_DIRS}
-)

src/PIM/Conversion/SpatialToGraphviz/SpatialToGraphviz.cpp

@@ -1,259 +0,0 @@
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/Dialect/Tosa/IR/TosaOps.h"
-#include "mlir/IR/Block.h"
-#include "mlir/IR/Diagnostics.h"
-#include "mlir/IR/Value.h"
-#include "mlir/Pass/Pass.h"
-#include "mlir/Support/LLVM.h"
-#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-
-#include "llvm/Support/Casting.h"
-#include "llvm/Support/Format.h"
-
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-#include "src/Accelerators/PIM/Pass/PIMPasses.h"
-#include "src/Dialect/ONNX/ONNXOps.hpp"
-
-#define FORMAT_OPERATION(op) 'x' << llvm::format_hex_no_prefix(reinterpret_cast<size_t>(op), 0)
-#define FORMAT_ARGUMENT(computeOpPointer, argumentNum) llvm::format("Arg_%p_%u", computeOpPointer, argumentNum)
-
-using namespace mlir;
-
-namespace onnx_mlir {
-
-namespace {
-
-struct SpatialToGraphvizPass : public PassWrapper<SpatialToGraphvizPass, OperationPass<ModuleOp>> {
-
-  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(SpatialToGraphvizPass)
-
-  StringRef getArgument() const override { return "convert-spatial-to-graphviz"; }
-
-  StringRef getDescription() const override { return "Lower ONNX ops to Spatial ops."; }
-
-  SpatialToGraphvizPass(raw_ostream& os = llvm::errs())
-  : os(os) {}
-  SpatialToGraphvizPass(const SpatialToGraphvizPass& pass)
-  : SpatialToGraphvizPass(pass.os) {}
-  void runOnOperation() final;
-
-private:
-  raw_ostream& os;
-
-  /**
-   * Draws the subgraph for a given spatial::SpatCompute, including:
-   * 1. Input nodes (block arguments)
-   * 2. Operations
-   * 3. Edges between yield (output) and its users
-   *
-   * @param op The spatial::SpatCompute to draw the subgraph for.
-   * @param computeNum The number of the compute operation.
-   */
-  void drawComputeOpSubgraph(spatial::SpatCompute op, size_t computeNum) {
-    os << "\tsubgraph cluster" << computeNum << " {\n\t\tlabel=\"Compute" << computeNum << "\";\n"
-       << "\t\tstyle=filled;\n"
-       << "\t\tcolor=lightblue;\n";
-
-    Block& block = op.getBody().front();
-
-    // Inputs
-    size_t inputNum = 0;
-    for (BlockArgument& input : block.getArguments()) {
-
-      auto fromOp = FORMAT_ARGUMENT(op.getOperation(), inputNum);
-
-      os << "\t\t" << fromOp << " [label=\"Arg" << inputNum << "\",shape=box];\n";
-      for (auto userOp : input.getUsers())
-        os << "\t\t" << fromOp << " -> " << FORMAT_OPERATION(userOp) << ";\n";
-      inputNum++;
-    }
-
-    // Iterate operations
-    for (auto& childOp : block.getOperations()) {
-      os << "\t\t" << FORMAT_OPERATION(&childOp) << " [label=\"" << childOp.getName() << "\"];\n";
-
-      drawEdgesFromOpToItsUsers(&childOp);
-    }
-
-    os << "\t}\n";
-
-    // Draw edges from the yield to the users of this computeOp
-    Operation* yieldOp = block.getTerminator();
-    if (!isa<spatial::SpatYieldOp>(yieldOp)) {
-      yieldOp->emitError("Terminator of block must be YieldOp ???");
-      signalPassFailure();
-      return;
-    }
-
-    for (auto computeOpResult : op->getResults()) {
-      for (auto& computeOpUse : computeOpResult.getUses()) {
-        auto toOp = FORMAT_ARGUMENT(computeOpUse.getOwner(), computeOpUse.getOperandNumber());
-        os << "\t" << FORMAT_OPERATION(yieldOp) << " -> " << toOp << ";\n";
-      }
-    }
-  }
-
-  /**
-   * @brief Draws the subgraph for a concatOp.
-   *
-   * This function draws a subgraph for a concatOp. The subgraph consists of a
-   * node for each input of the concatOp, as well as an output node. Edges are
-   * created from the output node to each user of the concatOp.
-   *
-   * @param concatOp The concatOp for which the subgraph is drawn.
-   * @param concatOpNum The number of the concatOp.
-   */
-  void drawConcatOpSubgraph(Operation* concatOp, size_t concatOpNum) {
-    os << "\tsubgraph clusterconcat" << concatOpNum << " {\n\t\tlabel=\"ConcatOp" << concatOpNum << "\";\n"
-       << "\t\tstyle=filled;\n"
-       << "\t\tcolor=orange;\n";
-
-    // Inputs
-    size_t inputNum = 0;
-    for (Value input : concatOp->getOperands()) {
-      auto fromOp = FORMAT_ARGUMENT(concatOp, inputNum);
-
-      os << "\t\t" << fromOp << " [label=\"Input" << inputNum << "\"];\n";
-      for (auto userOp : input.getUsers())
-        os << "\t\t" << fromOp << " -> " << FORMAT_OPERATION(userOp) << ";\n";
-      inputNum++;
-    }
-
-    // Output
-    os << "\t\t" << FORMAT_OPERATION(concatOp) << " [label=Out];\n";
-
-    os << "\t}\n";
-
-    // Edges from output to users
-
-    for (auto& computeOpUse : concatOp->getResult(0).getUses()) {
-      os << "\t" << FORMAT_OPERATION(concatOp) << " -> "
-         << FORMAT_ARGUMENT(computeOpUse.getOwner(), computeOpUse.getOperandNumber()) << ";\n";
-    }
-  }
-
-  /**
-   * Draws the ExtractSliceOp in the graph visualization.
-   *
-   * This function takes a tensor::ExtractSliceOp and adds the corresponding
-   * node and edges to the graph visualization. It creates a node with the
-   * label as the static offsets attribute of the sliceOp, and connects it to
-   * the compute operations that use the result of the sliceOp.
-   *
-   * @param sliceOp The tensor::ExtractSliceOp to be drawn in the graph
-   * visualization.
-   */
-  void drawExtractSliceOp(tensor::ExtractSliceOp sliceOp) {
-    auto nodeId = FORMAT_ARGUMENT(sliceOp.getOperation(), 0);
-    os << "\t" << nodeId << " [label=\"Slice: ";
-    sliceOp.getStaticOffsetsAttr().print(os);
-    os << "\",color=lawngreen];\n";
-
-    for (auto& computeOpUse : sliceOp.getResult().getUses()) {
-      os << "\t" << nodeId << " -> " << FORMAT_ARGUMENT(computeOpUse.getOwner(), computeOpUse.getOperandNumber())
-         << ";\n";
-    }
-  }
-
-  void drawBiasTileOp(tensor::ExtractSliceOp sliceOp) {
-    auto nodeId = FORMAT_ARGUMENT(sliceOp.getOperation(), 0);
-    os << "\t" << nodeId << " [label=\"Bias: ";
-    sliceOp.getStaticOffsetsAttr().print(os);
-    os << "\",color=lightpink];\n";
-
-    for (auto user : sliceOp.getResult().getUsers())
-      os << "\t" << nodeId << " -> " << FORMAT_OPERATION(user) << ";\n";
-  }
-
-  /**
-   * Draws edges from the given operation to its users.
-   *
-   * @param fromOp The operation from which the edges are drawn.
-   */
-  void drawEdgesFromOpToItsUsers(mlir::Operation* fromOp) {
-    for (auto result : fromOp->getResults())
-      for (auto userOp : result.getUsers())
-        os << "\t\t" << FORMAT_OPERATION(fromOp) << " -> " << FORMAT_OPERATION(userOp) << ";\n";
-  }
-
-  /**
-   * Draws input node and edges for the given `funcOp`.
-   *
-   * @param funcOp The `funcOp` for which to draw input nodes and edges.
-   */
-  void drawInputNodesAndEdges(func::FuncOp& funcOp) {
-    os << "\tinput [label=\"Module Input\",color=green];\n";
-
-    size_t funcOpArgNum = 0;
-    for (BlockArgument& arg : funcOp.getArguments()) {
-
-      for (auto& useOp : arg.getUses()) {
-        os << "\tinput -> " << FORMAT_ARGUMENT(useOp.getOwner(), useOp.getOperandNumber()) << "[label=" << funcOpArgNum
-           << "];\n";
-      }
-      funcOpArgNum++;
-    }
-  }
-};
-
-void SpatialToGraphvizPass::runOnOperation() {
-  ModuleOp module = getOperation();
-
-  auto entryFunc = getPimEntryFunc(module);
-  if (failed(entryFunc)) {
-    module.emitError("failed to locate the PIM entry function for Spatial graph visualization");
-    signalPassFailure();
-    return;
-  }
-  func::FuncOp func = *entryFunc;
-
-  os << "digraph G {\n"
-     << "\tnode [style=filled,color=white];\n";
-
-  size_t computeNum = 0;
-  size_t concatNum = 0;
-
-  // Iterate over the ComputeOps within FuncOp:
-  // 1. Print their subgraph
-  // 2. Print the edges from its inputs to its outputs
-  for (Operation& op : func.getOps()) {
-    if (auto computeOp = dyn_cast<spatial::SpatCompute>(op)) {
-      drawComputeOpSubgraph(computeOp, computeNum++);
-    }
-    else if (auto concatOp = dyn_cast<tensor::ConcatOp>(op)) {
-      drawConcatOpSubgraph(concatOp, concatNum++);
-    }
-    else if (auto extractSliceOp = dyn_cast<tensor::ExtractSliceOp>(op)) {
-      auto producerOp = extractSliceOp->getOperand(0).getDefiningOp();
-      if (producerOp) {
-        // Skip extractSliceOp if producer is constant weights (ONNXConstantOp)
-        if (llvm::isa<ONNXConstantOp>(producerOp))
-          continue;
-        // If produced by tosa::ReshapeOp (i.e. it is a bias tile) connect
-        // directly to its user, which is not a ComputeOp argument.
-        if (llvm::isa<tosa::ReshapeOp>(producerOp)) {
-          drawBiasTileOp(extractSliceOp);
-          continue;
-        }
-      }
-
-      drawExtractSliceOp(extractSliceOp);
-    }
-  }
-
-  // Draw input node, and edges to it users
-  drawInputNodesAndEdges(func);
-
-  // Draw output node (use the return Operation - argument number=0 - as nodeId)
-  auto returnOp = func.getBody().front().getTerminator();
-  os << '\t' << FORMAT_ARGUMENT(returnOp, 0) << " [label=\"Module Output\",color=green];\n";
-
-  os << "}\n";
-}
-
-} // namespace
-
-std::unique_ptr<Pass> createSpatialToGraphvizPass() { return std::make_unique<SpatialToGraphvizPass>(); }
-
-} // namespace onnx_mlir

src/PIM/Conversion/SpatialToPim/BatchCoreLoweringPatterns.cpp

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

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

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

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

@@ -302,76 +302,87 @@ void PimBufferizationPass::annotateWeightsMemrefs(ModuleOp moduleOp, func::FuncO
 
 LogicalResult PimBufferizationPass::verifyContiguousRuntimeOperands(ModuleOp moduleOp) const {
   bool hasFailure = false;
-  moduleOp.walk([&](Operation* op) {
-    auto verifyOperand = [&](Value operand, unsigned operandIndex) {
-      if (!isa<BaseMemRefType>(operand.getType()))
-        return;
-      if (succeeded(resolveContiguousAddress(operand)) || succeeded(compileContiguousAddressExpr(operand)))
-        return;
-      op->emitOpError() << "operand #" << operandIndex
-                        << " is not backed by contiguous addressable storage after PIM bufferization";
-      hasFailure = true;
-    };
 
-    if (auto memCopyOp = dyn_cast<PimMemCopyOp>(op)) {
-      if (!pim::isNormalizedCopyOp(memCopyOp)) {
-        memCopyOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
-        hasFailure = true;
-      }
-      verifyOperand(memCopyOp.getTarget(), 0);
-      verifyOperand(memCopyOp.getSource(), 1);
-      return;
-    }
-    if (auto loadOp = dyn_cast<PimMemCopyHostToDevOp>(op)) {
-      if (!pim::isNormalizedCopyOp(loadOp)) {
-        loadOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
-        hasFailure = true;
-      }
-      verifyOperand(loadOp.getDeviceTarget(), 2);
-      verifyOperand(loadOp.getHostSource(), 3);
-      return;
-    }
-    if (auto storeOp = dyn_cast<PimMemCopyDevToHostOp>(op)) {
-      if (!pim::isNormalizedCopyOp(storeOp)) {
-        storeOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
-        hasFailure = true;
-      }
-      verifyOperand(storeOp.getHostTarget(), 2);
-      verifyOperand(storeOp.getDeviceSource(), 3);
-      return;
-    }
-    if (auto sendOp = dyn_cast<PimSendOp>(op)) {
-      verifyOperand(sendOp.getInput(), 0);
-      return;
-    }
-    if (auto receiveOp = dyn_cast<PimReceiveOp>(op)) {
-      verifyOperand(receiveOp.getOutputBuffer(), 0);
-      return;
-    }
-    if (auto concatOp = dyn_cast<PimConcatOp>(op)) {
-      verifyOperand(concatOp.getOutputBuffer(), 0);
-      for (auto inputAndIndex : llvm::enumerate(concatOp.getInputs()))
-        verifyOperand(inputAndIndex.value(), inputAndIndex.index() + 1);
-      return;
-    }
-    if (isa<PimTransposeOp,
-            PimVMMOp,
-            PimVVAddOp,
-            PimVVSubOp,
-            PimVVMulOp,
-            PimVVMaxOp,
-            PimVVDMulOp,
-            PimVAvgOp,
-            PimVReluOp,
-            PimVTanhOp,
-            PimVSigmOp,
-            PimVSoftmaxOp>(op)) {
-      for (auto operandAndIndex : llvm::enumerate(op->getOperands())) {
-        if (auto vmmOp = dyn_cast<PimVMMOp>(op); vmmOp && operandAndIndex.index() == 0)
-          continue;
-        verifyOperand(operandAndIndex.value(), operandAndIndex.index());
-      }
-    }
+  auto verifyWithKnowledge = [&](auto coreLikeOp, const StaticValueKnowledge& initialKnowledge) {
+    (void) walkPimCoreBlockStructurally(
+      coreLikeOp.getBody().front(), initialKnowledge, [&](Operation& op, const StaticValueKnowledge& knowledge) {
+        auto verifyOperand = [&](Value operand, unsigned operandIndex) {
+          if (!isa<BaseMemRefType>(operand.getType()))
+            return;
+          if (succeeded(resolveContiguousAddress(operand, knowledge)) || succeeded(compileContiguousAddressExpr(operand)))
+            return;
+          op.emitOpError() << "operand #" << operandIndex
+                           << " is not backed by contiguous addressable storage after PIM bufferization";
+          hasFailure = true;
+        };
+
+        if (auto memCopyOp = dyn_cast<PimMemCopyOp>(&op)) {
+          if (!pim::isNormalizedCopyOp(memCopyOp)) {
+            memCopyOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
+            hasFailure = true;
+          }
+          verifyOperand(memCopyOp.getTarget(), 0);
+          verifyOperand(memCopyOp.getSource(), 1);
+          return success();
+        }
+        if (auto loadOp = dyn_cast<PimMemCopyHostToDevOp>(&op)) {
+          if (!pim::isNormalizedCopyOp(loadOp)) {
+            loadOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
+            hasFailure = true;
+          }
+          verifyOperand(loadOp.getDeviceTarget(), 2);
+          verifyOperand(loadOp.getHostSource(), 3);
+          return success();
+        }
+        if (auto storeOp = dyn_cast<PimMemCopyDevToHostOp>(&op)) {
+          if (!pim::isNormalizedCopyOp(storeOp)) {
+            storeOp.emitOpError("must use base memref operands plus explicit byte offsets after bufferization");
+            hasFailure = true;
+          }
+          verifyOperand(storeOp.getHostTarget(), 2);
+          verifyOperand(storeOp.getDeviceSource(), 3);
+          return success();
+        }
+        if (auto sendOp = dyn_cast<PimSendOp>(&op)) {
+          verifyOperand(sendOp.getInput(), 0);
+          return success();
+        }
+        if (auto receiveOp = dyn_cast<PimReceiveOp>(&op)) {
+          verifyOperand(receiveOp.getOutputBuffer(), 0);
+          return success();
+        }
+        if (auto concatOp = dyn_cast<PimConcatOp>(&op)) {
+          verifyOperand(concatOp.getOutputBuffer(), 0);
+          for (auto inputAndIndex : llvm::enumerate(concatOp.getInputs()))
+            verifyOperand(inputAndIndex.value(), inputAndIndex.index() + 1);
+          return success();
+        }
+        if (isa<PimTransposeOp,
+                PimVMMOp,
+                PimVVAddOp,
+                PimVVSubOp,
+                PimVVMulOp,
+                PimVVMaxOp,
+                PimVVDMulOp,
+                PimVAvgOp,
+                PimVReluOp,
+                PimVTanhOp,
+                PimVSigmOp,
+                PimVSoftmaxOp>(&op)) {
+          for (auto operandAndIndex : llvm::enumerate(op.getOperands())) {
+            if (auto vmmOp = dyn_cast<PimVMMOp>(&op); vmmOp && operandAndIndex.index() == 0)
+              continue;
+            verifyOperand(operandAndIndex.value(), operandAndIndex.index());
+          }
+        }
+        return success();
+      });
+  };
+
+  moduleOp.walk([&](pim::PimCoreOp coreOp) { verifyWithKnowledge(coreOp, seedCoreKnowledge(coreOp)); });
+  moduleOp.walk([&](pim::PimCoreBatchOp coreBatchOp) {
+    StaticValueKnowledge knowledge = seedCoreBatchKnowledge(coreBatchOp, 0);
+    verifyWithKnowledge(coreBatchOp, knowledge);
   });
 
   if (hasFailure) {

src/PIM/Dialect/Spatial/CMakeLists.txt

@@ -7,12 +7,16 @@ add_pim_library(SpatialOps
   SpatialOpsVerify.cpp
   SpatialOpsCanonicalization.cpp
   ${PIM_SRC_ROOT}/Conversion/ONNXToSpatial/CompileTime.cpp
-  Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp
-  Transforms/MergeComputeNodes/HostOutputFinalization.cpp
-  Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp
-  Transforms/MergeComputeNodes/ProjectedFragments.cpp
   Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp
   Transforms/MergeComputeNodes/Scheduling/ComputeInstanceUtils.cpp
+  Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp
+  Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp
+  Transforms/MergeComputeNodes/DeferredCommunicationDeadlock.cpp
+  Transforms/MergeComputeNodes/DeferredCommunicationRealization.cpp
+  Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp
+  Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp
+  Transforms/MergeComputeNodes/ScheduledComputeReport.cpp
+  Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp
   Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp
   Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp
 

src/PIM/Dialect/Spatial/Spatial.td

@@ -6,6 +6,7 @@ include "mlir/IR/OpAsmInterface.td"
 include "mlir/IR/BuiltinTypes.td"
 include "mlir/IR/AttrTypeBase.td"
 include "mlir/IR/RegionKindInterface.td"
+include "mlir/Interfaces/ControlFlowInterfaces.td"
 include "mlir/Interfaces/ParallelCombiningOpInterface.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 
@@ -27,7 +28,7 @@ def SpatTensor :
 //===----------------------------------------------------------------------===//
 
 class SpatComputeLikeBase<string mnemonic> : SpatOp<mnemonic,
-    [SingleBlock, AttrSizedOperandSegments,
+    [AttrSizedOperandSegments,
      DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmBlockArgumentNames"]>]> {
   let summary = "Compute region with attached constant weights";
 
@@ -40,7 +41,7 @@ class SpatComputeLikeBase<string mnemonic> : SpatOp<mnemonic,
     Variadic<SpatTensor>:$outputs
   );
 
-  let regions = (region SizedRegion<1>:$body);
+  let regions = (region MinSizedRegion<1>:$body);
 
   let hasVerifier = 1;
   let hasFolder = 1;
@@ -76,7 +77,7 @@ def SpatScheduledCompute : SpatComputeLikeBase<"scheduled_compute"> {
 }
 
 class SpatComputeBatchLikeBase<string mnemonic> : SpatOp<mnemonic,
-    [SingleBlock, AttrSizedOperandSegments,
+    [AttrSizedOperandSegments,
      DeclareOpInterfaceMethods<OpAsmOpInterface, ["getAsmBlockArgumentNames"]>]> {
   let summary = "Tensor-native batch of equivalent compute lanes with shared weights and packed inputs";
 
@@ -90,13 +91,14 @@ class SpatComputeBatchLikeBase<string mnemonic> : SpatOp<mnemonic,
     Variadic<SpatTensor>:$outputs
   );
 
-  let regions = (region SizedRegion<1>:$body);
+  let regions = (region MinSizedRegion<1>:$body);
 
   let hasVerifier = 1;
   let hasCustomAssemblyFormat = 1;
 }
 
 def SpatGraphComputeBatch : SpatComputeBatchLikeBase<"graph_compute_batch"> {
+  let hasCanonicalizer = 1;
   let extraClassDeclaration = [{
     std::optional<::mlir::BlockArgument> getLaneArgument();
     std::optional<::mlir::BlockArgument> getWeightArgument(unsigned idx);
@@ -113,6 +115,7 @@ def SpatGraphComputeBatch : SpatComputeBatchLikeBase<"graph_compute_batch"> {
 }
 
 def SpatScheduledComputeBatch : SpatComputeBatchLikeBase<"scheduled_compute_batch"> {
+  let hasCanonicalizer = 1;
   let extraClassDeclaration = [{
     std::optional<::mlir::BlockArgument> getLaneArgument();
     std::optional<::mlir::BlockArgument> getWeightArgument(unsigned idx);
@@ -161,6 +164,41 @@ def SpatYieldOp : SpatOp<"yield", [Terminator]> {
   let hasCustomAssemblyFormat = 1;
 }
 
+def SpatBlockYieldOp : SpatOp<"block_yield", [
+    Terminator,
+    DeclareOpInterfaceMethods<BranchOpInterface, ["getSuccessorForOperands"]>
+  ]> {
+  let summary = "Terminate a scheduled structural compute block";
+
+  let arguments = (ins
+    Variadic<AnyType>:$outputs
+  );
+
+  let successors = (successor
+    VariadicSuccessor<AnySuccessor>:$next
+  );
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def SpatDeferredCommunicationOp : SpatOp<"deferred_communication", [SingleBlock]> {
+  let summary = "Temporary scheduled payload derivation placeholder";
+
+  let arguments = (ins
+    Variadic<SpatTensor>:$sources
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let regions = (region SizedRegion<1>:$body);
+
+  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";
 
@@ -232,6 +270,22 @@ def SpatReluPlanOp : SpatOp<"relu_plan", []> {
   let hasVerifier = 1;
 }
 
+def SpatBiasAddPlanOp : SpatOp<"bias_add_plan", []> {
+  let summary = "Layout-aware Conv-style bias add planning op";
+
+  let arguments = (ins
+    SpatTensor:$input,
+    SpatTensor:$bias,
+    StrAttr:$logicalLayout
+  );
+
+  let results = (outs
+    SpatTensor:$output
+  );
+
+  let hasVerifier = 1;
+}
+
 def SpatBlueprintOp : SpatOp<"blueprint", []> {
   let summary = "Blueprint for assembling logical tensors from published fragments";
 
@@ -245,6 +299,7 @@ def SpatBlueprintOp : SpatOp<"blueprint", []> {
     StrAttr:$indexMap,
     OptionalAttr<StrAttr>:$mode,
     OptionalAttr<DenseI64ArrayAttr>:$fragmentOperandIndices,
+    OptionalAttr<DenseI64ArrayAttr>:$fragmentSourceSlots,
     OptionalAttr<DenseI64ArrayAttr>:$fragmentSourceOffsets,
     OptionalAttr<DenseI64ArrayAttr>:$fragmentStrides,
     OptionalAttr<StrAttr>:$conflictPolicy,

src/PIM/Dialect/Spatial/SpatialOps.cpp

@@ -10,6 +10,18 @@ using namespace mlir;
 
 namespace onnx_mlir {
 namespace spatial {
+
+RankedTensorType getGraphBatchPhysicalResultType(int64_t laneCount, RankedTensorType fragmentType) {
+  SmallVector<int64_t> shape {laneCount};
+  llvm::append_range(shape, fragmentType.getShape());
+  return RankedTensorType::get(shape, fragmentType.getElementType(), fragmentType.getEncoding());
+}
+
+FailureOr<RankedTensorType> getGraphBatchFragmentType(RankedTensorType physicalType, int64_t expectedLaneCount) {
+  if (!physicalType || physicalType.getRank() < 1 || physicalType.getDimSize(0) != expectedLaneCount)
+    return failure();
+  return RankedTensorType::get(physicalType.getShape().drop_front(), physicalType.getElementType(), physicalType.getEncoding());
+}
 namespace {
 
 std::optional<BlockArgument> getBlockArgument(Region& body, unsigned argIdx) {
@@ -238,6 +250,15 @@ void SpatScheduledCompute::getAsmBlockArgumentNames(Region& region, OpAsmSetValu
   setComputeAsmBlockArgumentNames(*this, region, setNameFn);
 }
 
+SuccessorOperands SpatBlockYieldOp::getSuccessorOperands(unsigned index) {
+  assert(index == 0 && "invalid successor index");
+  return SuccessorOperands(getOutputsMutable());
+}
+
+Block* SpatBlockYieldOp::getSuccessorForOperands(ArrayRef<Attribute>) {
+  return getOperation()->getNumSuccessors() == 0 ? nullptr : getOperation()->getSuccessor(0);
+}
+
 std::optional<BlockArgument> SpatGraphComputeBatch::getLaneArgument() { return getBlockArgument(getBody(), 0); }
 std::optional<BlockArgument> SpatGraphComputeBatch::getWeightArgument(unsigned idx) {
   return getBlockArgument(getBody(), 1 + idx);

src/PIM/Dialect/Spatial/SpatialOps.hpp

@@ -30,6 +30,10 @@
 namespace onnx_mlir {
 namespace spatial {
 
+mlir::RankedTensorType getGraphBatchPhysicalResultType(int64_t laneCount, mlir::RankedTensorType fragmentType);
+mlir::FailureOr<mlir::RankedTensorType>
+getGraphBatchFragmentType(mlir::RankedTensorType physicalType, int64_t expectedLaneCount);
+
 bool isGraphComputeLike(mlir::Operation* op);
 bool isGraphBatchComputeLike(mlir::Operation* op);
 bool isScheduledComputeLike(mlir::Operation* op);

src/PIM/Dialect/Spatial/SpatialOpsAsm.cpp

@@ -160,7 +160,7 @@ void printComputeLikeOp(ComputeOpTy op, OpAsmPrinter& printer) {
   printer << " -> ";
   printCompressedTypeSequence(printer, op.getResultTypes());
   printer << " ";
-  printer.printRegion(op.getBody(), /*printEntryBlockArgs=*/false);
+  printer.printRegion(op.getBody(), /*printEntryBlockArgs=*/!op.getBody().hasOneBlock());
 }
 
 template <typename ComputeOpTy>
@@ -290,7 +290,7 @@ void printComputeBatchLikeOp(ComputeBatchOpTy op, OpAsmPrinter& printer) {
   printer << " -> ";
   printCompressedTypeSequence(printer, op.getResultTypes());
   printer << " ";
-  printer.printRegion(op.getBody(), /*printEntryBlockArgs=*/false);
+  printer.printRegion(op.getBody(), /*printEntryBlockArgs=*/!op.getBody().hasOneBlock());
 }
 
 template <typename ComputeBatchOpTy>
@@ -407,6 +407,89 @@ ParseResult SpatYieldOp::parse(OpAsmParser& parser, OperationState& result) {
   return parser.resolveOperands(outputs, outputTypes, parser.getCurrentLocation(), result.operands);
 }
 
+void SpatBlockYieldOp::print(OpAsmPrinter& printer) {
+  printer << " ";
+  printCompressedValueSequence(printer, getOutputs());
+  if (getOperation()->getNumSuccessors() != 0) {
+    printer << " next ";
+    printer.printSuccessor(getOperation()->getSuccessor(0));
+  }
+  printer.printOptionalAttrDict((*this)->getAttrs());
+  printer << " : ";
+  printCompressedTypeSequence(printer, getOutputs().getTypes());
+}
+
+ParseResult SpatBlockYieldOp::parse(OpAsmParser& parser, OperationState& result) {
+  SmallVector<OpAsmParser::UnresolvedOperand> outputs;
+  SmallVector<Type> outputTypes;
+  Block* successor = nullptr;
+
+  OpAsmParser::UnresolvedOperand firstOutput;
+  OptionalParseResult firstOutputResult = parser.parseOptionalOperand(firstOutput);
+  if (firstOutputResult.has_value()) {
+    if (failed(*firstOutputResult))
+      return failure();
+    if (parseCompressedOperandEntryWithFirst(parser, firstOutput, outputs))
+      return failure();
+    while (succeeded(parser.parseOptionalComma()))
+      if (parseOneCompressedOperandEntry(parser, outputs))
+        return failure();
+  }
+
+  if (succeeded(parser.parseOptionalKeyword("next")) && parser.parseSuccessor(successor))
+    return failure();
+
+  if (parser.parseOptionalAttrDict(result.attributes) || parser.parseColon()
+      || parseCompressedTypeSequence(parser, outputTypes, /*allowEmpty=*/true))
+    return failure();
+
+  if (outputs.size() != outputTypes.size())
+    return parser.emitError(parser.getCurrentLocation(), "number of outputs and output types must match");
+  if (parser.resolveOperands(outputs, outputTypes, parser.getCurrentLocation(), result.operands))
+    return failure();
+  if (successor)
+    result.addSuccessors(successor);
+  return success();
+}
+
+void SpatDeferredCommunicationOp::print(OpAsmPrinter& printer) {
+  printer << " ";
+  printCompressedValueSequence(printer, getSources());
+  printer.printOptionalAttrDict((*this)->getAttrs());
+  printer << " : ";
+  printer.printFunctionalType(getSources().getTypes(), getOperation()->getResultTypes());
+  printer << " ";
+  printer.printRegion(getBody(), /*printEntryBlockArgs=*/false);
+}
+
+ParseResult SpatDeferredCommunicationOp::parse(OpAsmParser& parser, OperationState& result) {
+  SmallVector<OpAsmParser::UnresolvedOperand> sources;
+  Type functionTypeStorage;
+
+  if (parseCompressedOperandSequence(parser, sources) || parser.parseOptionalAttrDict(result.attributes)
+      || parser.parseColon() || parser.parseType(functionTypeStorage))
+    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())
+    return parser.emitError(parser.getCurrentLocation(), "number of sources and source types must match");
+
+  if (parser.resolveOperands(sources, functionType.getInputs(), parser.getCurrentLocation(), result.operands))
+    return failure();
+  result.addTypes(functionType.getResults());
+
+  Region* body = result.addRegion();
+  SmallVector<OpAsmParser::Argument> bodyArgs;
+  for (Type type : functionType.getInputs()) {
+    OpAsmParser::Argument argument;
+    argument.type = type;
+    bodyArgs.push_back(argument);
+  }
+  return parser.parseRegion(*body, bodyArgs);
+}
+
 void SpatExtractRowsOp::print(OpAsmPrinter& printer) {
   printer << " ";
   printer.printOperand(getInput());
@@ -493,6 +576,10 @@ void SpatBlueprintOp::print(OpAsmPrinter& printer) {
     printer << " operandIndices ";
     printCompressedIntegerList(printer, *operandIndices);
   }
+  if (std::optional<ArrayRef<int64_t>> sourceSlots = getFragmentSourceSlots()) {
+    printer << " sourceSlots ";
+    printCompressedIntegerList(printer, *sourceSlots);
+  }
   if (std::optional<ArrayRef<int64_t>> sourceOffsets = getFragmentSourceOffsets()) {
     printer << " sourceOffsets ";
     printCompressedIntegerList(printer, *sourceOffsets);
@@ -514,6 +601,7 @@ void SpatBlueprintOp::print(OpAsmPrinter& printer) {
                                  getIndexMapAttrName().getValue(),
                                  getModeAttrName().getValue(),
                                  getFragmentOperandIndicesAttrName().getValue(),
+                                 getFragmentSourceSlotsAttrName().getValue(),
                                  getFragmentSourceOffsetsAttrName().getValue(),
                                  getFragmentStridesAttrName().getValue(),
                                  getConflictPolicyAttrName().getValue(),
@@ -537,6 +625,7 @@ ParseResult SpatBlueprintOp::parse(OpAsmParser& parser, OperationState& result)
   SmallVector<int64_t> fragmentOffsets;
   SmallVector<int64_t> fragmentSizes;
   SmallVector<int64_t> fragmentOperandIndices;
+  SmallVector<int64_t> fragmentSourceSlots;
   SmallVector<int64_t> fragmentSourceOffsets;
   SmallVector<int64_t> fragmentStrides;
 
@@ -554,6 +643,9 @@ ParseResult SpatBlueprintOp::parse(OpAsmParser& parser, OperationState& result)
   if (succeeded(parser.parseOptionalKeyword("operandIndices"))
       && parseCompressedIntegerList(parser, fragmentOperandIndices))
     return failure();
+  if (succeeded(parser.parseOptionalKeyword("sourceSlots"))
+      && parseCompressedIntegerList(parser, fragmentSourceSlots))
+    return failure();
   if (succeeded(parser.parseOptionalKeyword("sourceOffsets"))
       && parseCompressedIntegerList(parser, fragmentSourceOffsets))
     return failure();
@@ -584,6 +676,8 @@ ParseResult SpatBlueprintOp::parse(OpAsmParser& parser, OperationState& result)
     result.addAttribute("mode", mode);
   if (!fragmentOperandIndices.empty())
     result.addAttribute("fragmentOperandIndices", builder.getDenseI64ArrayAttr(fragmentOperandIndices));
+  if (!fragmentSourceSlots.empty())
+    result.addAttribute("fragmentSourceSlots", builder.getDenseI64ArrayAttr(fragmentSourceSlots));
   if (!fragmentSourceOffsets.empty())
     result.addAttribute("fragmentSourceOffsets", builder.getDenseI64ArrayAttr(fragmentSourceOffsets));
   if (!fragmentStrides.empty())

src/PIM/Dialect/Spatial/SpatialOpsCanonicalization.cpp

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

src/PIM/Dialect/Spatial/SpatialOpsVerify.cpp

@@ -1,5 +1,6 @@
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Block.h"
@@ -59,6 +60,21 @@ static LogicalResult verifyStaticWeights(ComputeOpTy computeOp, StringRef kind)
   return success();
 }
 
+static bool isStaticScfForInductionVar(Value value) {
+  auto blockArg = dyn_cast<BlockArgument>(value);
+  if (!blockArg)
+    return false;
+
+  auto loop = dyn_cast_or_null<scf::ForOp>(blockArg.getOwner()->getParentOp());
+  if (!loop || loop.getInductionVar() != value)
+    return false;
+
+  std::optional<int64_t> lowerBound = matchConstantIndexValue(loop.getLowerBound());
+  std::optional<int64_t> upperBound = matchConstantIndexValue(loop.getUpperBound());
+  std::optional<int64_t> step = matchConstantIndexValue(loop.getStep());
+  return lowerBound && upperBound && step && *step > 0 && *upperBound >= *lowerBound;
+}
+
 static bool isStaticIndexExpr(Value value) {
   if (matchConstantIndexValue(value))
     return true;
@@ -80,7 +96,7 @@ static bool isStaticIndexExpr(Value value) {
 }
 
 static bool isSupportedLaneOffsetExpr(Value value, BlockArgument laneArg) {
-  if (value == laneArg || isStaticIndexExpr(value))
+  if (value == laneArg || isStaticIndexExpr(value) || isStaticScfForInductionVar(value))
     return true;
 
   auto affineApply = value.getDefiningOp<affine::AffineApplyOp>();
@@ -176,12 +192,18 @@ static bool isConstantExternalValue(Value value) {
   return definingOp && definingOp->hasTrait<OpTrait::ConstantLike>();
 }
 
+static bool isRecordedDeferredCommunicationSource(Operation* op, Value value) {
+  auto transfer = dyn_cast<SpatDeferredCommunicationOp>(op);
+  return transfer && llvm::is_contained(transfer.getSources(), value);
+}
+
 static LogicalResult verifyOnlyConstantExternalValues(Operation* ownerOp, Region& region, StringRef kind) {
   bool hasFailure = false;
   region.walk([&](Operation* op) {
     for (OpOperand& operand : op->getOpOperands()) {
       Value value = operand.get();
-      if (isDefinedInsideRegion(value, region) || isConstantExternalValue(value))
+      if (isDefinedInsideRegion(value, region) || isConstantExternalValue(value)
+          || isRecordedDeferredCommunicationSource(op, value))
         continue;
 
       InFlightDiagnostic diagnostic =
@@ -203,8 +225,35 @@ 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) {
+static LogicalResult verifyBatchBody(ComputeBatchOpTy batchOp, Block& block, bool verifyLaneSliceOffsets = true) {
   if (batchOp.getNumResults() == 0) {
     auto yieldOp = dyn_cast_or_null<SpatYieldOp>(block.getTerminator());
     if (!yieldOp)
@@ -219,11 +268,12 @@ static LogicalResult verifyBatchBody(ComputeBatchOpTy batchOp, Block& block) {
   auto laneArg = batchOp.getLaneArgument();
   if (!laneArg)
     return batchOp.emitError("compute_batch body must have a lane block argument");
-  for (auto& bodyOp : block) {
-    if (auto extractSlice = dyn_cast<tensor::ExtractSliceOp>(&bodyOp))
-      if (failed(verifyStaticUnitStrideExtractSliceOp(extractSlice, *laneArg, "tensor.extract_slice")))
-        return failure();
-  }
+  if (verifyLaneSliceOffsets)
+    for (auto& bodyOp : block) {
+      if (auto extractSlice = dyn_cast<tensor::ExtractSliceOp>(&bodyOp))
+        if (failed(verifyStaticUnitStrideExtractSliceOp(extractSlice, *laneArg, "tensor.extract_slice")))
+          return failure();
+    }
   return success();
 }
 
@@ -436,6 +486,39 @@ LogicalResult SpatReluPlanOp::verify() {
   return success();
 }
 
+LogicalResult SpatBiasAddPlanOp::verify() {
+  if (failed(verifyPlanTensorTypes(getOperation(), getInput(), getOutput(), "spat.bias_add_plan")))
+    return failure();
+  if (!isKnownLogicalLayout(getLogicalLayout()))
+    return emitError("requires a known logical layout");
+
+  auto inputType = dyn_cast<RankedTensorType>(getInput().getType());
+  auto biasType = dyn_cast<RankedTensorType>(getBias().getType());
+  auto outputType = dyn_cast<RankedTensorType>(getOutput().getType());
+  if (!inputType || !biasType || !outputType)
+    return emitError("requires ranked tensor input, bias, and output");
+  if (!inputType.hasStaticShape() || !biasType.hasStaticShape() || !outputType.hasStaticShape())
+    return emitError("requires static tensor input, bias, and output");
+  if (inputType != outputType)
+    return emitError("requires matching input and output tensor types");
+  if (outputType.getRank() != 4)
+    return emitError("requires rank-4 input/output tensors");
+  if (getLogicalLayout() != "nchw")
+    return emitError("requires logical layout \"nchw\"");
+  if (biasType.getElementType() != outputType.getElementType())
+    return emitError("requires bias element type to match the output element type");
+
+  ArrayRef<int64_t> biasShape = biasType.getShape();
+  const int64_t channels = outputType.getDimSize(1);
+  const bool supported = biasShape.empty() || (biasShape.size() == 1 && biasShape[0] == channels)
+                         || (biasShape.size() == 2 && biasShape[0] == 1 && biasShape[1] == channels)
+                         || (biasShape.size() == 4 && biasShape[0] == 1 && biasShape[1] == channels
+                             && biasShape[2] == 1 && biasShape[3] == 1);
+  if (!supported)
+    return emitError("requires scalar or per-channel bias broadcastable over NCHW");
+  return success();
+}
+
 LogicalResult SpatBlueprintOp::verify() {
   auto modeAttr = getModeAttr();
   bool isFragmentAssembly = modeAttr && modeAttr.getValue() == "fragment_assembly";
@@ -491,20 +574,26 @@ LogicalResult SpatBlueprintOp::verify() {
 
   auto stridesAttr = getFragmentStridesAttr();
   auto operandIndicesAttr = getFragmentOperandIndicesAttr();
+  auto sourceSlotsAttr = getFragmentSourceSlotsAttr();
   auto sourceOffsetsAttr = getFragmentSourceOffsetsAttr();
   if (!operandIndicesAttr)
     return emitError("fragment assembly blueprint requires fragment operand indices");
+  if (!sourceSlotsAttr)
+    return emitError("fragment assembly blueprint requires physical fragment source slots");
   if (!sourceOffsetsAttr)
     return emitError("fragment assembly blueprint requires fragment source offsets");
   if (!stridesAttr)
     return emitError("fragment assembly blueprint requires fragment strides");
   ArrayRef<int64_t> operandIndices = operandIndicesAttr.asArrayRef();
+  ArrayRef<int64_t> sourceSlots = sourceSlotsAttr.asArrayRef();
   ArrayRef<int64_t> sourceOffsets = sourceOffsetsAttr.asArrayRef();
   ArrayRef<int64_t> strides = stridesAttr.asArrayRef();
   if (strides.size() != offsets.size())
     return emitError("fragment stride and offset arrays must have the same length");
   if (sourceOffsets.size() != operandIndices.size())
     return emitError("fragment source offset count must match fragment operand index count");
+  if (sourceSlots.size() != operandIndices.size())
+    return emitError("fragment source slot count must match fragment operand index count");
   if (!getConflictPolicyAttr() || !getCoveragePolicyAttr())
     return emitError("fragment assembly blueprint requires conflict and coverage policies");
   if (getConflictPolicy() != "disjoint")
@@ -539,14 +628,19 @@ LogicalResult SpatBlueprintOp::verify() {
     int64_t operandIndex = operandIndices[fragmentIndex];
     if (operandIndex < 0 || operandIndex >= operandCount)
       return emitError("fragment assembly operand index is out of range");
+    if (sourceSlots[fragmentIndex] < 0)
+      return emitError("fragment assembly physical source slot must be nonnegative");
     if (sourceOffsets[fragmentIndex] < 0)
       return emitError("fragment assembly source offsets must be nonnegative");
 
     auto operandType = dyn_cast<RankedTensorType>(operands[operandIndex].getType());
     if (!operandType || !operandType.hasStaticShape())
       return emitError("fragment assembly blueprint requires static ranked tensor operands");
-    if (operandType.getRank() != rank)
-      return emitError("fragment assembly blueprint requires operand/result rank match");
+    if (operandType.getRank() != rank + 1)
+      return emitError("fragment assembly physical operand must have one leading source-slot dimension");
+    if (sourceSlots[fragmentIndex] >= operandType.getDimSize(0))
+      return emitError("fragment assembly physical source slot is out of range");
+    auto fragmentType = RankedTensorType::get(operandType.getShape().drop_front(), operandType.getElementType(), operandType.getEncoding());
 
     SmallVector<int64_t, 4> fragmentOffsets;
     SmallVector<int64_t, 4> fragmentSizes;
@@ -562,12 +656,12 @@ LogicalResult SpatBlueprintOp::verify() {
     int64_t fragmentElements = 1;
     for (int64_t dim = 0; dim < rank; ++dim)
       fragmentElements *= fragmentSizes[dim];
-    if (sourceOffsets[fragmentIndex] + fragmentElements > operandType.getNumElements())
-      return emitError("fragment assembly source offset exceeds the operand bounds");
+    if (sourceOffsets[fragmentIndex] + fragmentElements > fragmentType.getNumElements())
+      return emitError("fragment assembly source offset exceeds the selected physical fragment bounds");
     SmallVector<int64_t, 4> sourceSliceOffsets =
-      expandFlatElementIndex(sourceOffsets[fragmentIndex], operandType.getShape());
+      expandFlatElementIndex(sourceOffsets[fragmentIndex], fragmentType.getShape());
     for (int64_t dim = 0; dim < rank; ++dim)
-      if (sourceSliceOffsets[dim] + fragmentSizes[dim] > operandType.getDimSize(dim))
+      if (sourceSliceOffsets[dim] + fragmentSizes[dim] > fragmentType.getDimSize(dim))
         return emitError("fragment assembly source offset must describe a valid unit-stride slice");
 
     for (const auto& [existingOffsets, existingSizes] : slices) {
@@ -630,7 +724,9 @@ LogicalResult verifyComputeResultsUses(Operation* op) {
   if (!isAnySpatialComputeLike(op))
     return op->emitError("verifyComputeResultUses: op is not a Spatial compute-like operation");
   if (!llvm::all_of(op->getResults(), [](Value result) {
-        return llvm::all_of(result.getUsers(), [](Operation* op) {
+        return llvm::all_of(result.getUsers(), [result](Operation* op) {
+          if (isRecordedDeferredCommunicationSource(op, result))
+            return true;
           return !isAnySpatialComputeLike(op->getParentOp());
         });
       })) {
@@ -641,57 +737,68 @@ LogicalResult verifyComputeResultsUses(Operation* op) {
 
 template <typename ComputeOpTy>
 LogicalResult verifyComputeLikeOp(ComputeOpTy compute, StringRef opName) {
-  auto& block = compute.getBody().front();
   unsigned expectedArgCount = compute.getWeights().size() + compute.getInputs().size();
-  if (block.getNumArguments() != expectedArgCount)
-    return compute.emitOpError("compute body must have weight and input block arguments");
+  bool isScheduled = isa<SpatScheduledCompute>(compute.getOperation());
+  if (compute.getBody().empty())
+    return compute.emitOpError("compute body must have at least one block");
 
-  for (auto [weightIndex, weight] : llvm::enumerate(compute.getWeights())) {
-    auto blockArg = compute.getWeightArgument(weightIndex);
-    if (!blockArg || blockArg->getType() != weight.getType())
-      return compute.emitOpError("compute weight block argument types must match weight operand types exactly");
-  }
-  for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs())) {
-    auto blockArg = compute.getInputArgument(inputIndex);
-    if (!blockArg || blockArg->getType() != input.getType())
-      return compute.emitOpError("compute input block argument types must match input operand types exactly");
-  }
+  SmallVector<Type> yieldedTypes;
+  for (Block& block : compute.getBody()) {
+    if ((!isScheduled && block.getNumArguments() != expectedArgCount)
+        || (isScheduled && block.getNumArguments() < expectedArgCount))
+      return compute.emitOpError("compute body must have weight and input block arguments");
 
-  if (block.mightHaveTerminator()) {
-    auto yieldOp = dyn_cast_or_null<SpatYieldOp>(block.getTerminator());
-    if (!yieldOp)
+    for (auto [weightIndex, weight] : llvm::enumerate(compute.getWeights()))
+      if (block.getArgument(weightIndex).getType() != weight.getType())
+        return compute.emitOpError("compute weight block argument types must match weight operand types exactly");
+    for (auto [inputIndex, input] : llvm::enumerate(compute.getInputs()))
+      if (block.getArgument(compute.getWeights().size() + inputIndex).getType() != input.getType())
+        return compute.emitOpError("compute input block argument types must match input operand types exactly");
+
+    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)
       return compute.emitOpError("ComputeOp must have a single yield operation");
+    if (blockYield->getNumSuccessors() == 0)
+      llvm::append_range(yieldedTypes, blockYield->getOperandTypes());
+  }
 
-    auto resultTypes = compute.getResultTypes();
-    auto yieldTypes = yieldOp->getOperandTypes();
-    if (resultTypes.size() != yieldTypes.size())
-      return compute.emitOpError("ComputeOp must have same number of results as yieldOp operands");
+  auto resultTypes = compute.getResultTypes();
+  if (resultTypes.size() != yieldedTypes.size())
+    return compute.emitOpError("ComputeOp must have same number of results as yielded operands");
 
-    for (auto it : llvm::reverse(llvm::zip(resultTypes, yieldTypes))) {
-      auto resultType = std::get<0>(it);
-      auto yieldType = std::get<1>(it);
+  for (auto it : llvm::reverse(llvm::zip(resultTypes, yieldedTypes))) {
+    auto resultType = std::get<0>(it);
+    auto yieldType = std::get<1>(it);
 
-      if (resultType != yieldType || failed(verifyCompatibleShape(resultType, yieldType)))
-        return compute.emitOpError("ComputeOp output must be of the same type as yieldOp operand");
+    if (resultType != yieldType || failed(verifyCompatibleShape(resultType, yieldType)))
+      return compute.emitOpError("ComputeOp output must be of the same type as yieldOp operand");
 
-      if (auto resultRankedType = dyn_cast<RankedTensorType>(resultType)) {
-        if (auto yieldRankedType = dyn_cast<RankedTensorType>(yieldType)) {
-          if (resultRankedType.getEncoding() != yieldRankedType.getEncoding())
-            return compute.emitOpError("ComputeOp output must have the same encoding as yieldOp operand");
-        }
-        else {
+    if (auto resultRankedType = dyn_cast<RankedTensorType>(resultType)) {
+      if (auto yieldRankedType = dyn_cast<RankedTensorType>(yieldType)) {
+        if (resultRankedType.getEncoding() != yieldRankedType.getEncoding())
           return compute.emitOpError("ComputeOp output has an encoding while yieldOp operand does not have one");
-        }
       }
-      else if (dyn_cast<RankedTensorType>(yieldType)) {
-        return compute.emitOpError("ComputeOp output must not have an encoding if yieldOp operand has one");
+      else {
+        return compute.emitOpError("ComputeOp output must have the same encoding as yieldOp operand");
       }
     }
+    else if (dyn_cast<RankedTensorType>(yieldType)) {
+      return compute.emitOpError("ComputeOp output must not have an encoding if yieldOp operand has one");
+    }
   }
 
-  for (unsigned inputIndex = 0; inputIndex < compute.getInputs().size(); ++inputIndex)
-    if (auto inputArg = compute.getInputArgument(inputIndex); !inputArg || inputArg->use_empty())
-      return compute.emitOpError("ComputeOp block argument is not used");
+  if (compute.getBody().hasOneBlock())
+    for (unsigned inputIndex = 0; inputIndex < compute.getInputs().size(); ++inputIndex)
+      if (auto inputArg = compute.getInputArgument(inputIndex); !inputArg || inputArg->use_empty())
+        return compute.emitOpError("ComputeOp block argument is not used");
   if (failed(verifyStaticWeights(compute, opName)))
     return failure();
   if (failed(verifyOnlyConstantExternalValues(compute.getOperation(), compute.getBody(), opName)))
@@ -705,6 +812,41 @@ LogicalResult SpatGraphCompute::verify() { return verifyComputeLikeOp(*this, "sp
 
 LogicalResult SpatScheduledCompute::verify() { return verifyComputeLikeOp(*this, "spat.scheduled_compute"); }
 
+LogicalResult SpatBlockYieldOp::verify() {
+  if (getOperation()->getNumSuccessors() > 1)
+    return emitOpError("may target at most one next scheduled block");
+  Operation* parent = getOperation()->getParentOp();
+  if (!isa_and_nonnull<SpatScheduledCompute>(parent))
+    return emitOpError("expected spat.scheduled_compute parent");
+  if (getOperation()->getNumSuccessors() == 1) {
+    Block* next = getOperation()->getSuccessor(0);
+    if (getOperation()->getNumOperands() != next->getNumArguments())
+      return emitOpError("successor operand count must match next block argument count");
+    for (auto [operand, argument] : llvm::zip(getOperation()->getOperands(), next->getArguments()))
+      if (operand.getType() != argument.getType())
+        return emitOpError("successor operand types must match next block argument types");
+  }
+  return success();
+}
+
+LogicalResult SpatDeferredCommunicationOp::verify() {
+  if (getSources().empty())
+    return emitOpError("requires at least one source");
+  static constexpr StringLiteral staleAttributes[] = {
+    "exchangeId", "logicalProducer", "logicalConsumer", "sourceClass", "targetClass", "sourceCore",
+    "targetCore", "sourceLane", "targetLane", "transferKind", "resultIndex", "projectedTransfer",
+    "hostOutputOwner", "source_cpus", "source_classes", "source_lane_ranges", "target_cpus",
+    "target_classes", "target_lane_ranges", "batched", "source_operand_for_scheduled_lane",
+    "multi_source_payload"};
+  for (StringLiteral name : staleAttributes)
+    if (getOperation()->hasAttr(name))
+      return emitOpError() << "does not accept stale routing attribute '" << name
+                           << "'; source selection and shaping belong in the body and routing is derived in Phase 2";
+  if (failed(verifyRegionArguments(getOperation(), getBody(), getSources(), "spat.deferred_communication")))
+    return failure();
+  return verifyYieldTypes(getOperation(), getBody(), getOperation()->getResultTypes(), "spat.deferred_communication");
+}
+
 template <typename ComputeBatchOpTy>
 LogicalResult verifyComputeBatchLikeOp(ComputeBatchOpTy batch, StringRef opName) {
   int32_t count = batch.getLaneCount();
@@ -727,30 +869,33 @@ LogicalResult verifyComputeBatchLikeOp(ComputeBatchOpTy batch, StringRef opName)
         return batch.emitOpError("compute_batch coreIds values must be unique");
   }
 
-  Block& block = batch.getBody().front();
-  if (block.getNumArguments() == 0)
-    return batch.emitOpError("compute_batch body must have exactly one lane block argument");
-  unsigned expectedArgCount = 1 + batch.getWeights().size() + batch.getInputs().size() + batch.getNumResults();
-  if (block.getNumArguments() != expectedArgCount)
-    return batch.emitOpError("compute_batch body block arguments must match lane, weight, input, and output operands/results");
-  auto laneArg = batch.getLaneArgument();
-  if (!laneArg || !laneArg->getType().isIndex())
-    return batch.emitOpError("compute_batch first block argument must have index type");
+  if (batch.getBody().empty())
+    return batch.emitOpError("compute_batch body must have at least one block");
 
-  for (auto [weightIndex, weight] : llvm::enumerate(batch.getWeights())) {
-    auto blockArg = batch.getWeightArgument(weightIndex);
-    if (!blockArg || blockArg->getType() != weight.getType())
-      return batch.emitOpError("compute_batch weight block argument types must match weight operand types exactly");
-  }
-  for (auto [inputIndex, input] : llvm::enumerate(batch.getInputs())) {
-    auto blockArg = batch.getInputArgument(inputIndex);
-    if (!blockArg || blockArg->getType() != input.getType())
-      return batch.emitOpError("compute_batch input block argument types must match input operand types exactly");
-  }
-  for (auto [resultIndex, resultType] : llvm::enumerate(batch.getResultTypes())) {
-    auto blockArg = batch.getOutputArgument(resultIndex);
-    if (!blockArg || blockArg->getType() != resultType)
-      return batch.emitOpError("compute_batch output block argument types must match result types exactly");
+  unsigned expectedArgCount = 1 + batch.getWeights().size() + batch.getInputs().size() + batch.getNumResults();
+  bool verifyLaneSliceOffsets = !isa<SpatScheduledComputeBatch>(batch.getOperation());
+  for (Block& block : batch.getBody()) {
+    if (block.getNumArguments() == 0)
+      return batch.emitOpError("compute_batch body must have exactly one lane block argument");
+    if (block.getNumArguments() != expectedArgCount)
+      return batch.emitOpError(
+        "compute_batch body block arguments must match lane, weight, input, and output operands/results");
+    if (!block.getArgument(0).getType().isIndex())
+      return batch.emitOpError("compute_batch first block argument must have index type");
+
+    for (auto [weightIndex, weight] : llvm::enumerate(batch.getWeights()))
+      if (block.getArgument(1 + weightIndex).getType() != weight.getType())
+        return batch.emitOpError("compute_batch weight block argument types must match weight operand types exactly");
+    for (auto [inputIndex, input] : llvm::enumerate(batch.getInputs()))
+      if (block.getArgument(1 + batch.getWeights().size() + inputIndex).getType() != input.getType())
+        return batch.emitOpError("compute_batch input block argument types must match input operand types exactly");
+    for (auto [resultIndex, resultType] : llvm::enumerate(batch.getResultTypes()))
+      if (block.getArgument(1 + batch.getWeights().size() + batch.getInputs().size() + resultIndex).getType()
+          != resultType)
+        return batch.emitOpError("compute_batch output block argument types must match result types exactly");
+
+    if (failed(verifyBatchBody(batch, block, verifyLaneSliceOffsets)))
+      return failure();
   }
 
   if (failed(verifyComputeResultsUses(batch.getOperation())))
@@ -759,7 +904,7 @@ LogicalResult verifyComputeBatchLikeOp(ComputeBatchOpTy batch, StringRef opName)
     return failure();
   if (failed(verifyOnlyConstantExternalValues(batch.getOperation(), batch.getBody(), opName)))
     return failure();
-  return verifyBatchBody(batch, block);
+  return success();
 }
 
 LogicalResult SpatGraphComputeBatch::verify() { return verifyComputeBatchLikeOp(*this, "spat.graph_compute_batch"); }

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

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

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

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

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

@@ -0,0 +1,461 @@
+#include "DeferredCommunicationPlanning.hpp"
+
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/SCF/IR/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/IR/IRMapping.h"
+
+#include "llvm/ADT/SmallPtrSet.h"
+
+namespace onnx_mlir::spatial {
+using namespace mlir;
+namespace {
+
+static Value getBlockOperand(Block &block, ValueRange operands, Value value, unsigned firstArgument = 0) {
+  auto it = llvm::find(operands, value);
+  assert(it != operands.end() && "missing scheduled operand");
+  return block.getArgument(firstArgument + std::distance(operands.begin(), it));
+}
+
+static FailureOr<Value> getOriginalProducerValue(const ProducerValueRef &producer) {
+  auto outputs = getComputeInstanceOutputValues(producer.instance);
+  if (producer.resultIndex >= outputs.size())
+    return failure();
+  return outputs[producer.resultIndex];
+}
+
+static SmallVector<Value> getBlueprintFragments(SpatBlueprintOp blueprint) {
+  SmallVector<Value> fragments {blueprint.getInput()};
+  llvm::append_range(fragments, blueprint.getFragments());
+  return fragments;
+}
+
+static FailureOr<Value> buildBlueprintReconstruction(
+    OpBuilder &builder, Location loc, SpatBlueprintOp blueprint,
+    ValueRange sourceBlockArgs) {
+  auto resultType = dyn_cast<RankedTensorType>(blueprint.getOutput().getType());
+  auto operandIndices = blueprint.getFragmentOperandIndices();
+  auto sourceSlots = blueprint.getFragmentSourceSlots();
+  auto sourceOffsets = blueprint.getFragmentSourceOffsets();
+  auto strides = blueprint.getFragmentStrides();
+  if (!resultType || !resultType.hasStaticShape() || !operandIndices ||
+      !sourceSlots || !sourceOffsets || !strides)
+    return blueprint.emitOpError("phase 1 requires complete static fragment assembly metadata"), failure();
+  int64_t rank = resultType.getRank();
+  ArrayRef<int64_t> offsets = blueprint.getFragmentOffsets();
+  ArrayRef<int64_t> sizes = blueprint.getFragmentSizes();
+  if (offsets.size() != sizes.size() || offsets.size() != strides->size() ||
+      offsets.size() != operandIndices->size() * rank ||
+      sourceSlots->size() != operandIndices->size() ||
+      sourceOffsets->size() != operandIndices->size())
+    return blueprint.emitOpError("phase 1 fragment assembly metadata has inconsistent sizes"), failure();
+
+  Value result = tensor::EmptyOp::create(builder, loc, resultType.getShape(),
+                                         resultType.getElementType());
+  for (auto [fragmentIndex, operandIndex] : llvm::enumerate(*operandIndices)) {
+    if (operandIndex < 0 || operandIndex >= static_cast<int64_t>(sourceBlockArgs.size()))
+      return blueprint.emitOpError("phase 1 fragment assembly operand index is out of range"), failure();
+    auto physicalType = dyn_cast<RankedTensorType>(sourceBlockArgs[operandIndex].getType());
+    if (!physicalType || !physicalType.hasStaticShape() || physicalType.getRank() != rank + 1)
+      return blueprint.emitOpError("phase 1 fragment assembly source is not a physical fragment batch"), failure();
+    SmallVector<int64_t> fragmentShape(physicalType.getShape().drop_front());
+    int64_t linearOffset = (*sourceOffsets)[fragmentIndex];
+    SmallVector<int64_t> sourceCoordinates(rank);
+    for (int64_t dim = rank - 1; dim >= 0; --dim) {
+      sourceCoordinates[dim] = linearOffset % fragmentShape[dim];
+      linearOffset /= fragmentShape[dim];
+    }
+    if (linearOffset != 0)
+      return blueprint.emitOpError("phase 1 fragment source offset is out of range"), failure();
+
+    SmallVector<OpFoldResult> sliceOffsets, sliceSizes, sliceStrides;
+    sliceOffsets.push_back(builder.getIndexAttr((*sourceSlots)[fragmentIndex]));
+    sliceSizes.push_back(builder.getIndexAttr(1));
+    sliceStrides.push_back(builder.getIndexAttr(1));
+    SmallVector<int64_t> selectedShape {1};
+    for (int64_t dim = 0; dim < rank; ++dim) {
+      int64_t index = fragmentIndex * rank + dim;
+      int64_t size = sizes[index];
+      if ((*strides)[index] != 1 || sourceCoordinates[dim] < 0 || size <= 0 ||
+          sourceCoordinates[dim] + size > fragmentShape[dim])
+        return blueprint.emitOpError("phase 1 fragment geometry is unsupported"), failure();
+      sliceOffsets.push_back(builder.getIndexAttr(sourceCoordinates[dim]));
+      sliceSizes.push_back(builder.getIndexAttr(size));
+      sliceStrides.push_back(builder.getIndexAttr(1));
+      selectedShape.push_back(size);
+    }
+    auto selectedType = RankedTensorType::get(selectedShape, resultType.getElementType());
+    Value selected = tensor::ExtractSliceOp::create(
+        builder, loc, selectedType, sourceBlockArgs[operandIndex], sliceOffsets,
+        sliceSizes, sliceStrides);
+    SmallVector<int64_t> fragmentResultShape(selectedShape.begin() + 1,
+                                             selectedShape.end());
+    auto fragmentType = RankedTensorType::get(fragmentResultShape,
+                                              resultType.getElementType());
+    SmallVector<ReassociationIndices> reassociation {{0, 1}};
+    for (int64_t dim = 1; dim < rank; ++dim)
+      reassociation.push_back({dim + 1});
+    Value fragment = tensor::CollapseShapeOp::create(
+        builder, loc, fragmentType, selected, reassociation);
+    SmallVector<OpFoldResult> targetOffsets, targetSizes, targetStrides;
+    for (int64_t dim = 0; dim < rank; ++dim) {
+      int64_t index = fragmentIndex * rank + dim;
+      targetOffsets.push_back(builder.getIndexAttr(offsets[index]));
+      targetSizes.push_back(builder.getIndexAttr(sizes[index]));
+      targetStrides.push_back(builder.getIndexAttr((*strides)[index]));
+    }
+    result = tensor::InsertSliceOp::create(builder, loc, fragment, result,
+                                           targetOffsets, targetSizes,
+                                           targetStrides);
+  }
+  return result;
+}
+
+static bool isSupportedDeferredShapingOp(Operation *op) {
+  return isa<tensor::ExtractSliceOp, tensor::InsertSliceOp, tensor::CollapseShapeOp,
+             tensor::ExpandShapeOp, tensor::CastOp, tensor::EmptyOp, tensor::ExtractOp,
+             arith::ConstantOp, arith::IndexCastOp, arith::AddIOp, arith::SubIOp,
+             arith::MulIOp, affine::AffineApplyOp>(op);
+}
+
+static FailureOr<Value> buildSelectedDeferredSource(OpBuilder &builder, Location loc,
+                                                    SpatDeferredCommunicationOp transfer,
+                                                    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();
+  Value table = createI64LookupTableConstant(builder, transfer.getOperation(), sourceOperandForScheduledLane);
+  Value i64 = tensor::ExtractOp::create(builder, loc, table, ValueRange {scheduledLane}).getResult();
+  Value index = arith::IndexCastOp::create(builder, loc, builder.getIndexType(), i64).getResult();
+  auto type = dyn_cast<RankedTensorType>(sourceBlockArgs.front().getType());
+  if (!type || !type.hasStaticShape())
+    return transfer.emitOpError("multiple deferred sources require static ranked tensors"), failure();
+  for (Value source : sourceBlockArgs)
+    if (source.getType() != type)
+      return transfer.emitOpError("multiple deferred sources require identical tensor types"), failure();
+  SmallVector<int64_t> shape {static_cast<int64_t>(sourceBlockArgs.size())};
+  llvm::append_range(shape, type.getShape());
+  auto stacked = createEmptyTensorForType(builder, loc, RankedTensorType::get(shape, type.getElementType()));
+  if (failed(stacked))
+    return failure();
+  Value value = *stacked;
+  SmallVector<OpFoldResult> sizes, strides(shape.size(), builder.getIndexAttr(1));
+  sizes.push_back(builder.getIndexAttr(1));
+  for (int64_t dim : type.getShape()) sizes.push_back(builder.getIndexAttr(dim));
+  for (auto [i, source] : llvm::enumerate(sourceBlockArgs)) {
+    SmallVector<int64_t> expandedShape {1}; llvm::append_range(expandedShape, type.getShape());
+    SmallVector<ReassociationIndices> reassociation {{0, 1}};
+    for (int64_t dim = 1; dim < type.getRank(); ++dim) reassociation.push_back({dim + 1});
+    Value expanded = tensor::ExpandShapeOp::create(builder, loc,
+      RankedTensorType::get(expandedShape, type.getElementType()), source, reassociation).getResult();
+    SmallVector<OpFoldResult> offsets(shape.size(), builder.getIndexAttr(0));
+    offsets[0] = builder.getIndexAttr(i);
+    value = tensor::InsertSliceOp::create(builder, loc, expanded, value, offsets, sizes, strides).getResult();
+  }
+  SmallVector<OpFoldResult> offsets(shape.size(), builder.getIndexAttr(0)); offsets[0] = index;
+  SmallVector<int64_t> sliceShape {1}; llvm::append_range(sliceShape, type.getShape());
+  auto slice = tensor::ExtractSliceOp::create(builder, loc,
+    RankedTensorType::get(sliceShape, type.getElementType()), value, offsets, sizes, strides);
+  // extract has a leading unit dimension; remove it without changing the payload.
+  SmallVector<ReassociationIndices> reassociation {{0, 1}};
+  for (int64_t dim = 1; dim < type.getRank(); ++dim) reassociation.push_back({dim + 1});
+  return tensor::CollapseShapeOp::create(builder, loc, type, slice.getResult(), reassociation).getResult();
+}
+
+static bool isTopLevelShaping(Operation *op, Block &body) {
+  return op->getBlock() == &body && isSupportedDeferredShapingOp(op);
+}
+
+static bool isEligible(Value value, Block &body, const DeferredInputPlan &plan,
+                       llvm::SmallPtrSetImpl<Operation *> &seen) {
+  if (value == plan.graphInput || value == plan.graphLane || value == plan.scheduledLane)
+    return true;
+  auto arg = dyn_cast<BlockArgument>(value);
+  if (arg)
+    return false;
+  Operation *op = value.getDefiningOp();
+  if (op && op->hasTrait<OpTrait::ConstantLike>())
+    return true;
+  if (!op || !isTopLevelShaping(op, body) || !seen.insert(op).second)
+    return op && seen.contains(op);
+  return llvm::all_of(op->getOperands(), [&](Value operand) { return isEligible(operand, body, plan, seen); });
+}
+
+static FailureOr<Value> clonePayloadRoot(Value root, Block &body, const DeferredInputPlan &plan,
+                                         OpBuilder &builder, SpatDeferredCommunicationOp transfer,
+                                         Value selectedSource, Value boundGraphLane) {
+  IRMapping mapping;
+  mapping.map(plan.graphInput, selectedSource);
+  std::function<FailureOr<Value>(Value)> cloneScheduledLane = [&](Value value) -> FailureOr<Value> {
+    if (mapping.contains(value)) return mapping.lookup(value);
+    if (value == plan.scheduledLane) return value;
+    if (isa<BlockArgument>(value))
+      return transfer.emitOpError("phase 1 payload shaping captures an unsupported block argument"), failure();
+    Operation *op = value.getDefiningOp();
+    if (!op || !isSupportedDeferredShapingOp(op))
+      return transfer.emitOpError("phase 1 cannot clone the scheduled graph-lane expression"), failure();
+    for (Value operand : op->getOperands()) if (failed(cloneScheduledLane(operand))) return failure();
+    Operation *copy = builder.clone(*op, mapping);
+    for (auto pair : llvm::zip(op->getResults(), copy->getResults())) mapping.map(std::get<0>(pair), std::get<1>(pair));
+    return mapping.lookup(value);
+  };
+  std::function<FailureOr<Value>(Value)> clone = [&](Value value) -> FailureOr<Value> {
+    if (mapping.contains(value)) return mapping.lookup(value);
+    if (value == plan.graphLane) {
+      auto mappedLane = cloneScheduledLane(boundGraphLane ? boundGraphLane : plan.scheduledGraphLane);
+      if (failed(mappedLane)) return failure();
+      mapping.map(value, *mappedLane);
+      return *mappedLane;
+    }
+    if (isa<BlockArgument>(value))
+      return transfer.emitOpError("phase 1 payload shaping captures an unsupported block argument"), failure();
+    Operation *op = value.getDefiningOp();
+    if (!op || (!isTopLevelShaping(op, body) && !op->hasTrait<OpTrait::ConstantLike>()))
+      return transfer.emitOpError("phase 1 payload shaping contains an unsupported operation"), failure();
+    for (Value operand : op->getOperands()) if (failed(clone(operand))) return failure();
+    Operation *copy = builder.clone(*op, mapping);
+    for (auto pair : llvm::zip(op->getResults(), copy->getResults())) mapping.map(std::get<0>(pair), std::get<1>(pair));
+    return mapping.lookup(value);
+  };
+  return clone(root);
+}
+
+static bool dependsOnGraphLane(Value value, Value graphLane, Block &body,
+                               llvm::SmallPtrSetImpl<Operation *> &seen) {
+  if (value == graphLane)
+    return true;
+  Operation *op = value.getDefiningOp();
+  if (!op || !isTopLevelShaping(op, body) || !seen.insert(op).second)
+    return false;
+  return llvm::any_of(op->getOperands(), [&](Value operand) {
+    return dependsOnGraphLane(operand, graphLane, body, seen);
+  });
+}
+
+static FailureOr<Value> buildPayloadAggregate(OpBuilder &builder, Location loc,
+                                              ArrayRef<Value> payloads) {
+  auto payloadType = dyn_cast<RankedTensorType>(payloads.front().getType());
+  if (!payloadType || !payloadType.hasStaticShape())
+    return failure();
+  SmallVector<int64_t> shape {static_cast<int64_t>(payloads.size())};
+  llvm::append_range(shape, payloadType.getShape());
+  auto aggregateType = RankedTensorType::get(shape, payloadType.getElementType());
+  auto empty = createEmptyTensorForType(builder, loc, aggregateType);
+  if (failed(empty)) return failure();
+  Value aggregate = *empty;
+  SmallVector<OpFoldResult> sizes, strides(shape.size(), builder.getIndexAttr(1));
+  sizes.push_back(builder.getIndexAttr(1));
+  for (int64_t dim : payloadType.getShape()) sizes.push_back(builder.getIndexAttr(dim));
+  SmallVector<ReassociationIndices> reassociation {{0, 1}};
+  for (int64_t dim = 1; dim < payloadType.getRank(); ++dim) reassociation.push_back({dim + 1});
+  SmallVector<int64_t> expandedShape {1}; llvm::append_range(expandedShape, payloadType.getShape());
+  auto expandedType = RankedTensorType::get(expandedShape, payloadType.getElementType());
+  for (auto [index, payload] : llvm::enumerate(payloads)) {
+    Value expanded = tensor::ExpandShapeOp::create(builder, loc, expandedType, payload, reassociation).getResult();
+    SmallVector<OpFoldResult> offsets(shape.size(), builder.getIndexAttr(0));
+    offsets[0] = builder.getIndexAttr(index);
+    aggregate = tensor::InsertSliceOp::create(builder, loc, expanded, aggregate, offsets, sizes, strides).getResult();
+  }
+  return aggregate;
+}
+
+static FailureOr<Value> selectPayloadAggregate(OpBuilder &builder, Location loc, Value aggregate,
+                                               Value localLane) {
+  auto aggregateType = cast<RankedTensorType>(aggregate.getType());
+  SmallVector<int64_t> payloadShape(aggregateType.getShape().begin() + 1, aggregateType.getShape().end());
+  auto payloadType = RankedTensorType::get(payloadShape, aggregateType.getElementType());
+  SmallVector<OpFoldResult> offsets(aggregateType.getRank(), builder.getIndexAttr(0)); offsets[0] = localLane;
+  SmallVector<OpFoldResult> sizes, strides(aggregateType.getRank(), builder.getIndexAttr(1));
+  sizes.push_back(builder.getIndexAttr(1));
+  for (int64_t dim : payloadShape) sizes.push_back(builder.getIndexAttr(dim));
+  SmallVector<int64_t> unitShape {1}; llvm::append_range(unitShape, payloadShape);
+  Value unit = tensor::ExtractSliceOp::create(builder, loc,
+    RankedTensorType::get(unitShape, aggregateType.getElementType()), aggregate, offsets, sizes, strides).getResult();
+  SmallVector<ReassociationIndices> reassociation {{0, 1}};
+  for (int64_t dim = 1; dim < payloadType.getRank(); ++dim) reassociation.push_back({dim + 1});
+  return tensor::CollapseShapeOp::create(builder, loc, payloadType, unit, reassociation).getResult();
+}
+
+static void collectClosure(Value value, Block &body, const DeferredInputPlan &plan,
+                           llvm::SmallPtrSetImpl<Operation *> &ops) {
+  Operation *op = value.getDefiningOp();
+  if (!op || !isTopLevelShaping(op, body) || !ops.insert(op).second) return;
+  for (Value operand : op->getOperands())
+    if (operand != plan.graphInput && operand != plan.graphLane) collectClosure(operand, body, plan, ops);
+}
+
+} // namespace
+
+bool isDeferredFragmentAssemblyInput(
+    Value input, const ComputeInstance &consumerInstance) {
+  auto blueprint = input.getDefiningOp<SpatBlueprintOp>();
+  if (!blueprint || blueprint.getMode() != "fragment_assembly")
+    return false;
+  return llvm::all_of(getBlueprintFragments(blueprint), [&](Value fragment) {
+    return getProducerValueRef(fragment, &consumerInstance).has_value();
+  });
+}
+
+LogicalResult prepareSingleCpuInput(OpBuilder &, Location loc, Value input, BlockArgument graphInput,
+                                    const ComputeInstance &consumerInstance, const MergeScheduleResult &,
+                                    ValueRange scheduledInputs, Block &block, unsigned firstInputArgument,
+                                    ArrayRef<ProducerValueKey> carriedKeys, Value graphLane, Value scheduledGraphLane,
+                                    DeferredInputPlan &plan) {
+  plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, {}, {}, {}, {}, 1, nullptr};
+  if (isDeferredFragmentAssemblyInput(input, consumerInstance)) {
+    plan.blueprint = input.getDefiningOp<SpatBlueprintOp>();
+    plan.originalSources = getBlueprintFragments(plan.blueprint);
+    return success();
+  }
+  auto producer = getProducerValueRef(input, &consumerInstance);
+  if (!producer) { plan.availableValue = getBlockOperand(block, scheduledInputs, input, firstInputArgument); return success(); }
+  ProducerValueKey key {producer->instance, producer->resultIndex};
+  auto carried = llvm::find(carriedKeys, key);
+  if (carried != carriedKeys.end()) {
+    plan.availableValue = block.getArgument(firstInputArgument + scheduledInputs.size() + std::distance(carriedKeys.begin(), carried));
+    return success();
+  }
+  auto source = getOriginalProducerValue(*producer);
+  if (failed(source)) return emitError(loc) << "cannot resolve original graph producer value";
+  plan.originalSources.push_back(*source);
+  return success();
+}
+
+LogicalResult prepareMultiCpuTupleInput(OpBuilder &, Location loc, Value input, BlockArgument graphInput,
+                                        const ComputeStepTuple &tuple, const PeftClassPlan &,
+                                        const MergeScheduleResult &, ValueRange scheduledInputs, Block &block,
+                                        unsigned firstInputArgument, Value graphLane, Value scheduledGraphLane, Value scheduledLane,
+                                        DeferredInputPlan &plan) {
+  const ComputeInstance &representative = tuple.instances.front();
+  plan = {graphInput, {}, {}, {}, graphLane, scheduledGraphLane, scheduledLane, {}, {}, {}, 1, nullptr};
+  if (isDeferredFragmentAssemblyInput(input, representative)) {
+    plan.blueprint = input.getDefiningOp<SpatBlueprintOp>();
+    plan.originalSources = getBlueprintFragments(plan.blueprint);
+    return success();
+  }
+  auto producer = getProducerValueRef(input, &representative);
+  if (!producer) { plan.availableValue = getBlockOperand(block, scheduledInputs, input, firstInputArgument); return success(); }
+  auto inputs = getComputeInstanceInputs(representative);
+  auto it = llvm::find(inputs, input);
+  if (it == inputs.end()) return emitError(loc) << "cannot resolve scheduled batch step input";
+  unsigned inputIndex = std::distance(inputs.begin(), it);
+  for (const ComputeInstance &instance : tuple.instances) {
+    auto laneInputs = getComputeInstanceInputs(instance);
+    if (inputIndex >= laneInputs.size()) return emitError(loc) << "scheduled batch step input out of range";
+    auto laneProducer = getProducerValueRef(laneInputs[inputIndex], &instance);
+    if (!laneProducer) return emitError(loc) << "scheduled batch step mixes host and producer inputs";
+    auto source = getOriginalProducerValue(*laneProducer);
+    if (failed(source)) return emitError(loc) << "cannot resolve original graph producer value";
+    auto sourceIt = llvm::find(plan.originalSources, *source);
+    if (sourceIt == plan.originalSources.end()) { plan.sourceOperandForScheduledLane.push_back(plan.originalSources.size()); plan.originalSources.push_back(*source); }
+    else plan.sourceOperandForScheduledLane.push_back(std::distance(plan.originalSources.begin(), sourceIt));
+  }
+  return success();
+}
+
+LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc, Block &body,
+                                                 ArrayRef<DeferredInputPlan> plans, IRMapping &mapper,
+                                                 llvm::SmallPtrSetImpl<Operation *> &absorbed) {
+  for (const DeferredInputPlan &plan : plans) {
+    if (plan.availableValue) { mapper.map(plan.graphInput, plan.availableValue); continue; }
+    SmallVector<Value> roots;
+    bool needsIdentity = false;
+    SmallVector<Value> worklist {plan.graphInput};
+    llvm::SmallDenseSet<Value, 32> seen;
+    while (!worklist.empty()) {
+      Value value = worklist.pop_back_val();
+      if (!seen.insert(value).second) continue;
+      for (OpOperand &use : value.getUses()) {
+        Operation *user = use.getOwner();
+        if (!isTopLevelShaping(user, body)) { needsIdentity = true; continue; }
+        llvm::SmallPtrSet<Operation *, 16> eligibility;
+        if (!isEligible(user->getResult(0), body, plan, eligibility)) { needsIdentity = true; continue; }
+        for (Value result : user->getResults()) {
+          bool hasShapingUse = llvm::any_of(result.getUses(), [&](OpOperand &next) { return isTopLevelShaping(next.getOwner(), body); });
+          bool hasOtherUse = llvm::any_of(result.getUses(), [&](OpOperand &next) { return !isTopLevelShaping(next.getOwner(), body); });
+          if (hasOtherUse) roots.push_back(result);
+          if (hasShapingUse) worklist.push_back(result);
+        }
+      }
+    }
+    if (needsIdentity) roots.push_back(plan.graphInput);
+    llvm::sort(roots, [](Value a, Value b) { return a.getAsOpaquePointer() < b.getAsOpaquePointer(); });
+    roots.erase(std::unique(roots.begin(), roots.end()), roots.end());
+    for (Value root : roots) {
+      llvm::SmallPtrSet<Operation *, 16> laneDependencies;
+      bool scalarize = plan.scalarizedGraphLaneBase
+                    && dependsOnGraphLane(root, plan.graphLane, body, laneDependencies);
+      OpBuilder::InsertPoint restore = builder.saveInsertionPoint();
+      Operation *loop = nullptr;
+      if (scalarize) {
+        loop = builder.getInsertionBlock()->getParentOp();
+        if (loop && !isa<scf::ForOp>(loop))
+          loop = loop->getParentOfType<scf::ForOp>();
+        if (loop)
+          builder.setInsertionPoint(loop);
+        else if (plan.scalarizedHoistBlock)
+          builder.setInsertionPointToEnd(plan.scalarizedHoistBlock);
+        else
+          return emitError(loc) << "phase 1 scalarized deferred payload is missing a hoist point";
+      }
+      SmallVector<Value> payloads;
+      unsigned count = scalarize ? plan.scalarizedLaneCount : 1;
+      for (unsigned offset = 0; offset < count; ++offset) {
+      auto transfer = SpatDeferredCommunicationOp::create(builder, loc, root.getType(), plan.originalSources);
+      Block *deferred = builder.createBlock(&transfer.getBody(), transfer.getBody().end(),
+        TypeRange {transfer.getSources().getTypes()}, SmallVector<Location>(transfer.getSources().size(), loc));
+      builder.setInsertionPointToStart(deferred);
+      auto selected = plan.blueprint
+                        ? buildBlueprintReconstruction(builder, loc, plan.blueprint,
+                                                       deferred->getArguments())
+                        : buildSelectedDeferredSource(builder, loc, transfer,
+                                                      plan.scheduledLane,
+                                                      deferred->getArguments(),
+                                                      plan.sourceOperandForScheduledLane);
+      if (failed(selected)) return failure();
+      Value boundGraphLane;
+      if (scalarize) {
+        Value offsetValue = arith::ConstantIndexOp::create(builder, loc, offset);
+        boundGraphLane = offset ? arith::AddIOp::create(builder, loc, plan.scalarizedGraphLaneBase, offsetValue).getResult()
+                                : plan.scalarizedGraphLaneBase;
+      }
+      auto payload = clonePayloadRoot(root, body, plan, builder, transfer, *selected, boundGraphLane);
+      if (failed(payload)) return failure();
+      SpatYieldOp::create(builder, loc, *payload);
+      payloads.push_back(transfer.getOutput());
+      builder.setInsertionPointAfter(transfer);
+      }
+      if (scalarize) {
+        builder.restoreInsertionPoint(restore);
+        if (payloads.size() == 1) {
+          mapper.map(root, payloads.front());
+        } else {
+          if (loop) builder.setInsertionPoint(loop);
+          else builder.setInsertionPointToEnd(plan.scalarizedHoistBlock);
+          auto aggregate = buildPayloadAggregate(builder, loc, payloads);
+          if (failed(aggregate)) return failure();
+          builder.restoreInsertionPoint(restore);
+          auto selected = selectPayloadAggregate(builder, loc, *aggregate, plan.scalarizedLocalLane);
+          if (failed(selected)) return failure();
+          mapper.map(root, *selected);
+        }
+      } else {
+        mapper.map(root, payloads.front());
+      }
+      collectClosure(root, body, plan, absorbed);
+    }
+  }
+  for (Operation *op : absorbed) {
+    bool allResultsMapped = llvm::all_of(op->getResults(), [&](Value result) {
+      return mapper.contains(result) || llvm::all_of(result.getUses(), [&](OpOperand &use) { return absorbed.contains(use.getOwner()); });
+    });
+    if (!allResultsMapped) absorbed.erase(op);
+  }
+  return success();
+}
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,54 @@
+#pragma once
+
+#include "ScheduledComputePlan.hpp"
+
+namespace onnx_mlir {
+namespace spatial {
+
+// A graph input is either already available in the scheduled block, or is a
+// graph result whose individual deterministic payloads are materialized later.
+struct DeferredInputPlan {
+  BlockArgument graphInput;
+  Value availableValue;
+  SmallVector<Value> originalSources;
+  SmallVector<int64_t> sourceOperandForScheduledLane;
+  Value graphLane;
+  Value scheduledGraphLane;
+  Value scheduledLane;
+  SpatBlueprintOp blueprint;
+  Value scalarizedLocalLane;
+  Value scalarizedGraphLaneBase;
+  int64_t scalarizedLaneCount = 1;
+  Block *scalarizedHoistBlock = nullptr;
+};
+
+bool isDeferredFragmentAssemblyInput(Value input,
+                                     const ComputeInstance &consumerInstance);
+
+LogicalResult prepareSingleCpuInput(OpBuilder &builder, Location loc, Value input,
+                                    BlockArgument graphInput,
+                                    const ComputeInstance &consumerInstance,
+                                    const MergeScheduleResult &schedule,
+                                    ValueRange scheduledInputs, Block &block,
+                                    unsigned firstInputArgument,
+                                    ArrayRef<ProducerValueKey> carriedKeys,
+                                    Value graphLane, Value scheduledGraphLane,
+                                    DeferredInputPlan &plan);
+
+LogicalResult prepareMultiCpuTupleInput(OpBuilder &builder, Location loc, Value input,
+                                        BlockArgument graphInput,
+                                        const ComputeStepTuple &stepTuple,
+                                        const PeftClassPlan &peftClassPlan,
+                                        const MergeScheduleResult &schedule,
+                                        ValueRange scheduledInputs, Block &block,
+                                        unsigned firstInputArgument, Value graphLane, Value scheduledGraphLane,
+                                        Value scheduledLane, DeferredInputPlan &plan);
+
+LogicalResult materializeDeferredPayloadDemands(OpBuilder &builder, Location loc,
+                                                 Block &graphBody,
+                                                 ArrayRef<DeferredInputPlan> plans,
+                                                 IRMapping &mapper,
+                                                 llvm::SmallPtrSetImpl<Operation *> &absorbed);
+
+} // namespace spatial
+} // namespace onnx_mlir

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

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

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

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

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

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

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

@@ -1,134 +0,0 @@
-#include "HostOutputFinalization.hpp"
-
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/IR/BuiltinTypes.h"
-#include "mlir/IR/PatternMatch.h"
-
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/DenseSet.h"
-#include "llvm/ADT/STLExtras.h"
-
-#include "MaterializedClassState.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-
-using namespace mlir;
-
-namespace onnx_mlir::spatial {
-
-LogicalResult finalizeProjectedHostOutputFragments(MaterializerState& state) {
-  if (state.pendingProjectedHostOutputFragments.empty())
-    return success();
-
-  DenseMap<Value, SmallVector<PendingProjectedHostOutputFragment*, 16>> byOutput;
-  for (PendingProjectedHostOutputFragment& fragment : state.pendingProjectedHostOutputFragments)
-    byOutput[fragment.originalOutput].push_back(&fragment);
-
-  SmallVector<Value, 8> outputs;
-  outputs.reserve(byOutput.size());
-
-  auto returnOp = dyn_cast<func::ReturnOp>(state.func.getBody().front().getTerminator());
-  if (!returnOp)
-    return state.func.emitError("expected func.return terminator while finalizing projected host output fragments");
-
-  DenseSet<Value> seenOutputs;
-  for (Value returned : returnOp.getOperands()) {
-    if (!byOutput.contains(returned) || !seenOutputs.insert(returned).second)
-      continue;
-    outputs.push_back(returned);
-  }
-  if (outputs.size() != byOutput.size())
-    return state.func.emitError("projected host output fragments must be keyed by returned logical host outputs");
-
-  for (Value originalOutput : outputs) {
-    if (isa_and_present<SpatScheduledCompute, SpatScheduledComputeBatch>(originalOutput.getDefiningOp())) {
-      return state.func.emitError(
-        "projected host output assembly must be keyed by the original logical host output, not by a materialized scheduled result");
-    }
-
-    auto resultType = dyn_cast<RankedTensorType>(originalOutput.getType());
-    if (!resultType || !resultType.hasStaticShape())
-      return state.func.emitError("projected host output must have static ranked tensor type");
-
-    SmallVector<PendingProjectedHostOutputFragment*, 16>& fragments = byOutput[originalOutput];
-    llvm::sort(fragments, [](const PendingProjectedHostOutputFragment* lhs,
-                             const PendingProjectedHostOutputFragment* rhs) {
-      if (lhs->sourceClass != rhs->sourceClass)
-        return lhs->sourceClass < rhs->sourceClass;
-      if (lhs->publicationResultIndex != rhs->publicationResultIndex)
-        return lhs->publicationResultIndex < rhs->publicationResultIndex;
-      if (lhs->sourceFragmentOrdinal != rhs->sourceFragmentOrdinal)
-        return lhs->sourceFragmentOrdinal < rhs->sourceFragmentOrdinal;
-      return std::lexicographical_compare(lhs->offsets.begin(),
-                                          lhs->offsets.end(),
-                                          rhs->offsets.begin(),
-                                          rhs->offsets.end());
-    });
-
-    state.rewriter.setInsertionPoint(returnOp);
-    Location loc = fragments.front()->loc;
-    SmallVector<Value, 16> blueprintOperands;
-    SmallVector<int64_t, 16> fragmentOperandIndices;
-    SmallVector<int64_t, 16> fragmentSourceOffsets;
-    SmallVector<int64_t, 64> flatOffsets;
-    SmallVector<int64_t, 64> flatSizes;
-    SmallVector<int64_t, 64> flatStrides;
-    DenseMap<Value, int64_t> operandIndicesByValue;
-
-    for (PendingProjectedHostOutputFragment* fragmentRecord : fragments) {
-      if (fragmentRecord->sourceClass >= state.classes.size())
-        return state.func.emitError("projected host output fragment references an invalid source class");
-
-      MaterializedClass& sourceClass = state.classes[fragmentRecord->sourceClass];
-      if (fragmentRecord->publicationResultIndex >= sourceClass.op->getNumResults()) {
-        return sourceClass.op->emitError("projected host output fragment references an invalid publication result")
-               << " sourceClass=" << sourceClass.id
-               << " resultIndex=" << fragmentRecord->publicationResultIndex
-               << " resultCount=" << sourceClass.op->getNumResults();
-      }
-
-      Value operand = sourceClass.op->getResult(fragmentRecord->publicationResultIndex);
-      auto [operandIt, inserted] =
-        operandIndicesByValue.try_emplace(operand, static_cast<int64_t>(blueprintOperands.size()));
-      if (inserted)
-        blueprintOperands.push_back(operand);
-      fragmentOperandIndices.push_back(operandIt->second);
-      fragmentSourceOffsets.push_back(fragmentRecord->sourceElementOffset);
-      llvm::append_range(flatOffsets, fragmentRecord->offsets);
-      llvm::append_range(flatSizes, fragmentRecord->sizes);
-      llvm::append_range(flatStrides, fragmentRecord->strides);
-
-      auto operandType = dyn_cast<RankedTensorType>(operand.getType());
-      if (!operandType || !operandType.hasStaticShape())
-        return state.func.emitError("projected host output assembly requires static ranked tensor operands");
-    }
-
-    if (blueprintOperands.empty())
-      return state.func.emitError("missing projected host output fragments");
-
-    Value input = blueprintOperands.front();
-    ValueRange extraFragments = ValueRange(blueprintOperands).drop_front();
-    auto blueprint = SpatBlueprintOp::create(
-      state.rewriter,
-      loc,
-      resultType,
-      input,
-      extraFragments,
-      state.rewriter.getStringAttr("nchw"),
-      state.rewriter.getStringAttr("fragmented"),
-      state.rewriter.getDenseI64ArrayAttr(flatOffsets),
-      state.rewriter.getDenseI64ArrayAttr(flatSizes),
-      state.rewriter.getStringAttr("identity"),
-      state.rewriter.getStringAttr("fragment_assembly"),
-      state.rewriter.getDenseI64ArrayAttr(fragmentOperandIndices),
-      state.rewriter.getDenseI64ArrayAttr(fragmentSourceOffsets),
-      state.rewriter.getDenseI64ArrayAttr(flatStrides),
-      state.rewriter.getStringAttr("disjoint"),
-      state.rewriter.getStringAttr("complete"));
-
-    state.hostReplacements[originalOutput] = blueprint.getOutput();
-  }
-
-  return success();
-}
-
-} // namespace onnx_mlir::spatial

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

@@ -1,11 +0,0 @@
-#pragma once
-
-#include "mlir/Support/LogicalResult.h"
-
-namespace onnx_mlir::spatial {
-
-struct MaterializerState;
-
-mlir::LogicalResult finalizeProjectedHostOutputFragments(MaterializerState& state);
-
-} // namespace onnx_mlir::spatial

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

@@ -1,17 +0,0 @@
-#pragma once
-
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/Support/LogicalResult.h"
-
-#include "Scheduling/MergeSchedule.hpp"
-
-namespace onnx_mlir {
-namespace spatial {
-
-class MergeScheduleMaterializer {
-public:
-  mlir::LogicalResult run(mlir::func::FuncOp func, const MergeScheduleResult& schedule, int64_t& nextChannelId);
-};
-
-} // namespace spatial
-} // namespace onnx_mlir

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

@@ -1,252 +0,0 @@
-#pragma once
-
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/IR/PatternMatch.h"
-#include "mlir/Transforms/FoldUtils.h"
-
-#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/DenseSet.h"
-#include "llvm/ADT/SmallVector.h"
-
-#include <optional>
-
-#include "MaterializeMergeSchedule.hpp"
-#include "MergeMessages.hpp"
-#include "MergeScheduleKeys.hpp"
-#include "ProjectedFragments.hpp"
-#include "Scheduling/ComputeInstanceUtils.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
-
-namespace onnx_mlir::spatial {
-
-struct MaterializedClass {
-  ClassId id = 0;
-  llvm::SmallVector<CpuId, 8> cpus;
-  mlir::Operation* op = nullptr;
-  mlir::Block* body = nullptr;
-  bool isBatch = false;
-
-  llvm::DenseMap<CpuId, unsigned> cpuToLane;
-  llvm::SmallVector<mlir::Value, 8> weights;
-  llvm::SmallVector<mlir::Value, 8> inputs;
-  llvm::SmallVector<mlir::Value, 4> hostOutputs;
-  llvm::DenseMap<mlir::Value, unsigned> publicationOutputToResultIndex;
-  llvm::DenseMap<mlir::Value, mlir::BlockArgument> weightArgs;
-  llvm::DenseMap<mlir::Value, mlir::BlockArgument> inputArgs;
-  llvm::DenseMap<mlir::Value, unsigned> hostOutputToResultIndex;
-};
-
-struct PackedScalarRunSlot {
-  llvm::SmallVector<ProducerKey, 8> keys;
-};
-
-enum class PackedScalarRunKind {
-  Materialized,
-  DeferredReceive,
-  DeferredLocalCompute
-};
-
-struct PackedScalarRunValue {
-  ClassId targetClass = 0;
-  mlir::Operation* sourceOp = nullptr;
-  size_t resultIndex = 0;
-  PackedScalarRunKind kind = PackedScalarRunKind::Materialized;
-
-  mlir::Value packed;
-
-  mlir::RankedTensorType fragmentType;
-  llvm::SmallVector<PackedScalarRunSlot, 8> slots;
-  MessageVector messages;
-};
-
-struct IndexedBatchRunValue {
-  ClassId targetClass = 0;
-  mlir::Operation* sourceOp = nullptr;
-  size_t resultIndex = 0;
-  mlir::Value packed;
-  mlir::RankedTensorType fragmentType;
-  llvm::SmallVector<PackedScalarRunSlot, 8> slots;
-  MessageVector messages;
-};
-
-struct LogicalSlotRange {
-  SlotId start = 0;
-  SlotId count = 0;
-};
-
-struct MaterializationRunSlot {
-  llvm::SmallVector<ComputeInstance, 8> peers;
-};
-
-using MaterializationRun = llvm::SmallVector<MaterializationRunSlot, 8>;
-
-struct OutputDestinationGroup {
-  llvm::SmallVector<size_t, 4> resultIndices;
-  llvm::SmallVector<ClassId, 4> destinationClasses;
-};
-
-struct BatchRunSendPlan {
-  size_t resultIndex = 0;
-  ClassId destinationClass = 0;
-  MessageVector messages;
-};
-
-enum class TensorDemandActionKind {
-  DestinationFanout,
-  SameClassIndexedFragment,
-  TerminalBlueprintPublication,
-  WholeTensorBarrier
-};
-
-enum class WholeTensorBarrierReason {
-  FunctionReturnWithoutBlueprint,
-  DenseLogicalConsumer
-};
-
-struct TensorDemandAction {
-  TensorDemandActionKind kind = TensorDemandActionKind::DestinationFanout;
-  std::optional<ClassId> destinationClass;
-  std::optional<WholeTensorBarrierReason> barrierReason;
-};
-
-struct RunOutputDemand {
-  size_t resultIndex = 0;
-  mlir::Value originalOutput;
-  mlir::RankedTensorType fragmentType;
-  llvm::SmallVector<TensorDemandAction, 4> actions;
-};
-
-struct CompactRunPlan {
-  llvm::SmallVector<RunOutputDemand, 4> outputs;
-};
-
-enum class BatchInputDemandKind {
-  LaneFragment,
-  ProjectedFragment,
-  WholeTensorBarrier
-};
-
-struct BatchInputDemand {
-  BatchInputDemandKind kind = BatchInputDemandKind::LaneFragment;
-  std::optional<ProducerKey> wholeTensorProducer;
-};
-
-struct CloneIndexingContext {
-  std::optional<mlir::Value> runSlotIndex;
-  std::optional<mlir::Value> projectionSlotIndex;
-};
-
-struct MaterializerState;
-
-class AvailableValueStore {
-public:
-  struct ExactBatchFragmentRecord {
-    ProducerKey key;
-    mlir::Value value;
-  };
-
-  void record(ProducerKey key, ClassId classId, mlir::Value value) {
-    exactValues[key][classId] = value;
-
-    auto batch = mlir::dyn_cast_or_null<SpatComputeBatch>(key.instance.op);
-    if (!batch || key.instance.laneCount == 0)
-      return;
-
-    WholeBatchAssemblyLookupKey lookupKey {batch.getOperation(), key.resultIndex, classId};
-    llvm::SmallVector<ExactBatchFragmentRecord, 16>& bucket = exactBatchFragmentsByProducerResultClass[lookupKey];
-    for (ExactBatchFragmentRecord& record : bucket) {
-      if (!(record.key == key))
-        continue;
-      record.value = value;
-      return;
-    }
-    bucket.push_back({key, value});
-  }
-
-  void recordPackedRun(PackedScalarRunValue run) {
-    size_t runIndex = packedScalarRuns.size();
-    packedScalarRuns.push_back(std::move(run));
-    const PackedScalarRunValue& storedRun = packedScalarRuns[runIndex];
-    WholeBatchAssemblyLookupKey lookupKey {storedRun.sourceOp, storedRun.resultIndex, storedRun.targetClass};
-    packedRunsByProducerResultClass[lookupKey].push_back(runIndex);
-  }
-
-  void recordIndexedBatchRun(IndexedBatchRunValue run) { indexedBatchRuns.push_back(std::move(run)); }
-
-  std::optional<mlir::Value> lookupExact(ProducerKey key, ClassId classId) const;
-  std::optional<mlir::Value> lookup(MaterializerState& state, ProducerKey key, ClassId classId);
-  IndexedBatchRunValue* lookupIndexedBatchRun(ProducerKey key, ClassId classId);
-
-  llvm::ArrayRef<size_t> getPackedRunIndicesForWholeBatch(WholeBatchAssemblyLookupKey key) const {
-    auto it = packedRunsByProducerResultClass.find(key);
-    if (it == packedRunsByProducerResultClass.end())
-      return {};
-    return it->second;
-  }
-
-  llvm::ArrayRef<ExactBatchFragmentRecord> getExactFragmentsForWholeBatch(WholeBatchAssemblyLookupKey key) const {
-    auto it = exactBatchFragmentsByProducerResultClass.find(key);
-    if (it == exactBatchFragmentsByProducerResultClass.end())
-      return {};
-    return it->second;
-  }
-
-  PackedScalarRunValue& getPackedRun(size_t index) { return packedScalarRuns[index]; }
-
-private:
-  std::optional<mlir::Value> lookupPackedRun(MaterializerState& state, ProducerKey key, ClassId classId);
-
-  llvm::DenseMap<ProducerKey, llvm::DenseMap<ClassId, mlir::Value>, ProducerKeyInfo> exactValues;
-  llvm::SmallVector<PackedScalarRunValue, 8> packedScalarRuns;
-  llvm::SmallVector<IndexedBatchRunValue, 8> indexedBatchRuns;
-  llvm::DenseMap<WholeBatchAssemblyLookupKey,
-                 llvm::SmallVector<ExactBatchFragmentRecord, 16>,
-                 WholeBatchAssemblyLookupKeyInfo>
-    exactBatchFragmentsByProducerResultClass;
-  llvm::DenseMap<WholeBatchAssemblyLookupKey, llvm::SmallVector<size_t, 16>, WholeBatchAssemblyLookupKeyInfo>
-    packedRunsByProducerResultClass;
-};
-
-struct MaterializerState {
-  mlir::func::FuncOp func;
-  const MergeScheduleResult& schedule;
-  mlir::IRRewriter rewriter;
-  mlir::OperationFolder constantFolder;
-  int64_t& nextChannelId;
-  llvm::SmallVector<MaterializedClass, 8> classes;
-  llvm::DenseMap<CpuId, ClassId> cpuToClass;
-  llvm::DenseMap<CpuId, llvm::SmallVector<ComputeInstance, 32>> logicalInstancesByCpu;
-  llvm::DenseMap<ComputeInstance, LogicalSlotRange> scheduledInstanceToLogicalSlots;
-  llvm::DenseMap<ComputeInstance, ComputeInstance> logicalInstanceToScheduledChunk;
-  llvm::DenseSet<ClassSlotKey> materializedLogicalSlots;
-
-  llvm::DenseMap<ProducerKey, llvm::SmallVector<ClassId, 4>, ProducerKeyInfo> producerDestClasses;
-  llvm::DenseMap<SameClassConsumerLookupKey, llvm::SmallVector<ProducerKey, 4>, SameClassConsumerLookupKeyInfo>
-    sameClassConsumerIndex;
-  llvm::DenseMap<ProjectedBatchInputKey, AffineProjectedInputSliceMatch, ProjectedBatchInputKeyInfo>
-    projectedInputMatches;
-  llvm::DenseSet<ProjectedBatchInputKey, ProjectedBatchInputKeyInfo> nonProjectedInputs;
-  llvm::DenseMap<mlir::Value, bool> liveExternalUseCache;
-  llvm::DenseMap<mlir::Operation*, llvm::SmallVector<mlir::Type, 4>> batchOutputFragmentTypesCache;
-  llvm::DenseMap<ComputeInstance, llvm::SmallVector<mlir::Value, 4>, llvm::DenseMapInfo<ComputeInstance>>
-    computeInstanceOutputsCache;
-  llvm::DenseMap<ProducerKey, llvm::DenseMap<ClassId, ProjectedTransferDescriptor>, ProducerKeyInfo>
-    projectedTransfers;
-  llvm::DenseMap<mlir::Operation*, llvm::DenseMap<ClassId, ProjectedExtractReplacement>>
-    projectedExtractReplacements;
-  AvailableValueStore availableValues;
-  llvm::DenseMap<mlir::Value, mlir::Value> hostReplacements;
-  llvm::DenseMap<mlir::Value, ClassId> hostOutputOwners;
-  llvm::SmallVector<PendingProjectedHostOutputFragment, 32> pendingProjectedHostOutputFragments;
-  llvm::DenseSet<mlir::Operation*> oldComputeOps;
-
-  MaterializerState(mlir::func::FuncOp func, const MergeScheduleResult& schedule, int64_t& nextChannelId)
-  : func(func),
-    schedule(schedule),
-    rewriter(func.getContext()),
-    constantFolder(func.getContext()),
-    nextChannelId(nextChannelId) {}
-};
-
-} // namespace onnx_mlir::spatial

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

@@ -1,391 +1,112 @@
-#include "mlir/Analysis/TopologicalSortUtils.h"
-#include "mlir/Dialect/Func/IR/FuncOps.h"
-#include "mlir/IR/IRMapping.h"
-#include "mlir/IR/Location.h"
-#include "mlir/IR/PatternMatch.h"
-#include "mlir/IR/Region.h"
-#include "mlir/IR/Value.h"
-#include "mlir/IR/ValueRange.h"
+#include "ScheduledComputeMaterialization.hpp"
+#include "ScheduledComputeReport.hpp"
+#include "ScheduledComputeVerification.hpp"
+#include "DeferredCommunicationRealization.hpp"
+#include "DeferredCommunicationDeadlock.hpp"
+
 #include "mlir/Pass/Pass.h"
-#include "mlir/Support/LLVM.h"
 
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/STLExtras.h"
-#include "llvm/ADT/SmallSet.h"
-#include "llvm/ADT/SmallVector.h"
-#include "llvm/Support/raw_os_ostream.h"
-
-#include <algorithm>
-#include <cstddef>
-#include <cstdint>
-#include <fstream>
-#include <memory>
-#include <optional>
-#include <utility>
-#include <vector>
-
-#include "MaterializeMergeSchedule.hpp"
-#include "Scheduling/ComputeGraph.hpp"
-#include "Scheduling/ComputeInstanceUtils.hpp"
 #include "Scheduling/MergeSchedulingAnalysis.hpp"
-#include "src/Accelerators/PIM/Common/IR/BatchCoreUtils.hpp"
-#include "src/Accelerators/PIM/Common/IR/CompactAsmUtils.hpp"
-#include "src/Accelerators/PIM/Common/PimCommon.hpp"
-#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
-#include "src/Accelerators/PIM/Common/Support/ReportUtils.hpp"
 #include "src/Accelerators/PIM/Conversion/ONNXToSpatial/ONNXToSpatialVerifier.hpp"
-#include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
+#include "src/Accelerators/PIM/Common/PimCommon.hpp"
+#include "src/Accelerators/PIM/Common/Support/DebugDump.hpp"
+#include "src/Accelerators/PIM/Pass/PIMPasses.h"
 
 using namespace mlir;
 
 namespace onnx_mlir {
+namespace spatial {
 namespace {
-using namespace onnx_mlir::compact_asm;
-using SpatCompute = spatial::SpatGraphCompute;
-using SpatComputeBatch = spatial::SpatGraphComputeBatch;
 
-bool isTrivialSerialMergeCandidate(SpatCompute compute) {
-  if (!compute->hasOneUse())
-    return false;
-
-  auto& use = *compute->getUses().begin();
-  auto user = dyn_cast<SpatCompute>(use.getOwner());
-  return user && user.getInputs().size() == 1 && use.getOperandNumber() >= user.getWeights().size();
-}
-
-SmallVector<size_t> appendMissingWeightsAndBuildIndexMap(SmallVectorImpl<Value>& targetWeights,
-                                                         ValueRange sourceWeights) {
-  DenseMap<Value, SmallVector<size_t, 4>> targetWeightIndices;
-  for (auto [weightIndex, weight] : llvm::enumerate(targetWeights))
-    targetWeightIndices[weight].push_back(weightIndex);
-
-  DenseMap<Value, size_t> usedSourceWeightOccurrences;
-  SmallVector<size_t> sourceToTargetIndex;
-  sourceToTargetIndex.reserve(sourceWeights.size());
-  for (Value weight : sourceWeights) {
-    size_t occurrence = usedSourceWeightOccurrences[weight]++;
-    auto& matchingIndices = targetWeightIndices[weight];
-    if (occurrence >= matchingIndices.size()) {
-      size_t newIndex = targetWeights.size();
-      targetWeights.push_back(weight);
-      matchingIndices.push_back(newIndex);
-      sourceToTargetIndex.push_back(newIndex);
-      continue;
-    }
-    sourceToTargetIndex.push_back(matchingIndices[occurrence]);
-  }
-  return sourceToTargetIndex;
-}
-
-void mergeTriviallyConnectedComputes(func::FuncOp funcOp) {
-  Location loc = funcOp.getLoc();
-  IRRewriter rewriter(funcOp->getContext());
-  SmallVector<SpatCompute> trivialComputes;
-  llvm::SmallSet<SpatCompute, 8> toErase;
-
-  for (auto compute : funcOp.getOps<SpatCompute>())
-    if (isTrivialSerialMergeCandidate(compute))
-      trivialComputes.push_back(compute);
-
-  while (!trivialComputes.empty()) {
-    auto compute = trivialComputes.front();
-
-    if (compute.use_empty()) {
-      std::swap(trivialComputes.front(), trivialComputes.back());
-      trivialComputes.pop_back();
-      continue;
-    }
-
-    auto& computeUse = *compute->getUses().begin();
-    auto child = cast<SpatCompute>(computeUse.getOwner());
-    auto usedResult = cast<OpResult>(computeUse.get()).getResultNumber();
-    auto childInputIndex = computeUse.getOperandNumber() - child.getWeights().size();
-
-    rewriter.setInsertionPointAfter(compute.getOperation());
-    SmallVector<Value> mergedWeights(compute.getWeights().begin(), compute.getWeights().end());
-    SmallVector<size_t> childWeightToNewIndex = appendMissingWeightsAndBuildIndexMap(mergedWeights, child.getWeights());
-    SmallVector<Value> mergedInputs(compute.getInputs().begin(), compute.getInputs().end());
-    auto newCompute = SpatCompute::create(rewriter, loc, child.getResultTypes(), mergedWeights, mergedInputs);
-    Block* newBody = rewriter.createBlock(&newCompute.getBodyRegion());
-    for (Value weight : mergedWeights)
-      newBody->addArgument(weight.getType(), loc);
-    for (Value input : mergedInputs)
-      newBody->addArgument(input.getType(), loc);
-
-    IRMapping mapper;
-    for (auto [weightIndex, _] : llvm::enumerate(compute.getWeights())) {
-      auto oldWeightArg = compute.getWeightArgument(weightIndex);
-      auto newWeightArg = newCompute.getWeightArgument(weightIndex);
-      assert(oldWeightArg && newWeightArg && "expected compute weight block arguments");
-      mapper.map(*oldWeightArg, *newWeightArg);
-    }
-    for (auto [inputIndex, _] : llvm::enumerate(compute.getInputs())) {
-      auto oldInputArg = compute.getInputArgument(inputIndex);
-      auto newInputArg = newCompute.getInputArgument(inputIndex);
-      assert(oldInputArg && newInputArg && "expected compute input block arguments");
-      mapper.map(*oldInputArg, *newInputArg);
-    }
-    for (auto [oldIndex, weight] : llvm::enumerate(child.getWeights())) {
-      auto oldWeightArg = child.getWeightArgument(oldIndex);
-      auto newWeightArg = newCompute.getWeightArgument(childWeightToNewIndex[oldIndex]);
-      assert(oldWeightArg && newWeightArg && "expected child compute weight block arguments");
-      mapper.map(*oldWeightArg, *newWeightArg);
-    }
-
-    rewriter.setInsertionPointToEnd(newBody);
-    auto computeYield = cast<spatial::SpatYieldOp>(compute.getBody().front().getTerminator());
-    for (Operation& op : compute.getBody().front().without_terminator())
-      rewriter.clone(op, mapper);
-    auto childInputArg = child.getInputArgument(childInputIndex);
-    assert(childInputArg && "expected child compute input block argument");
-    mapper.map(*childInputArg, mapper.lookupOrDefault(computeYield.getOperand(usedResult)));
-
-    rewriter.setInsertionPointToEnd(newBody);
-    for (auto& op : child.getBody().front())
-      rewriter.clone(op, mapper);
-
-    child.replaceAllUsesWith(newCompute);
-    toErase.insert(child);
-
-    std::swap(trivialComputes.front(), trivialComputes.back());
-    trivialComputes.pop_back();
-    toErase.insert(compute);
-
-    if (isTrivialSerialMergeCandidate(newCompute))
-      trivialComputes.push_back(newCompute);
-  }
-
-  for (auto compute : toErase) {
-    for (Value result : compute->getResults())
-      result.dropAllUses();
-    compute.erase();
-  }
-}
-
-void generateReport(func::FuncOp funcOp, const std::string& name, size_t usedCpuCount = 0) {
-  std::fstream file = openReportFile(name);
-  if (!file.is_open())
-    return;
-  llvm::raw_os_ostream os(file);
-
-  struct ReportRow {
-    uint64_t id = 0;
-    uint64_t logicalComputeCount = 0;
-    uint64_t crossbarCount = 0;
-    uint64_t instructionCount = 0;
-    bool isRebatched = false;
-    SmallVector<int32_t> coreIds;
-  };
-
-  //TODO Used for report refactor
-  struct CollectorConcatRow {
-    uint64_t computeId = 0;
-    int32_t coreId = -1;
-    uint64_t operandCount = 0;
-  };
-
-  uint64_t totalComputeOps = 0;
-  uint64_t totalLogicalComputes = 0;
-  uint64_t totalBatchComputeOps = 0;
-  uint64_t totalInstructionCount = 0;
-  uint64_t totalCrossbarCount = 0;
-  uint64_t nextBatchId = 0;
-  //TODO Used for report refactor
-  std::vector<ReportRow> collectedData;
-  //TODO Used for report refactor
-  std::vector<CollectorConcatRow> collectorConcatRows;
-
-  auto getPerInstanceCrossbarCount = [&](Operation* op) -> uint64_t {
-    return static_cast<uint64_t>(spatial::collectDistinctCrossbarWeights(op).size());
-  };
-
-  for (Operation& op : funcOp.getBody().front()) {
-    if (auto spatCompute = dyn_cast<spatial::SpatScheduledCompute>(&op)) {
-      uint64_t numInst = spatial::countComputeBodyInstructions(spatCompute.getBody());
-      uint64_t perInstanceCrossbarCount = getPerInstanceCrossbarCount(spatCompute.getOperation());
-      SmallVector<int32_t> coreIds;
-      auto coreId = getOptionalScheduledCoreId(spatCompute, "merge compute core id");
-      if (failed(coreId))
-        return;
-      if (*coreId)
-        coreIds.push_back(**coreId);
-      uint64_t computeId = totalComputeOps++;
-      collectedData.push_back({computeId, 1, perInstanceCrossbarCount, numInst, false, coreIds});
-      uint64_t maxConcatOperands = 0;
-      spatCompute.getBody().walk([&](spatial::SpatConcatOp concatOp) {
-        maxConcatOperands = std::max<uint64_t>(maxConcatOperands, concatOp.getInputs().size());
-      });
-      //TODO 128 is a magic number
-      if (maxConcatOperands >= 128 && !coreIds.empty())
-        collectorConcatRows.push_back({computeId, coreIds.front(), maxConcatOperands});
-      totalLogicalComputes += 1;
-      totalInstructionCount += numInst;
-      totalCrossbarCount += perInstanceCrossbarCount;
-      continue;
-    }
-    if (auto batch = dyn_cast<spatial::SpatScheduledComputeBatch>(&op)) {
-      uint64_t numInst = spatial::countComputeBodyInstructions(batch.getBody());
-      uint64_t logicalCount = static_cast<uint64_t>(batch.getLaneCount());
-      uint64_t perInstanceCrossbarCount = getPerInstanceCrossbarCount(batch.getOperation());
-      SmallVector<int32_t> coreIds;
-      auto optionalCoreIds = getOptionalScheduledBatchCoreIds(batch, "merge compute_batch core id");
-      if (failed(optionalCoreIds))
-        return;
-      if (*optionalCoreIds)
-        coreIds = std::move(**optionalCoreIds);
-      collectedData.push_back(
-        {nextBatchId++, logicalCount, perInstanceCrossbarCount * logicalCount, numInst, true, coreIds});
-      totalComputeOps += 1;
-      totalLogicalComputes += logicalCount;
-      totalBatchComputeOps += 1;
-      totalInstructionCount += numInst * logicalCount;
-      totalCrossbarCount += perInstanceCrossbarCount * logicalCount;
-    }
-  }
-
-  llvm::SmallVector<ReportField, 6> totalFields = {
-    {"Used cores",                            std::to_string(usedCpuCount)         },
-    {"Number of top-level compute ops",       std::to_string(totalComputeOps)      },
-    {"Number of logical computes",            std::to_string(totalLogicalComputes) },
-    {"Number of top-level batch compute ops", std::to_string(totalBatchComputeOps) },
-    {"Number of instructions",                std::to_string(totalInstructionCount)},
-    {"Number of used crossbars",              std::to_string(totalCrossbarCount)   }
-  };
-  printReportTotalsBlock(os, totalFields);
-  if (!collectedData.empty() || !collectorConcatRows.empty())
-    os << "\n";
-
-  if (!collectorConcatRows.empty()) {
-    os << "Collector concat materialization:\n";
-    for (const CollectorConcatRow& row : collectorConcatRows)
-      os << "\tmaterialization_kind = single_collector_concat, compute = " << row.computeId
-         << ", concat_operand_count = " << row.operandCount << ", collector_core = " << row.coreId << "\n";
-    os << "\n";
-  }
-
-  sortReportEntriesByFirstCore(collectedData);
-
-  for (uint64_t cI = 0; cI < totalComputeOps; ++cI) {
-    uint64_t lastIndex = cI;
-    ReportRow current = collectedData[cI];
-
-    for (uint64_t nI = cI + 1; nI < totalComputeOps; ++nI) {
-      ReportRow next = collectedData[nI];
-      if (current.isRebatched == next.isRebatched && current.crossbarCount == next.crossbarCount
-          && current.instructionCount == next.instructionCount
-          && current.logicalComputeCount == next.logicalComputeCount)
-        lastIndex = nI;
-      else
-        break;
-    }
-
-    if (current.isRebatched) {
-      os << "Batch ";
-      for (uint64_t index = cI; index <= lastIndex; ++index) {
-        if (index != cI)
-          os << ",\n      ";
-        os << collectedData[index].id << " (cores ";
-        if (collectedData[index].coreIds.empty())
-          os << "unknown";
-        else
-          printCompressedIntegerEntries(os, ArrayRef<int32_t>(collectedData[index].coreIds));
-        os << ")";
-      }
-    }
-    else {
-      os << "Compute ";
-      SmallVector<uint64_t> opIds;
-      opIds.reserve(lastIndex - cI + 1);
-      for (uint64_t index = cI; index <= lastIndex; ++index)
-        opIds.push_back(collectedData[index].id);
-      printCompressedIntegerEntries(os, ArrayRef<uint64_t>(opIds));
-    }
-
-    os << ":\n";
-    uint64_t perCoreLogicalComputeCount = current.isRebatched ? 1 : current.logicalComputeCount;
-    uint64_t perCoreInstructionCount = current.instructionCount;
-    uint64_t perCoreCrossbarCount =
-      current.logicalComputeCount == 0 ? 0 : current.crossbarCount / current.logicalComputeCount;
-    uint64_t totalEntryInstructionCount = current.instructionCount * current.logicalComputeCount;
-
-    llvm::SmallVector<ReportField, 3> perCoreFields = {
-      {"Number of logical computes", std::to_string(perCoreLogicalComputeCount)},
-      {"Number of instructions",     std::to_string(perCoreInstructionCount)   },
-      {"Number of used crossbars",   std::to_string(perCoreCrossbarCount)      }
-    };
-    if (current.isRebatched) {
-      llvm::SmallVector<ReportField, 3> totalEntryFields = {
-        {"Number of logical computes", std::to_string(current.logicalComputeCount)},
-        {"Number of instructions",     std::to_string(totalEntryInstructionCount) },
-        {"Number of used crossbars",   std::to_string(current.crossbarCount)      }
-      };
-      printReportPerCoreAndTotalFields(os, perCoreFields, totalEntryFields);
-    }
-    else {
-      printReportFlatFields(os, perCoreFields);
-    }
-    printReportEntrySeparator(os, lastIndex + 1 < totalComputeOps);
-    cI = lastIndex;
-  }
-
-  os.flush();
-  file.close();
-}
-
-struct MergeComputeNodesPass : PassWrapper<MergeComputeNodesPass, OperationPass<func::FuncOp>> {
-
-private:
-  int64_t nextChannelId = 0;
-
-public:
+struct MergeComputeNodesPass final : PassWrapper<MergeComputeNodesPass, OperationPass<ModuleOp>> {
   MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MergeComputeNodesPass)
 
-  StringRef getArgument() const override { return "pim-merge-compute-nodes-pass"; }
+  StringRef getArgument() const override { return "pim-merge-compute-nodes"; }
   StringRef getDescription() const override {
-    return "Merge Spatial-Compute-Nodes in order to reduce the total "
-           "execution time";
+    return "Materialize scheduled Spatial compute with deferred communication placeholders.";
   }
 
-  LogicalResult initialize(MLIRContext* context) override { return success(); }
-
   void runOnOperation() override {
-    func::FuncOp func = getOperation();
-    if (failed(verifyLogicalSpatialGraphInvariants(func))) {
-      func.emitOpError("logical Spatial graph verification failed at the start of MergeComputeNodes");
-      signalPassFailure();
-      return;
-    }
-    mergeTriviallyConnectedComputes(func);
-    if (failed(verifyLogicalSpatialGraphInvariants(func))) {
-      func.emitOpError("logical Spatial graph verification failed after trivial merge simplification");
+    ModuleOp moduleOp = getOperation();
+    auto entryFunc = getPimEntryFunc(moduleOp);
+    if (failed(entryFunc)) {
+      moduleOp.emitError("failed to locate the PIM entry function during MergeComputeNodes");
       signalPassFailure();
       return;
     }
 
-    const spatial::MergeScheduleResult* analysisResult = nullptr;
-    analysisResult = &getAnalysis<spatial::MergeSchedulingAnalysis>().getResult();
-    if (failed(spatial::MergeScheduleMaterializer().run(func, *analysisResult, nextChannelId))) {
+    func::FuncOp funcOp = *entryFunc;
+    MergeScheduleResult schedule = MergeSchedulingAnalysis(funcOp).getResult();
+    PatternRewriter rewriter(moduleOp.getContext());
+    FailureOr<ScheduledComputeMaterializationResult> materialization =
+      materializeScheduledCompute(funcOp, schedule, rewriter);
+    if (failed(materialization)) {
       signalPassFailure();
       return;
     }
 
-    if (!sortTopologically(&func.getBody().front())) {
-      func.emitOpError("failed to topologically order merged Spatial IR");
+    // Phase 1 is intentionally dumped before its verifier: malformed deferred
+    // payloads must be diagnosed from the producer-owned body.
+    dumpModule(moduleOp, "spatial2_scheduled_no_comm", /*assumeVerified=*/true);
+
+    if (failed(verifyMaterializedScheduleMapping(funcOp,
+                                                 schedule,
+                                                 materialization->peftClassPlans,
+                                                 materialization->graphComputeToBlockMap,
+                                                 materialization->materializedSchedules))) {
+      moduleOp.emitError("scheduled Spatial materialization mapping verification failed");
       signalPassFailure();
       return;
     }
-    if (failed(verifyScheduledSpatialInvariants(func))) {
-      func.emitOpError("scheduled Spatial verification failed after merge materialization");
+    if (failed(verifyDeferredTransferPhase1Invariants(funcOp))) {
+      moduleOp.emitError("scheduled Spatial deferred communication verification failed");
       signalPassFailure();
       return;
     }
-    dumpModule(cast<ModuleOp>(func->getParentOp()), "spatial1_merged");
-    generateReport(func, "spatial_merge_report", analysisResult->cpuToLastComputeMap.size());
+    if (failed(verifyScheduledMaterializationRecords(materialization->materializedSchedules))) {
+      moduleOp.emitError("scheduled Spatial materialization record verification failed");
+      signalPassFailure();
+      return;
+    }
+    if (failed(verifyPeftMaterializationReportSummary(funcOp,
+                                                      schedule,
+                                                      materialization->peftClassPlans,
+                                                      materialization->materializedSchedules))) {
+      moduleOp.emitError("scheduled Spatial report verification failed");
+      signalPassFailure();
+      return;
+    }
+    if (failed(verifyScheduledSpatialInvariants(funcOp))) {
+      moduleOp.emitError("scheduled Spatial phase 1 verification failed");
+      signalPassFailure();
+      return;
+    }
+
+    dumpScheduledComputeReportAndModule(moduleOp,
+                                        funcOp,
+                                        schedule,
+                                        materialization->peftClassPlans,
+                                        materialization->materializedSchedules);
+    if (failed(realizeDeferredCommunication(funcOp))) {
+      moduleOp.emitError("MergeComputeNodes phase 2 communication realization failed");
+      signalPassFailure();
+      return;
+    }
+    if (failed(verifyRealizedCommunicationDeadlockFree(funcOp))) {
+      moduleOp.emitError("MergeComputeNodes final communication verification failed");
+      signalPassFailure();
+      return;
+    }
+    if (failed(verifyScheduledSpatialInvariants(funcOp))) {
+      moduleOp.emitError("scheduled Spatial phase 2 verification failed");
+      signalPassFailure();
+    }
   }
 };
 
 } // namespace
+} // namespace spatial
 
-std::unique_ptr<Pass> createMergeComputeNodesPass() { return std::make_unique<MergeComputeNodesPass>(); }
+std::unique_ptr<Pass> createMergeComputeNodesPass() { return std::make_unique<spatial::MergeComputeNodesPass>(); }
 
 } // namespace onnx_mlir

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

@@ -1,67 +0,0 @@
-#pragma once
-
-#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/SmallVector.h"
-#include "llvm/ADT/StringRef.h"
-
-#include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
-
-namespace onnx_mlir::spatial {
-
-using CpuId = size_t;
-
-inline mlir::FailureOr<int32_t> getCheckedCoreId(mlir::Operation* anchor, CpuId cpu, llvm::StringRef fieldName) {
-  return pim::checkedI32(static_cast<uint64_t>(cpu), anchor, fieldName);
-}
-
-inline mlir::FailureOr<llvm::SmallVector<int32_t, 8>>
-getCheckedCoreIds(mlir::Operation* anchor, llvm::ArrayRef<CpuId> cpus, llvm::StringRef fieldName) {
-  llvm::SmallVector<int32_t, 8> coreIds;
-  coreIds.reserve(cpus.size());
-  for (CpuId cpu : cpus) {
-    auto checkedCoreId = getCheckedCoreId(anchor, cpu, fieldName);
-    if (mlir::failed(checkedCoreId))
-      return mlir::failure();
-    coreIds.push_back(*checkedCoreId);
-  }
-  return coreIds;
-}
-
-struct MessageVector {
-  llvm::SmallVector<int64_t, 16> channelIds;
-  llvm::SmallVector<int32_t, 16> sourceCoreIds;
-  llvm::SmallVector<int32_t, 16> targetCoreIds;
-
-  size_t size() const { return channelIds.size(); }
-  bool empty() const { return channelIds.empty(); }
-
-  mlir::LogicalResult verify(mlir::Operation* anchor) const {
-    if (channelIds.size() != sourceCoreIds.size() || channelIds.size() != targetCoreIds.size())
-      return anchor->emitError("message metadata is inconsistent");
-    return mlir::success();
-  }
-
-  void append(int64_t channelId, int32_t sourceCoreId, int32_t targetCoreId) {
-    channelIds.push_back(channelId);
-    sourceCoreIds.push_back(sourceCoreId);
-    targetCoreIds.push_back(targetCoreId);
-  }
-
-  void append(llvm::ArrayRef<int64_t> channels, llvm::ArrayRef<int32_t> sources, llvm::ArrayRef<int32_t> targets) {
-    assert(channels.size() == sources.size() && "channel/source count mismatch");
-    assert(channels.size() == targets.size() && "channel/target count mismatch");
-    llvm::append_range(channelIds, channels);
-    llvm::append_range(sourceCoreIds, sources);
-    llvm::append_range(targetCoreIds, targets);
-  }
-
-  MessageVector slice(size_t offset, size_t count) const {
-    MessageVector result;
-    result.append(llvm::ArrayRef<int64_t>(channelIds).slice(offset, count),
-                  llvm::ArrayRef<int32_t>(sourceCoreIds).slice(offset, count),
-                  llvm::ArrayRef<int32_t>(targetCoreIds).slice(offset, count));
-    return result;
-  }
-};
-
-} // namespace onnx_mlir::spatial

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

@@ -1,134 +0,0 @@
-#pragma once
-
-#include "llvm/ADT/DenseMap.h"
-#include "llvm/ADT/SmallVector.h"
-
-#include <cstddef>
-#include <cstdint>
-#include <limits>
-#include <utility>
-
-#include "Scheduling/ComputeInstanceUtils.hpp"
-
-namespace onnx_mlir::spatial {
-
-using ClassId = size_t;
-using SlotId = size_t;
-
-struct ProducerKey {
-  ComputeInstance instance;
-  size_t resultIndex = 0;
-
-  bool operator==(const ProducerKey& other) const {
-    return instance == other.instance && resultIndex == other.resultIndex;
-  }
-};
-
-struct ProducerKeyInfo {
-  static ProducerKey getEmptyKey() {
-    return {llvm::DenseMapInfo<ComputeInstance>::getEmptyKey(), std::numeric_limits<size_t>::max()};
-  }
-
-  static ProducerKey getTombstoneKey() {
-    return {llvm::DenseMapInfo<ComputeInstance>::getTombstoneKey(), std::numeric_limits<size_t>::max()};
-  }
-
-  static unsigned getHashValue(const ProducerKey& key) {
-    return llvm::hash_combine(llvm::DenseMapInfo<ComputeInstance>::getHashValue(key.instance), key.resultIndex);
-  }
-
-  static bool isEqual(const ProducerKey& lhs, const ProducerKey& rhs) { return lhs == rhs; }
-};
-
-struct SameClassConsumerLookupKey {
-  mlir::Operation* sourceOp = nullptr;
-  size_t resultIndex = 0;
-  ClassId classId = 0;
-
-  bool operator==(const SameClassConsumerLookupKey& other) const {
-    return sourceOp == other.sourceOp && resultIndex == other.resultIndex && classId == other.classId;
-  }
-};
-
-struct SameClassConsumerLookupKeyInfo {
-  static SameClassConsumerLookupKey getEmptyKey() {
-    return {llvm::DenseMapInfo<mlir::Operation*>::getEmptyKey(), std::numeric_limits<size_t>::max(),
-            std::numeric_limits<ClassId>::max()};
-  }
-
-  static SameClassConsumerLookupKey getTombstoneKey() {
-    return {llvm::DenseMapInfo<mlir::Operation*>::getTombstoneKey(), std::numeric_limits<size_t>::max(),
-            std::numeric_limits<ClassId>::max()};
-  }
-
-  static unsigned getHashValue(const SameClassConsumerLookupKey& key) {
-    return llvm::hash_combine(llvm::DenseMapInfo<mlir::Operation*>::getHashValue(key.sourceOp),
-                              key.resultIndex,
-                              key.classId);
-  }
-
-  static bool isEqual(const SameClassConsumerLookupKey& lhs, const SameClassConsumerLookupKey& rhs) {
-    return lhs == rhs;
-  }
-};
-
-struct WholeBatchAssemblyLookupKey {
-  mlir::Operation* sourceOp = nullptr;
-  size_t resultIndex = 0;
-  ClassId classId = 0;
-
-  bool operator==(const WholeBatchAssemblyLookupKey& other) const {
-    return sourceOp == other.sourceOp && resultIndex == other.resultIndex && classId == other.classId;
-  }
-};
-
-struct WholeBatchAssemblyLookupKeyInfo {
-  static WholeBatchAssemblyLookupKey getEmptyKey() {
-    return {llvm::DenseMapInfo<mlir::Operation*>::getEmptyKey(), std::numeric_limits<size_t>::max(),
-            std::numeric_limits<ClassId>::max()};
-  }
-
-  static WholeBatchAssemblyLookupKey getTombstoneKey() {
-    return {llvm::DenseMapInfo<mlir::Operation*>::getTombstoneKey(), std::numeric_limits<size_t>::max(),
-            std::numeric_limits<ClassId>::max()};
-  }
-
-  static unsigned getHashValue(const WholeBatchAssemblyLookupKey& key) {
-    return llvm::hash_combine(llvm::DenseMapInfo<mlir::Operation*>::getHashValue(key.sourceOp),
-                              key.resultIndex,
-                              key.classId);
-  }
-
-  static bool isEqual(const WholeBatchAssemblyLookupKey& lhs, const WholeBatchAssemblyLookupKey& rhs) {
-    return lhs == rhs;
-  }
-};
-
-using ClassSlotKey = std::pair<ClassId, SlotId>;
-
-struct ProjectedBatchInputKey {
-  mlir::Operation* consumerOp = nullptr;
-  unsigned inputIndex = 0;
-
-  bool operator==(const ProjectedBatchInputKey& other) const {
-    return consumerOp == other.consumerOp && inputIndex == other.inputIndex;
-  }
-};
-
-struct ProjectedBatchInputKeyInfo {
-  static ProjectedBatchInputKey getEmptyKey() {
-    return {llvm::DenseMapInfo<mlir::Operation*>::getEmptyKey(), std::numeric_limits<unsigned>::max()};
-  }
-
-  static ProjectedBatchInputKey getTombstoneKey() {
-    return {llvm::DenseMapInfo<mlir::Operation*>::getTombstoneKey(), std::numeric_limits<unsigned>::max()};
-  }
-
-  static unsigned getHashValue(const ProjectedBatchInputKey& key) {
-    return llvm::hash_combine(key.consumerOp, key.inputIndex);
-  }
-
-  static bool isEqual(const ProjectedBatchInputKey& lhs, const ProjectedBatchInputKey& rhs) { return lhs == rhs; }
-};
-
-} // namespace onnx_mlir::spatial

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

@@ -1,104 +0,0 @@
-#include "ProjectedFragments.hpp"
-
-#include "mlir/IR/BuiltinTypes.h"
-
-namespace onnx_mlir::spatial {
-
-static mlir::FailureOr<mlir::RankedTensorType> getPackedBatchTensorType(mlir::Type laneType, size_t laneCount) {
-  auto tensorType = mlir::dyn_cast<mlir::RankedTensorType>(laneType);
-  if (!tensorType || !tensorType.hasStaticShape() || tensorType.getRank() == 0)
-    return mlir::failure();
-
-  llvm::SmallVector<int64_t, 4> shape(tensorType.getShape());
-  shape[0] *= static_cast<int64_t>(laneCount);
-  return mlir::RankedTensorType::get(shape, tensorType.getElementType(), tensorType.getEncoding());
-}
-
-unsigned getProjectedFragmentsPerLogicalSlot(llvm::ArrayRef<int64_t> loopTripCounts) {
-  unsigned fragmentsPerLogicalSlot = 1;
-  for (int64_t tripCount : loopTripCounts) {
-    assert(tripCount > 0 && "projected loop trip counts must be positive");
-    fragmentsPerLogicalSlot *= static_cast<unsigned>(tripCount);
-  }
-  return fragmentsPerLogicalSlot;
-}
-
-mlir::LogicalResult verifyProjectedFragmentLayout(mlir::Operation* anchor, const ProjectedFragmentLayout& layout) {
-  if (!layout.fragmentType || layout.fragmentShape.empty())
-    return anchor->emitError("projected fragment layout is missing fragment type metadata");
-  if (layout.fragmentShape.size() != static_cast<size_t>(layout.fragmentType.getRank()))
-    return anchor->emitError("projected fragment layout rank does not match fragment type");
-  if (layout.payloadFragmentCount == 0 || layout.fragmentsPerLogicalSlot == 0)
-    return anchor->emitError("projected fragment layout has an invalid fragment count");
-  if (layout.payloadFragmentCount % layout.fragmentsPerLogicalSlot != 0)
-    return anchor->emitError("projected fragment layout payload fragment count is incompatible with logical slots");
-  return mlir::success();
-}
-
-mlir::FailureOr<mlir::RankedTensorType>
-getProjectedPayloadType(mlir::Operation* anchor, mlir::RankedTensorType fragmentType, unsigned payloadFragmentCount) {
-  auto packedType = getPackedBatchTensorType(fragmentType, payloadFragmentCount);
-  if (mlir::failed(packedType)) {
-    anchor->emitError("cannot create projected payload type");
-    return mlir::failure();
-  }
-  return *packedType;
-}
-
-llvm::SmallVector<llvm::SmallVector<int64_t, 16>, 4>
-buildProjectedFragmentOffsetsByDim(llvm::ArrayRef<llvm::SmallVector<int64_t, 4>> fragmentOffsets, size_t rank) {
-  llvm::SmallVector<llvm::SmallVector<int64_t, 16>, 4> fragmentOffsetsByDim(rank);
-  for (llvm::ArrayRef<int64_t> offsets : fragmentOffsets) {
-    assert(offsets.size() == rank && "projected offset rank mismatch");
-    for (size_t dim = 0; dim < rank; ++dim)
-      fragmentOffsetsByDim[dim].push_back(offsets[dim]);
-  }
-  return fragmentOffsetsByDim;
-}
-
-mlir::LogicalResult verifyProjectedTransferDescriptor(mlir::Operation* anchor,
-                                                      const ProjectedTransferDescriptor& descriptor) {
-  if (mlir::failed(verifyProjectedFragmentLayout(anchor, descriptor.layout)))
-    return mlir::failure();
-  if (!descriptor.payloadType)
-    return anchor->emitError("projected transfer descriptor is missing payload type");
-  if (descriptor.fragmentOffsets.empty())
-    return anchor->emitError("projected transfer descriptor expected at least one fragment offset");
-  if (descriptor.fragmentOffsetsByDim.size() != descriptor.layout.fragmentShape.size())
-    return anchor->emitError("projected transfer descriptor dimension-major offsets are inconsistent");
-  for (llvm::ArrayRef<int64_t> dimOffsets : descriptor.fragmentOffsetsByDim)
-    if (dimOffsets.size() != descriptor.fragmentOffsets.size())
-      return anchor->emitError("projected transfer descriptor dimension-major offsets are inconsistent");
-  for (llvm::ArrayRef<int64_t> offsets : descriptor.fragmentOffsets)
-    if (offsets.size() != descriptor.layout.fragmentShape.size())
-      return anchor->emitError("projected transfer offset rank does not match fragment rank");
-  return mlir::success();
-}
-
-mlir::LogicalResult verifyProjectedSendDescriptor(mlir::Operation* anchor,
-                                                  const ProjectedTransferDescriptor& descriptor,
-                                                  const MessageVector& messages) {
-  if (mlir::failed(verifyProjectedTransferDescriptor(anchor, descriptor)))
-    return mlir::failure();
-  if (messages.size() * descriptor.layout.payloadFragmentCount != descriptor.fragmentOffsets.size())
-    return anchor->emitError("projected send descriptor metadata is inconsistent");
-  return mlir::success();
-}
-
-mlir::LogicalResult finalizeProjectedTransferDescriptor(mlir::Operation* anchor,
-                                                        ProjectedTransferDescriptor& descriptor) {
-  descriptor.fragmentOffsetsByDim =
-    buildProjectedFragmentOffsetsByDim(descriptor.fragmentOffsets, descriptor.layout.fragmentShape.size());
-
-  auto payloadType =
-    getProjectedPayloadType(anchor, descriptor.layout.fragmentType, descriptor.layout.payloadFragmentCount);
-  if (mlir::failed(payloadType))
-    return mlir::failure();
-  if (descriptor.payloadType && descriptor.payloadType != *payloadType)
-    return anchor->emitError("projected transfer descriptor payload type does not match projected layout");
-  descriptor.payloadType = *payloadType;
-
-  return verifyProjectedTransferDescriptor(anchor, descriptor);
-}
-
-} // namespace onnx_mlir::spatial

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

@@ -1,87 +0,0 @@
-#pragma once
-
-#include "mlir/Dialect/Tensor/IR/Tensor.h"
-#include "mlir/IR/BuiltinAttributes.h"
-#include "mlir/IR/BuiltinTypes.h"
-#include "mlir/IR/Value.h"
-#include "mlir/IR/ValueRange.h"
-
-#include "llvm/ADT/ArrayRef.h"
-#include "llvm/ADT/SmallVector.h"
-
-#include <cstdint>
-
-#include "MergeMessages.hpp"
-#include "MergeScheduleKeys.hpp"
-
-namespace onnx_mlir::spatial {
-
-struct ProjectedFragmentLayout {
-  mlir::RankedTensorType fragmentType;
-  llvm::SmallVector<int64_t, 4> fragmentShape;
-  unsigned fragmentsPerLogicalSlot = 1;
-  unsigned payloadFragmentCount = 1;
-  llvm::SmallVector<int64_t, 4> loopLowerBounds;
-  llvm::SmallVector<int64_t, 4> loopSteps;
-  llvm::SmallVector<int64_t, 4> loopTripCounts;
-};
-
-struct StaticProjectedLoopInfo {
-  mlir::BlockArgument iv;
-  int64_t lowerBound = 0;
-  int64_t step = 1;
-  int64_t tripCount = 1;
-};
-
-struct ProjectedTransferDescriptor {
-  ProjectedBatchInputKey inputKey;
-  mlir::Operation* extractOp = nullptr;
-  ProjectedFragmentLayout layout;
-  mlir::RankedTensorType payloadType;
-  llvm::SmallVector<llvm::SmallVector<int64_t, 4>, 16> fragmentOffsets;
-  llvm::SmallVector<llvm::SmallVector<int64_t, 16>, 4> fragmentOffsetsByDim;
-};
-
-struct ProjectedExtractReplacement {
-  mlir::Value payload;
-  ProjectedFragmentLayout layout;
-};
-
-struct PendingProjectedHostOutputFragment {
-  mlir::Value originalOutput;
-  ClassId sourceClass = 0;
-  ProducerKey producerKey;
-  unsigned publicationResultIndex = 0;
-  int64_t sourceFragmentOrdinal = 0;
-  int64_t sourceElementOffset = 0;
-  llvm::SmallVector<int64_t, 4> offsets;
-  llvm::SmallVector<int64_t, 4> sizes;
-  llvm::SmallVector<int64_t, 4> strides;
-  uint32_t sourceLane = 0;
-  mlir::Location loc;
-};
-
-struct AffineProjectedInputSliceMatch {
-  mlir::tensor::ExtractSliceOp extract;
-  mlir::RankedTensorType sourceType;
-  mlir::RankedTensorType fragmentType;
-  llvm::SmallVector<int64_t, 4> fragmentShape;
-  llvm::SmallVector<mlir::OpFoldResult, 4> offsets;
-  llvm::SmallVector<StaticProjectedLoopInfo, 4> loops;
-};
-
-unsigned getProjectedFragmentsPerLogicalSlot(llvm::ArrayRef<int64_t> loopTripCounts);
-mlir::LogicalResult verifyProjectedFragmentLayout(mlir::Operation* anchor, const ProjectedFragmentLayout& layout);
-mlir::FailureOr<mlir::RankedTensorType>
-getProjectedPayloadType(mlir::Operation* anchor, mlir::RankedTensorType fragmentType, unsigned payloadFragmentCount);
-llvm::SmallVector<llvm::SmallVector<int64_t, 16>, 4>
-buildProjectedFragmentOffsetsByDim(llvm::ArrayRef<llvm::SmallVector<int64_t, 4>> fragmentOffsets, size_t rank);
-mlir::LogicalResult verifyProjectedTransferDescriptor(mlir::Operation* anchor,
-                                                      const ProjectedTransferDescriptor& descriptor);
-mlir::LogicalResult verifyProjectedSendDescriptor(mlir::Operation* anchor,
-                                                  const ProjectedTransferDescriptor& descriptor,
-                                                  const MessageVector& messages);
-mlir::LogicalResult finalizeProjectedTransferDescriptor(mlir::Operation* anchor,
-                                                        ProjectedTransferDescriptor& descriptor);
-
-} // namespace onnx_mlir::spatial

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -106,10 +106,11 @@ static std::optional<uint32_t> getConstantExtractLane(tensor::ExtractSliceOp ext
 
 static std::optional<ProducerValueRef> getResultfulBatchProducerValueRef(SpatComputeBatch batch,
                                                                          const ComputeInstance* consumerInstance) {
-  if (!consumerInstance)
-    return std::nullopt;
-  if (!isa<SpatComputeBatch>(consumerInstance->op))
-    return std::nullopt;
+  if (!consumerInstance || !isa<SpatComputeBatch>(consumerInstance->op))
+    return ProducerValueRef {
+      {batch.getOperation(), 0, static_cast<uint32_t>(batch.getLaneCount())},
+      0
+    };
   if (consumerInstance->laneStart + consumerInstance->laneCount > static_cast<uint32_t>(batch.getLaneCount()))
     return std::nullopt;
   return ProducerValueRef {

src/PIM/Pass/PIMPasses.h

@@ -11,8 +11,6 @@ std::unique_ptr<mlir::Pass> createONNXToSpatialPass();
 std::unique_ptr<mlir::Pass> createSpatialLayoutPlanningPass();
 std::unique_ptr<mlir::Pass> createLowerSpatialPlansPass();
 
-std::unique_ptr<mlir::Pass> createSpatialToGraphvizPass();
-
 std::unique_ptr<mlir::Pass> createSpatialToPimPass();
 
 std::unique_ptr<mlir::Pass> createPimBufferizationPass();

src/PIM/PimAccelerator.cpp

@@ -74,7 +74,6 @@ void PimAccelerator::registerPasses(int optLevel) const {
   registerPass(createONNXToSpatialPass);
   registerPass(createSpatialLayoutPlanningPass);
   registerPass(createLowerSpatialPlansPass);
-  registerPass(createSpatialToGraphvizPass);
   registerPass(createSpatialToPimPass);
   registerPass(createPimBufferizationPass);
   registerPass(createPimMemoryCoalescingPass);

tools/analyze_spatial_ir_cardinality.py

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