refactor

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/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;