Merge done

src/PIM/Dialect/Spatial/CMakeLists.txt

@@ -8,7 +8,9 @@ add_pim_library(SpatialOps
   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/Scheduling/MergeSchedulingAnalysis.cpp

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

@@ -0,0 +1,134 @@
+#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

@@ -0,0 +1,11 @@
+#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.cpp

@@ -24,14 +24,18 @@
 #include <string>
 #include <utility>
 
-#include "MaterializeMergeSchedule.hpp"
+#include "HostOutputFinalization.hpp"
+#include "MaterializedClassState.hpp"
+#include "MergeMessages.hpp"
+#include "MergeScheduleKeys.hpp"
+#include "ProjectedFragments.hpp"
 #include "Scheduling/ComputeInstanceUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp"
 #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp"
+#include "src/Accelerators/PIM/Common/IR/TensorSliceUtils.hpp"
 #include "src/Accelerators/PIM/Common/PimCommon.hpp"
 #include "src/Accelerators/PIM/Common/Support/CheckedArithmetic.hpp"
-#include "src/Accelerators/PIM/Conversion/ONNXToSpatial/Common/ShapeTilingUtils.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using namespace mlir;
@@ -39,353 +43,6 @@ using namespace mlir;
 namespace onnx_mlir {
 namespace spatial {
 namespace {
-
-using CpuId = size_t;
-using ClassId = size_t;
-using SlotId = size_t;
-
-static FailureOr<int32_t> getCheckedCoreId(Operation* anchor, CpuId cpu, StringRef fieldName) {
-  return pim::checkedI32(static_cast<uint64_t>(cpu), anchor, fieldName);
-}
-
-static FailureOr<SmallVector<int32_t, 8>>
-getCheckedCoreIds(Operation* anchor, ArrayRef<CpuId> cpus, StringRef fieldName) {
-  SmallVector<int32_t, 8> coreIds;
-  coreIds.reserve(cpus.size());
-  for (CpuId cpu : cpus) {
-    auto checkedCoreId = getCheckedCoreId(anchor, cpu, fieldName);
-    if (failed(checkedCoreId))
-      return failure();
-    coreIds.push_back(*checkedCoreId);
-  }
-  return coreIds;
-}
-
-struct MessageVector {
-  SmallVector<int64_t, 16> channelIds;
-  SmallVector<int32_t, 16> sourceCoreIds;
-  SmallVector<int32_t, 16> targetCoreIds;
-
-  size_t size() const { return channelIds.size(); }
-  bool empty() const { return channelIds.empty(); }
-
-  LogicalResult verify(Operation* anchor) const {
-    if (channelIds.size() != sourceCoreIds.size() || channelIds.size() != targetCoreIds.size())
-      return anchor->emitError("message metadata is inconsistent");
-    return 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(ArrayRef<int64_t> channels, ArrayRef<int32_t> sources, 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(ArrayRef<int64_t>(channelIds).slice(offset, count),
-                  ArrayRef<int32_t>(sourceCoreIds).slice(offset, count),
-                  ArrayRef<int32_t>(targetCoreIds).slice(offset, count));
-    return result;
-  }
-};
-
-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 {
-  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<Operation*>::getEmptyKey(),
-            std::numeric_limits<size_t>::max(),
-            std::numeric_limits<ClassId>::max()};
-  }
-
-  static SameClassConsumerLookupKey getTombstoneKey() {
-    return {llvm::DenseMapInfo<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<Operation*>::getHashValue(key.sourceOp), key.resultIndex, key.classId);
-  }
-
-  static bool isEqual(const SameClassConsumerLookupKey& lhs, const SameClassConsumerLookupKey& rhs) {
-    return lhs == rhs;
-  }
-};
-
-struct WholeBatchAssemblyLookupKey {
-  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<Operation*>::getEmptyKey(),
-            std::numeric_limits<size_t>::max(),
-            std::numeric_limits<ClassId>::max()};
-  }
-
-  static WholeBatchAssemblyLookupKey getTombstoneKey() {
-    return {llvm::DenseMapInfo<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<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 MaterializedClass {
-  ClassId id = 0;
-  SmallVector<CpuId, 8> cpus;
-  Operation* op = nullptr;
-  Block* body = nullptr;
-  bool isBatch = false;
-
-  DenseMap<CpuId, unsigned> cpuToLane;
-  SmallVector<Value, 8> weights;
-  SmallVector<Value, 8> inputs;
-  SmallVector<Value, 4> hostOutputs;
-  DenseMap<Value, unsigned> publicationOutputToResultIndex;
-  DenseMap<Value, BlockArgument> weightArgs;
-  DenseMap<Value, BlockArgument> inputArgs;
-  DenseMap<Value, unsigned> hostOutputToResultIndex;
-};
-
-struct PackedScalarRunSlot {
-  SmallVector<ProducerKey, 8> keys;
-};
-
-enum class PackedScalarRunKind {
-  Materialized,
-  DeferredReceive,
-  DeferredLocalCompute
-};
-
-struct PackedScalarRunValue {
-  ClassId targetClass = 0;
-  Operation* sourceOp = nullptr;
-  size_t resultIndex = 0;
-  PackedScalarRunKind kind = PackedScalarRunKind::Materialized;
-
-  Value packed;
-
-  RankedTensorType fragmentType;
-  SmallVector<PackedScalarRunSlot, 8> slots;
-  MessageVector messages;
-};
-
-struct IndexedBatchRunValue {
-  ClassId targetClass = 0;
-  Operation* sourceOp = nullptr;
-  size_t resultIndex = 0;
-  Value packed;
-  RankedTensorType fragmentType;
-  SmallVector<PackedScalarRunSlot, 8> slots;
-  MessageVector messages;
-};
-
-struct LogicalSlotRange {
-  SlotId start = 0;
-  SlotId count = 0;
-};
-
-struct MaterializationRunSlot {
-  SmallVector<ComputeInstance, 8> peers;
-};
-
-using MaterializationRun = SmallVector<MaterializationRunSlot, 8>;
-
-struct OutputDestinationGroup {
-  SmallVector<size_t, 4> resultIndices;
-  SmallVector<ClassId, 4> destinationClasses;
-};
-
-struct BatchRunSendPlan {
-  size_t resultIndex = 0;
-  ClassId destinationClass = 0;
-  MessageVector messages;
-};
-
-struct ProjectedBatchInputKey {
-  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<Operation*>::getEmptyKey(), std::numeric_limits<unsigned>::max()};
-  }
-
-  static ProjectedBatchInputKey getTombstoneKey() {
-    return {llvm::DenseMapInfo<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; }
-};
-
-struct ProjectedFragmentLayout {
-  RankedTensorType fragmentType;
-  SmallVector<int64_t, 4> fragmentShape;
-  unsigned fragmentsPerLogicalSlot = 1;
-  unsigned payloadFragmentCount = 1;
-  SmallVector<int64_t, 4> loopLowerBounds;
-  SmallVector<int64_t, 4> loopSteps;
-  SmallVector<int64_t, 4> loopTripCounts;
-};
-
-struct StaticProjectedLoopInfo {
-  BlockArgument iv;
-  int64_t lowerBound = 0;
-  int64_t step = 1;
-  int64_t tripCount = 1;
-};
-
-struct ProjectedTransferDescriptor {
-  ProjectedBatchInputKey inputKey;
-  Operation* extractOp = nullptr;
-  ProjectedFragmentLayout layout;
-  RankedTensorType payloadType;
-  SmallVector<SmallVector<int64_t, 4>, 16> fragmentOffsets;
-  SmallVector<SmallVector<int64_t, 16>, 4> fragmentOffsetsByDim;
-};
-
-struct ProjectedExtractReplacement {
-  Value payload;
-  ProjectedFragmentLayout layout;
-};
-
-struct PendingProjectedHostOutputFragment {
-  Value originalOutput;
-  ClassId sourceClass = 0;
-  ProducerKey producerKey;
-  unsigned publicationResultIndex = 0;
-  int64_t sourceFragmentOrdinal = 0;
-  int64_t sourceElementOffset = 0;
-  SmallVector<int64_t, 4> offsets;
-  SmallVector<int64_t, 4> sizes;
-  SmallVector<int64_t, 4> strides;
-  uint32_t sourceLane = 0;
-  Location loc;
-};
-
-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;
-  Value originalOutput;
-  RankedTensorType fragmentType;
-  SmallVector<TensorDemandAction, 4> actions;
-};
-
-struct CompactRunPlan {
-  SmallVector<RunOutputDemand, 4> outputs;
-};
-
-enum class BatchInputDemandKind {
-  LaneFragment,
-  ProjectedFragment,
-  WholeTensorBarrier
-};
-
-struct BatchInputDemand {
-  BatchInputDemandKind kind = BatchInputDemandKind::LaneFragment;
-  std::optional<ProducerKey> wholeTensorProducer;
-};
-
-struct AffineProjectedInputSliceMatch {
-  tensor::ExtractSliceOp extract;
-  RankedTensorType sourceType;
-  RankedTensorType fragmentType;
-  SmallVector<int64_t, 4> fragmentShape;
-  SmallVector<OpFoldResult, 4> offsets;
-  SmallVector<StaticProjectedLoopInfo, 4> loops;
-};
-
-struct CloneIndexingContext {
-  std::optional<Value> runSlotIndex;
-  std::optional<Value> projectionSlotIndex;
-};
-
-struct MaterializerState;
 FailureOr<bool> recordProjectedScalarHostFragmentsFromPackedValue(MaterializerState& state,
                                                                   MaterializedClass& sourceClass,
                                                                   ArrayRef<ProducerKey> keys,
@@ -444,111 +101,6 @@ FailureOr<Value> materializeProjectedWholeBatchExtractReplacement(MaterializerSt
                                                                   tensor::ExtractSliceOp extract,
                                                                   ProducerKey producer,
                                                                   IRMapping* mapper = nullptr);
-
-class AvailableValueStore {
-public:
-  struct ExactBatchFragmentRecord {
-    ProducerKey key;
-    Value value;
-  };
-
-  void record(ProducerKey key, ClassId classId, Value value) {
-    exactValues[key][classId] = value;
-
-    auto batch = dyn_cast_or_null<SpatComputeBatch>(key.instance.op);
-    if (!batch || key.instance.laneCount == 0)
-      return;
-
-    WholeBatchAssemblyLookupKey lookupKey {batch.getOperation(), key.resultIndex, classId};
-    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<Value> lookupExact(ProducerKey key, ClassId classId) const;
-
-  std::optional<Value> lookup(MaterializerState& state, ProducerKey key, ClassId classId);
-  IndexedBatchRunValue* lookupIndexedBatchRun(ProducerKey key, ClassId classId);
-
-  ArrayRef<size_t> getPackedRunIndicesForWholeBatch(WholeBatchAssemblyLookupKey key) const {
-    auto it = packedRunsByProducerResultClass.find(key);
-    if (it == packedRunsByProducerResultClass.end())
-      return {};
-    return it->second;
-  }
-
-  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<Value> lookupPackedRun(MaterializerState& state, ProducerKey key, ClassId classId);
-
-  DenseMap<ProducerKey, DenseMap<ClassId, Value>, ProducerKeyInfo> exactValues;
-  SmallVector<PackedScalarRunValue, 8> packedScalarRuns;
-  SmallVector<IndexedBatchRunValue, 8> indexedBatchRuns;
-  DenseMap<WholeBatchAssemblyLookupKey, SmallVector<ExactBatchFragmentRecord, 16>, WholeBatchAssemblyLookupKeyInfo>
-    exactBatchFragmentsByProducerResultClass;
-  DenseMap<WholeBatchAssemblyLookupKey, SmallVector<size_t, 16>, WholeBatchAssemblyLookupKeyInfo>
-    packedRunsByProducerResultClass;
-};
-
-struct MaterializerState {
-  func::FuncOp func;
-  const MergeScheduleResult& schedule;
-  IRRewriter rewriter;
-  OperationFolder constantFolder;
-  int64_t& nextChannelId;
-  SmallVector<MaterializedClass, 8> classes;
-  DenseMap<CpuId, ClassId> cpuToClass;
-  DenseMap<CpuId, SmallVector<ComputeInstance, 32>> logicalInstancesByCpu;
-  DenseMap<ComputeInstance, LogicalSlotRange> scheduledInstanceToLogicalSlots;
-  DenseMap<ComputeInstance, ComputeInstance> logicalInstanceToScheduledChunk;
-  DenseSet<ClassSlotKey> materializedLogicalSlots;
-
-  DenseMap<ProducerKey, SmallVector<ClassId, 4>, ProducerKeyInfo> producerDestClasses;
-  DenseMap<SameClassConsumerLookupKey, SmallVector<ProducerKey, 4>, SameClassConsumerLookupKeyInfo>
-    sameClassConsumerIndex;
-  DenseMap<ProjectedBatchInputKey, AffineProjectedInputSliceMatch, ProjectedBatchInputKeyInfo> projectedInputMatches;
-  DenseSet<ProjectedBatchInputKey, ProjectedBatchInputKeyInfo> nonProjectedInputs;
-  DenseMap<Value, bool> liveExternalUseCache;
-  DenseMap<Operation*, SmallVector<Type, 4>> batchOutputFragmentTypesCache;
-  DenseMap<ComputeInstance, SmallVector<Value, 4>, llvm::DenseMapInfo<ComputeInstance>> computeInstanceOutputsCache;
-  DenseMap<ProducerKey, DenseMap<ClassId, ProjectedTransferDescriptor>, ProducerKeyInfo> projectedTransfers;
-  DenseMap<Operation*, DenseMap<ClassId, ProjectedExtractReplacement>> projectedExtractReplacements;
-  AvailableValueStore availableValues;
-  DenseMap<Value, Value> hostReplacements;
-  DenseMap<Value, ClassId> hostOutputOwners;
-  SmallVector<PendingProjectedHostOutputFragment, 32> pendingProjectedHostOutputFragments;
-  DenseSet<Operation*> oldComputeOps;
-
-  MaterializerState(func::FuncOp func, const MergeScheduleResult& schedule, int64_t& nextChannelId)
-  : func(func),
-    schedule(schedule),
-    rewriter(func.getContext()),
-    constantFolder(func.getContext()),
-    nextChannelId(nextChannelId) {}
-};
-
 bool isConstantLike(Value value) {
   Operation* definingOp = value.getDefiningOp();
   return definingOp && definingOp->hasTrait<OpTrait::ConstantLike>();
@@ -1260,17 +812,6 @@ LogicalResult createEmptyMaterializedOps(MaterializerState& state) {
   return success();
 }
 
-void setInsertionPointForNewMaterializedOp(MaterializerState& state) {
-  Block& funcBlock = state.func.getBody().front();
-  for (Operation& op : funcBlock) {
-    if (state.oldComputeOps.contains(&op)) {
-      state.rewriter.setInsertionPoint(&op);
-      return;
-    }
-  }
-  state.rewriter.setInsertionPointToEnd(&funcBlock);
-}
-
 BlockArgument appendWeight(MaterializerState& state, MaterializedClass& materializedClass, Value weight) {
   auto it = materializedClass.weightArgs.find(weight);
   if (it != materializedClass.weightArgs.end())
@@ -1411,19 +952,6 @@ FailureOr<unsigned> appendBatchPublicationResult(MaterializerState& state,
 // Materialized-class value localization helpers.
 // -----------------------------------------------------------------------------
 
-Region* getParentRegion(Value value) {
-  if (auto blockArg = dyn_cast<BlockArgument>(value))
-    return blockArg.getOwner()->getParent();
-  if (Operation* definingOp = value.getDefiningOp())
-    return definingOp->getParentRegion();
-  return nullptr;
-}
-
-bool isDefinedInsideRegion(Value value, Region& region) {
-  Region* parentRegion = getParentRegion(value);
-  return parentRegion && (&region == parentRegion || region.isAncestor(parentRegion));
-}
-
 Operation* getEnclosingSpatialComputeLikeOp(Value value) {
   Block* block = nullptr;
   if (auto blockArg = dyn_cast<BlockArgument>(value))
@@ -2169,14 +1697,6 @@ buildProjectedFragmentOffsetsInClass(MaterializerState& state,
   return fragmentOffsets;
 }
 
-SmallVector<OpFoldResult, 4> getStaticIndexAttrs(Builder& builder, ArrayRef<int64_t> values) {
-  SmallVector<OpFoldResult, 4> attrs;
-  attrs.reserve(values.size());
-  for (int64_t value : values)
-    attrs.push_back(builder.getIndexAttr(value));
-  return attrs;
-}
-
 Value createDim0InsertSlice(
   MaterializerState& state, Location loc, Value fragment, Value destination, OpFoldResult firstOffset) {
   auto fragmentType = cast<RankedTensorType>(fragment.getType());
@@ -2293,6 +1813,8 @@ std::optional<Value> extractPackedProducerSlice(MaterializerState& state,
   return createDim0ExtractSlice(state, materializedClass.op->getLoc(), packed, firstOffset, rowCount);
 }
 
+} // namespace
+
 std::optional<Value> AvailableValueStore::lookupExact(ProducerKey key, ClassId classId) const {
   auto producerIt = exactValues.find(key);
   if (producerIt == exactValues.end())
@@ -2305,6 +1827,32 @@ std::optional<Value> AvailableValueStore::lookupExact(ProducerKey key, ClassId c
   return valueIt->second;
 }
 
+namespace {
+
+using IndexedFragmentBuilder = llvm::function_ref<FailureOr<Value>(Value flatIndex)>;
+using IndexedInsertOffsetBuilder = llvm::function_ref<FailureOr<Value>(Value flatIndex)>;
+
+SmallVector<ProducerKey, 16> flattenPackedScalarRunKeys(const PackedScalarRunValue& run);
+FailureOr<Value> emitIndexedFragmentInsertLoop(MaterializerState& state,
+                                               MaterializedClass& targetClass,
+                                               Value destination,
+                                               int64_t itemCount,
+                                               IndexedFragmentBuilder buildFragment,
+                                               IndexedInsertOffsetBuilder buildOffset,
+                                               Location loc);
+FailureOr<SmallVector<Value, 4>> cloneBatchBodyForLane(MaterializerState& state,
+                                                       MaterializedClass& targetClass,
+                                                       const ComputeInstance& instance,
+                                                       Value laneValue,
+                                                       ArrayRef<size_t> resultIndices,
+                                                       CloneIndexingContext indexing);
+Value createIndexedChannelId(
+  MaterializerState& state, Operation* anchor, const MessageVector& messages, Value index, Location loc);
+Value createIndexedSourceCoreId(
+  MaterializerState& state, Operation* anchor, const MessageVector& messages, Value index, Location loc);
+Value createIndexedTargetCoreId(
+  MaterializerState& state, Operation* anchor, const MessageVector& messages, Value index, Location loc);
+
 Value getPackedSliceForRunIndex(MaterializerState& state,
                                 Operation* anchor,
                                 Value packed,
@@ -2322,21 +1870,9 @@ Value getPackedSliceForDynamicRunIndex(
   return createDim0ExtractSlice(state, loc, packed, firstOffset, fragmentType.getDimSize(0));
 }
 
-FailureOr<Value> createReceiveConcatLoop(MaterializerState& state,
-                                         MaterializedClass& targetClass,
-                                         RankedTensorType concatType,
-                                         RankedTensorType fragmentType,
-                                         const MessageVector& messages,
-                                         Location loc);
-
 using IndexedFragmentBuilder = llvm::function_ref<FailureOr<Value>(Value flatIndex)>;
 using IndexedInsertOffsetBuilder = llvm::function_ref<FailureOr<Value>(Value flatIndex)>;
 
-FailureOr<Value> materializeDeferredLocalPackedScalarRunValue(MaterializerState& state,
-                                                              MaterializedClass& targetClass,
-                                                              PackedScalarRunValue& run,
-                                                              Location loc);
-
 bool isDeferredLocalPackedScalarRun(const PackedScalarRunValue& run) {
   return run.kind == PackedScalarRunKind::DeferredLocalCompute;
 }
@@ -2376,8 +1912,69 @@ FailureOr<Value> materializePackedScalarRunValue(MaterializerState& state,
   if (run.kind == PackedScalarRunKind::Materialized)
     return targetClass.op->emitError("materialized packed scalar run has no packed value");
 
-  if (isDeferredLocalPackedScalarRun(run))
-    return materializeDeferredLocalPackedScalarRunValue(state, targetClass, run, loc);
+  if (isDeferredLocalPackedScalarRun(run)) {
+    SmallVector<ProducerKey, 16> keys = flattenPackedScalarRunKeys(run);
+    if (keys.empty())
+      return failure();
+    FailureOr<RankedTensorType> packedType = getPackedBatchTensorType(run.fragmentType, keys.size());
+    if (failed(packedType))
+      return targetClass.op->emitError("cannot materialize deferred local packed run for non-static ranked tensor");
+
+    SmallVector<int64_t, 16> sourceLanes;
+    sourceLanes.reserve(keys.size());
+    for (ProducerKey key : keys) {
+      if (key.instance.laneCount != 1)
+        return failure();
+      sourceLanes.push_back(key.instance.laneStart);
+    }
+
+    SmallVector<size_t, 1> resultIndices {run.resultIndex};
+
+    state.rewriter.setInsertionPoint(targetClass.body->getTerminator());
+    Value init =
+      tensor::EmptyOp::create(state.rewriter, loc, packedType->getShape(), packedType->getElementType()).getResult();
+
+    Value lowerBound = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 0);
+    Value upperBound = getOrCreateIndexConstant(state.constantFolder, targetClass.op, static_cast<int64_t>(keys.size()));
+    Value step = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 1);
+
+    auto loop = buildNormalizedScfFor(
+      state.rewriter,
+      loc,
+      lowerBound,
+      upperBound,
+      step,
+      ValueRange {init},
+      [&](OpBuilder&, Location, Value loopIndex, ValueRange iterArgs, SmallVectorImpl<Value>& yielded) {
+        Value acc = iterArgs.front();
+        Value sourceLane = createIndexedIndexValue(state, targetClass.op, sourceLanes, loopIndex, loc);
+
+        FailureOr<SmallVector<Value, 4>> produced =
+          cloneBatchBodyForLane(state,
+                                targetClass,
+                                keys.front().instance,
+                                sourceLane,
+                                resultIndices,
+                                CloneIndexingContext {.runSlotIndex = std::nullopt, .projectionSlotIndex = loopIndex});
+        if (failed(produced) || produced->size() != 1)
+          return failure();
+
+        FailureOr<Value> firstOffset =
+          scaleIndexByDim0SizeInClass(state, targetClass, loopIndex, run.fragmentType.getDimSize(0), loc);
+        if (failed(firstOffset))
+          return failure();
+        FailureOr<Value> next =
+          createDim0InsertSliceInClass(state, targetClass, loc, produced->front(), acc, *firstOffset);
+        if (failed(next))
+          return failure();
+        yielded.push_back(*next);
+        return success();
+      });
+    if (failed(loop))
+      return failure();
+    run.packed = loop->results.front();
+    return run.packed;
+  }
 
   if (failed(validatePackedScalarRunMetadata(targetClass.op, run)))
     return failure();
@@ -2387,13 +1984,34 @@ FailureOr<Value> materializePackedScalarRunValue(MaterializerState& state,
   if (failed(fullPackedType))
     return targetClass.op->emitError("cannot create lazy packed scalar run receive type");
 
-  auto packed = createReceiveConcatLoop(state, targetClass, *fullPackedType, run.fragmentType, run.messages, loc);
+  state.rewriter.setInsertionPoint(targetClass.body->getTerminator());
+  Value init =
+    tensor::EmptyOp::create(state.rewriter, loc, fullPackedType->getShape(), fullPackedType->getElementType())
+      .getResult();
+  auto packed = emitIndexedFragmentInsertLoop(
+    state,
+    targetClass,
+    init,
+    static_cast<int64_t>(run.messages.size()),
+    [&](Value index) -> FailureOr<Value> {
+      Value channelId = createIndexedChannelId(state, targetClass.op, run.messages, index, loc);
+      Value sourceCoreId = createIndexedSourceCoreId(state, targetClass.op, run.messages, index, loc);
+      Value targetCoreId = createIndexedTargetCoreId(state, targetClass.op, run.messages, index, loc);
+      return SpatChannelReceiveOp::create(state.rewriter, loc, run.fragmentType, channelId, sourceCoreId, targetCoreId)
+        .getOutput();
+    },
+    [&](Value index) -> FailureOr<Value> {
+      return scaleIndexByDim0SizeInClass(state, targetClass, index, run.fragmentType.getDimSize(0), loc);
+    },
+    loc);
   if (failed(packed))
     return failure();
   run.packed = *packed;
   return run.packed;
 }
 
+} // namespace
+
 std::optional<Value> AvailableValueStore::lookupPackedRun(MaterializerState& state, ProducerKey key, ClassId classId) {
   for (PackedScalarRunValue& run : packedScalarRuns) {
     if (run.targetClass != classId || run.sourceOp != key.instance.op || run.resultIndex != key.resultIndex)
@@ -2491,6 +2109,8 @@ std::optional<Value> AvailableValueStore::lookup(MaterializerState& state, Produ
   return std::nullopt;
 }
 
+namespace {
+
 Value createIndexTensorConstant(MaterializerState& state, Operation* anchor, ArrayRef<int64_t> values) {
   SmallVector<APInt, 8> elements;
   elements.reserve(values.size());
@@ -2983,89 +2603,6 @@ isProjectedOffsetValue(Value value, Value laneArg, ArrayRef<StaticProjectedLoopI
 
 static std::optional<int64_t> getConstantIndex(OpFoldResult value);
 
-static unsigned getProjectedFragmentsPerLogicalSlot(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;
-}
-
-LogicalResult verifyProjectedFragmentLayout(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 success();
-}
-
-FailureOr<RankedTensorType>
-getProjectedPayloadType(Operation* anchor, RankedTensorType fragmentType, unsigned payloadFragmentCount) {
-  if (failed(
-        verifyPackableFragmentType(anchor, fragmentType, payloadFragmentCount, "cannot create projected payload type")))
-    return failure();
-  return getPackedBatchTensorType(fragmentType, payloadFragmentCount);
-}
-
-SmallVector<SmallVector<int64_t, 16>, 4>
-buildProjectedFragmentOffsetsByDim(ArrayRef<SmallVector<int64_t, 4>> fragmentOffsets, size_t rank) {
-  SmallVector<SmallVector<int64_t, 16>, 4> fragmentOffsetsByDim(rank);
-  for (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;
-}
-
-LogicalResult verifyProjectedTransferDescriptor(Operation* anchor, const ProjectedTransferDescriptor& descriptor) {
-  if (failed(verifyProjectedFragmentLayout(anchor, descriptor.layout)))
-    return 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 (ArrayRef<int64_t> dimOffsets : descriptor.fragmentOffsetsByDim)
-    if (dimOffsets.size() != descriptor.fragmentOffsets.size())
-      return anchor->emitError("projected transfer descriptor dimension-major offsets are inconsistent");
-  for (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 success();
-}
-
-LogicalResult verifyProjectedSendDescriptor(Operation* anchor,
-                                            const ProjectedTransferDescriptor& descriptor,
-                                            const MessageVector& messages) {
-  if (failed(verifyProjectedTransferDescriptor(anchor, descriptor)))
-    return failure();
-  if (messages.size() * descriptor.layout.payloadFragmentCount != descriptor.fragmentOffsets.size())
-    return anchor->emitError("projected send descriptor metadata is inconsistent");
-  return success();
-}
-
-LogicalResult finalizeProjectedTransferDescriptor(Operation* anchor, ProjectedTransferDescriptor& descriptor) {
-  descriptor.fragmentOffsetsByDim =
-    buildProjectedFragmentOffsetsByDim(descriptor.fragmentOffsets, descriptor.layout.fragmentShape.size());
-
-  FailureOr<RankedTensorType> payloadType =
-    getProjectedPayloadType(anchor, descriptor.layout.fragmentType, descriptor.layout.payloadFragmentCount);
-  if (failed(payloadType))
-    return 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);
-}
-
 static FailureOr<int64_t> evaluateProjectedOffsetValue(OpFoldResult value,
                                                        Value laneArg,
                                                        uint32_t lane,
@@ -4819,12 +4356,12 @@ FailureOr<Value> materializeIndexedBatchRunReceive(MaterializerState& state,
                                                    Value runSlotIndex,
                                                    Location loc);
 
+} // namespace
+
 FailureOr<Value> materializeDeferredLocalPackedScalarRunValue(MaterializerState& state,
                                                               MaterializedClass& targetClass,
                                                               PackedScalarRunValue& run,
                                                               Location loc) {
-  assert(isDeferredLocalPackedScalarRun(run) && "expected deferred local packed scalar run");
-
   SmallVector<ProducerKey, 16> keys = flattenPackedScalarRunKeys(run);
   if (keys.empty())
     return failure();
@@ -4888,6 +4425,8 @@ FailureOr<Value> materializeDeferredLocalPackedScalarRunValue(MaterializerState&
   return run.packed;
 }
 
+namespace {
+
 LogicalResult collectPackedRunsForWholeBatchInput(MaterializerState& state,
                                                   MaterializedClass& targetClass,
                                                   ProducerKey key,
@@ -5946,119 +5485,6 @@ FailureOr<bool> recordProjectedScalarHostFragmentsFromPackedValue(MaterializerSt
   return true;
 }
 
-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 = spatial::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();
-}
-
 FailureOr<Value> resolveInputValue(MaterializerState& state,
                                    MaterializedClass& targetClass,
                                    Value input,
@@ -8191,36 +7617,6 @@ LogicalResult materializeInstanceSlot(MaterializerState& state, const ComputeIns
   return success();
 }
 
-FailureOr<Value> createReceiveConcatLoop(MaterializerState& state,
-                                         MaterializedClass& targetClass,
-                                         RankedTensorType concatType,
-                                         RankedTensorType fragmentType,
-                                         const MessageVector& messages,
-                                         Location loc) {
-  assert(succeeded(messages.verify(targetClass.op)) && "message metadata is inconsistent");
-  assert(!messages.empty() && "expected at least one receive");
-
-  state.rewriter.setInsertionPoint(targetClass.body->getTerminator());
-  Value init =
-    tensor::EmptyOp::create(state.rewriter, loc, concatType.getShape(), concatType.getElementType()).getResult();
-  return emitIndexedFragmentInsertLoop(
-    state,
-    targetClass,
-    init,
-    static_cast<int64_t>(messages.size()),
-    [&](Value index) -> FailureOr<Value> {
-      Value channelId = createIndexedChannelId(state, targetClass.op, messages, index, loc);
-      Value sourceCoreId = createIndexedSourceCoreId(state, targetClass.op, messages, index, loc);
-      Value targetCoreId = createIndexedTargetCoreId(state, targetClass.op, messages, index, loc);
-      return SpatChannelReceiveOp::create(state.rewriter, loc, fragmentType, channelId, sourceCoreId, targetCoreId)
-        .getOutput();
-    },
-    [&](Value index) -> FailureOr<Value> {
-      return scaleIndexByDim0SizeInClass(state, targetClass, index, fragmentType.getDimSize(0), loc);
-    },
-    loc);
-}
-
 bool valueMayEvaluateToCore(Value value, int64_t coreId) {
   if (std::optional<int64_t> constant = getConstantIndexValue(value))
     return *constant == coreId;

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

@@ -0,0 +1,252 @@
+#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

@@ -28,6 +28,7 @@
 #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"
@@ -43,16 +44,6 @@ using namespace onnx_mlir::compact_asm;
 using SpatCompute = spatial::SpatGraphCompute;
 using SpatComputeBatch = spatial::SpatGraphComputeBatch;
 
-static std::optional<int32_t> getComputeCoreId(spatial::SpatScheduledCompute compute) {
-  if (auto coreIdAttr = compute->getAttrOfType<IntegerAttr>(onnx_mlir::kCoreIdAttrName)) {
-    auto checkedCoreId = pim::checkedI32(coreIdAttr.getInt(), compute, "merge compute core id");
-    if (failed(checkedCoreId))
-      return std::nullopt;
-    return *checkedCoreId;
-  }
-  return std::nullopt;
-}
-
 bool isTrivialSerialMergeCandidate(SpatCompute compute) {
   if (!compute->hasOneUse())
     return false;
@@ -213,8 +204,11 @@ void generateReport(func::FuncOp funcOp, const std::string& name, size_t usedCpu
       uint64_t numInst = spatial::countComputeBodyInstructions(spatCompute.getBody());
       uint64_t perInstanceCrossbarCount = getPerInstanceCrossbarCount(spatCompute.getOperation());
       SmallVector<int32_t> coreIds;
-      if (auto coreId = getComputeCoreId(spatCompute))
-        coreIds.push_back(*coreId);
+      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;
@@ -234,8 +228,11 @@ void generateReport(func::FuncOp funcOp, const std::string& name, size_t usedCpu
       uint64_t logicalCount = static_cast<uint64_t>(batch.getLaneCount());
       uint64_t perInstanceCrossbarCount = getPerInstanceCrossbarCount(batch.getOperation());
       SmallVector<int32_t> coreIds;
-      if (auto coreIdsAttr = batch->getAttrOfType<DenseI32ArrayAttr>(onnx_mlir::kCoreIdsAttrName))
-        llvm::append_range(coreIds, coreIdsAttr.asArrayRef());
+      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;

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

@@ -0,0 +1,67 @@
+#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

@@ -0,0 +1,134 @@
+#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

@@ -0,0 +1,104 @@
+#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

@@ -0,0 +1,87 @@
+#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/Scheduling/Utils.hpp

@@ -12,7 +12,6 @@
 #include <utility>
 #include <vector>
 
-#include "src/Accelerators/PIM/Common/LabeledList.hpp"
 #include "src/Accelerators/PIM/Dialect/Spatial/SpatialOps.hpp"
 
 using CPU = int;