cose brutte

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

@@ -3,86 +3,255 @@
 #include "mlir/IR/Diagnostics.h"
 
 #include "llvm/ADT/DenseSet.h"
-#include "llvm/ADT/TypeSwitch.h"
+#include "llvm/ADT/STLExtras.h"
+
+#include <algorithm>
 
 using namespace mlir;
 
 namespace onnx_mlir::spatial {
 namespace {
 
-struct GroupSignature {
-  CommunicationPhase phase = CommunicationPhase::Local;
-  ClassId sourceClass = 0;
-  ClassId targetClass = 0;
-  Type payloadType;
-  bool sourceIsBatch = false;
-  bool targetIsBatch = false;
-
-  bool operator==(const GroupSignature& other) const {
-    return phase == other.phase && sourceClass == other.sourceClass && targetClass == other.targetClass
-        && payloadType == other.payloadType && sourceIsBatch == other.sourceIsBatch
-        && targetIsBatch == other.targetIsBatch;
+StringRef stringifyExchangeKind(CommunicationExchangeKind kind) {
+  switch (kind) {
+  case CommunicationExchangeKind::HostForward: return "HostForward";
+  case CommunicationExchangeKind::OrdinaryValue: return "OrdinaryValue";
+  case CommunicationExchangeKind::PackedScalarRun: return "PackedScalarRun";
+  case CommunicationExchangeKind::IndexedBatchRun: return "IndexedBatchRun";
+  case CommunicationExchangeKind::ProjectedInput: return "ProjectedInput";
   }
-};
+  llvm_unreachable("unknown communication exchange kind");
+}
 
-struct GroupSignatureInfo {
-  static GroupSignature getEmptyKey() {
-    return {CommunicationPhase::Local,
-            llvm::DenseMapInfo<ClassId>::getEmptyKey(),
-            llvm::DenseMapInfo<ClassId>::getEmptyKey(),
-            llvm::DenseMapInfo<Type>::getEmptyKey(),
-            false,
-            false};
+unsigned getExchangeKindOrder(CommunicationExchangeKind kind) {
+  switch (kind) {
+  case CommunicationExchangeKind::HostForward: return 0;
+  case CommunicationExchangeKind::OrdinaryValue: return 1;
+  case CommunicationExchangeKind::PackedScalarRun: return 2;
+  case CommunicationExchangeKind::IndexedBatchRun: return 3;
+  case CommunicationExchangeKind::ProjectedInput: return 4;
   }
-
-  static GroupSignature getTombstoneKey() {
-    return {CommunicationPhase::Local,
-            llvm::DenseMapInfo<ClassId>::getTombstoneKey(),
-            llvm::DenseMapInfo<ClassId>::getTombstoneKey(),
-            llvm::DenseMapInfo<Type>::getTombstoneKey(),
-            false,
-            false};
-  }
-
-  static unsigned getHashValue(const GroupSignature& key) {
-    return llvm::hash_combine(static_cast<unsigned>(key.phase),
-                              key.sourceClass,
-                              key.targetClass,
-                              key.payloadType,
-                              key.sourceIsBatch,
-                              key.targetIsBatch);
-  }
-
-  static bool isEqual(const GroupSignature& lhs, const GroupSignature& rhs) { return lhs == rhs; }
-};
+  llvm_unreachable("unknown communication exchange kind");
+}
 
 void emitPlanError(Operation* anchor, StringRef message) {
   if (anchor)
     anchor->emitError(message);
 }
 
+struct PeerPair {
+  int64_t minCore = 0;
+  int64_t maxCore = 0;
+
+  bool operator==(const PeerPair& other) const {
+    return minCore == other.minCore && maxCore == other.maxCore;
+  }
+};
+
+struct PeerPairInfo {
+  static PeerPair getEmptyKey() {
+    return {llvm::DenseMapInfo<int64_t>::getEmptyKey(), llvm::DenseMapInfo<int64_t>::getEmptyKey()};
+  }
+  static PeerPair getTombstoneKey() {
+    return {llvm::DenseMapInfo<int64_t>::getTombstoneKey(), llvm::DenseMapInfo<int64_t>::getTombstoneKey()};
+  }
+  static unsigned getHashValue(PeerPair key) { return llvm::hash_combine(key.minCore, key.maxCore); }
+  static bool isEqual(PeerPair lhs, PeerPair rhs) { return lhs == rhs; }
+};
+
+PeerPair getPeerPair(const CommunicationExchange& exchange) {
+  return {std::min(exchange.sourceCore, exchange.targetCore), std::max(exchange.sourceCore, exchange.targetCore)};
+}
+
+bool exchangeOrderLess(const CommunicationExchange& lhs, const CommunicationExchange& rhs) {
+  if (getExchangeKindOrder(lhs.kind) != getExchangeKindOrder(rhs.kind))
+    return getExchangeKindOrder(lhs.kind) < getExchangeKindOrder(rhs.kind);
+  if (lhs.sourceClass != rhs.sourceClass)
+    return lhs.sourceClass < rhs.sourceClass;
+  if (lhs.targetClass != rhs.targetClass)
+    return lhs.targetClass < rhs.targetClass;
+  if (lhs.producerOrder != rhs.producerOrder)
+    return lhs.producerOrder < rhs.producerOrder;
+  if (lhs.producer.resultIndex != rhs.producer.resultIndex)
+    return lhs.producer.resultIndex < rhs.producer.resultIndex;
+  if (lhs.producer.instance.laneStart != rhs.producer.instance.laneStart)
+    return lhs.producer.instance.laneStart < rhs.producer.instance.laneStart;
+  if (lhs.producer.instance.laneCount != rhs.producer.instance.laneCount)
+    return lhs.producer.instance.laneCount < rhs.producer.instance.laneCount;
+  if (lhs.sourceLane != rhs.sourceLane)
+    return lhs.sourceLane < rhs.sourceLane;
+  if (lhs.targetLane != rhs.targetLane)
+    return lhs.targetLane < rhs.targetLane;
+  if (lhs.ordinal != rhs.ordinal)
+    return lhs.ordinal < rhs.ordinal;
+  return lhs.channelId < rhs.channelId;
+}
+
+void sortExchangeIds(ArrayRef<CommunicationExchange> exchanges,
+                     SmallVectorImpl<CommunicationExchangeId>& ids) {
+  llvm::stable_sort(ids, [&](CommunicationExchangeId lhs, CommunicationExchangeId rhs) {
+    return exchangeOrderLess(exchanges[lhs.value], exchanges[rhs.value]);
+  });
+}
+
+void addEvent(CommunicationPlan& plan,
+              DenseMap<ClassId, SmallVector<CommunicationEvent, 16>>& eventsByClass,
+              CommunicationExchangeId id,
+              CommunicationEventKind kind,
+              ClassId classId,
+              int64_t core,
+              unsigned& order) {
+  eventsByClass[classId].push_back(CommunicationEvent {id, kind, classId, core, order++});
+  (void)plan;
+}
+
+void buildEventStreams(CommunicationPlan& plan,
+                       SmallVectorImpl<CommunicationExchange>& exchanges,
+                       DenseMap<ClassId, SmallVector<CommunicationEvent, 16>>& eventsByClass) {
+  DenseMap<PeerPair, SmallVector<CommunicationExchangeId, 8>, PeerPairInfo> exchangesByPair;
+  for (const CommunicationExchange& exchange : exchanges)
+    exchangesByPair[getPeerPair(exchange)].push_back(exchange.id);
+
+  SmallVector<PeerPair, 8> pairs;
+  pairs.reserve(exchangesByPair.size());
+  for (const auto& entry : exchangesByPair)
+    pairs.push_back(entry.first);
+  llvm::sort(pairs, [](PeerPair lhs, PeerPair rhs) {
+    if (lhs.minCore != rhs.minCore)
+      return lhs.minCore < rhs.minCore;
+    return lhs.maxCore < rhs.maxCore;
+  });
+
+  unsigned order = 0;
+  for (PeerPair pair : pairs) {
+    SmallVector<CommunicationExchangeId, 8> hiToLo;
+    SmallVector<CommunicationExchangeId, 8> loToHi;
+    for (CommunicationExchangeId id : exchangesByPair.lookup(pair)) {
+      const CommunicationExchange& exchange = exchanges[id.value];
+      if (exchange.sourceCore == pair.maxCore && exchange.targetCore == pair.minCore)
+        hiToLo.push_back(id);
+      else
+        loToHi.push_back(id);
+    }
+
+    sortExchangeIds(exchanges, hiToLo);
+    sortExchangeIds(exchanges, loToHi);
+
+    for (CommunicationExchangeId id : hiToLo) {
+      const CommunicationExchange& exchange = exchanges[id.value];
+      addEvent(plan, eventsByClass, id, CommunicationEventKind::Receive, exchange.targetClass, exchange.targetCore, order);
+    }
+    for (CommunicationExchangeId id : hiToLo) {
+      const CommunicationExchange& exchange = exchanges[id.value];
+      addEvent(plan, eventsByClass, id, CommunicationEventKind::Send, exchange.sourceClass, exchange.sourceCore, order);
+    }
+    for (CommunicationExchangeId id : loToHi) {
+      const CommunicationExchange& exchange = exchanges[id.value];
+      addEvent(plan, eventsByClass, id, CommunicationEventKind::Send, exchange.sourceClass, exchange.sourceCore, order);
+    }
+    for (CommunicationExchangeId id : loToHi) {
+      const CommunicationExchange& exchange = exchanges[id.value];
+      addEvent(plan, eventsByClass, id, CommunicationEventKind::Receive, exchange.targetClass, exchange.targetCore, order);
+    }
+  }
+}
+
 } // namespace
 
 FailureOr<CommunicationPlan>
-CommunicationPlan::buildProjectedInputPlan(Operation* anchor,
-                                           ArrayRef<CommunicationClassInfo> classes,
-                                           const DenseMap<CpuId, ClassId>& cpuToClass,
-                                           const ProjectedInputPlanMap& projectedPlans) {
+CommunicationPlan::build(Operation* anchor,
+                         ArrayRef<CommunicationClassInfo> classes,
+                         const DenseMap<CpuId, ClassId>& cpuToClass,
+                         const ProducerSourceClassMap& producerSourceClasses,
+                         const DenseMap<ProducerKey, Type, ProducerKeyInfo>& producerPayloadTypes,
+                         const DenseMap<ProducerKey, SmallVector<ClassId, 4>, ProducerKeyInfo>& producerDestClasses,
+                         const DenseMap<ProducerKey, SmallVector<ClassId, 4>, ProducerKeyInfo>& ordinaryProducerDestClasses,
+                         ProjectedInputPlanMap& projectedPlans,
+                         const DenseMap<Value, ClassId>& hostOutputOwners,
+                         const DenseMap<Value, bool>& liveExternalUseCache,
+                         int64_t& nextChannelId) {
   CommunicationPlan plan;
 
   DenseMap<ClassId, CommunicationClassInfo> classById;
   for (CommunicationClassInfo klass : classes)
-    classById[klass.classId] = klass;
+    classById[klass.classId] = std::move(klass);
 
-  DenseSet<int64_t> remoteChannels;
-  DenseSet<const ProjectedInputTransferFragment*> seenFragments;
-  DenseMap<GroupSignature, unsigned, GroupSignatureInfo> groupBySignature;
-  unsigned nextPlaceholder = 0;
+  auto appendExchange = [&](CommunicationExchangeKind kind,
+                            ProducerKey producer,
+                            ClassId sourceClass,
+                            ClassId targetClass,
+                            Type payloadType,
+                            int64_t sourceCore,
+                            int64_t targetCore,
+                            unsigned sourceLane,
+                            unsigned targetLane,
+                            unsigned ordinal,
+                            Value originalHostOutput = Value(),
+                            const ProjectedInputTransferFragment* fragment = nullptr) -> FailureOr<CommunicationExchangeId> {
+    if (sourceCore == targetCore)
+      return failure();
 
-  for (const auto& extractEntry : projectedPlans) {
-    for (const auto& classEntry : extractEntry.second) {
-      ClassId targetClass = classEntry.first;
-      const ProjectedInputTransferPlan& inputPlan = classEntry.second;
+    auto sourceInfoIt = classById.find(sourceClass);
+    auto targetInfoIt = classById.find(targetClass);
+    if (sourceInfoIt == classById.end() || targetInfoIt == classById.end())
+      return failure();
+
+    CommunicationExchangeId exchangeId {plan.exchanges.size()};
+    plan.exchanges.push_back(CommunicationExchange {
+      exchangeId,
+      kind,
+      PlaceholderId {static_cast<unsigned>(plan.demands.size())},
+      producer,
+      sourceClass,
+      targetClass,
+      sourceInfoIt->second.rank,
+      targetInfoIt->second.rank,
+      payloadType,
+      nextChannelId++,
+      sourceCore,
+      targetCore,
+      sourceLane,
+      targetLane,
+      ordinal,
+      exchangeId.value,
+      originalHostOutput,
+      fragment,
+    });
+    return exchangeId;
+  };
+
+  auto getPayloadType = [&](ProducerKey producer) -> Type {
+    auto typeIt = producerPayloadTypes.find(producer);
+    if (typeIt != producerPayloadTypes.end())
+      return typeIt->second;
+    if (!producer.instance.op || producer.resultIndex >= producer.instance.op->getNumResults())
+      return Type();
+    return producer.instance.op->getResult(producer.resultIndex).getType();
+  };
+
+  for (const auto& entry : producerDestClasses) {
+    ProducerKey producer = entry.first;
+    auto sourceClassIt = producerSourceClasses.find(producer);
+    if (sourceClassIt == producerSourceClasses.end()) {
+      emitPlanError(anchor, "communication plan could not find producer source class");
+      return failure();
+    }
+    ClassId sourceClass = sourceClassIt->second;
+    auto sourceInfoIt = classById.find(sourceClass);
+    if (sourceInfoIt == classById.end()) {
+      emitPlanError(anchor, "communication plan references an unknown producer source class");
+      return failure();
+    }
+    const CommunicationClassInfo& sourceInfo = sourceInfoIt->second;
+    Type payloadType = getPayloadType(producer);
+    if (!payloadType) {
+      emitPlanError(anchor, "communication plan could not determine producer payload type");
+      return failure();
+    }
+
+    for (ClassId targetClass : entry.second) {
+      if (targetClass == sourceClass)
+        continue;
 
       auto targetInfoIt = classById.find(targetClass);
       if (targetInfoIt == classById.end()) {
@@ -91,7 +260,190 @@ CommunicationPlan::buildProjectedInputPlan(Operation* anchor,
       }
       const CommunicationClassInfo& targetInfo = targetInfoIt->second;
 
-      for (const ProjectedInputTransferFragment& fragment : inputPlan.fragments) {
+      CommunicationExchangeKind kind = CommunicationExchangeKind::OrdinaryValue;
+      if (sourceInfo.isBatch && !targetInfo.isBatch)
+        kind = CommunicationExchangeKind::PackedScalarRun;
+      else if (sourceInfo.isBatch && targetInfo.isBatch)
+        kind = CommunicationExchangeKind::IndexedBatchRun;
+      else if (!llvm::is_contained(ordinaryProducerDestClasses.lookup(producer), targetClass))
+        continue;
+
+      SmallVector<CommunicationExchangeId, 8>* bucket = nullptr;
+      ProducerTargetKey lookupKey {producer, targetClass};
+      if (kind == CommunicationExchangeKind::OrdinaryValue)
+        bucket = &plan.ordinaryExchangesByProducerTarget[lookupKey];
+      else if (kind == CommunicationExchangeKind::PackedScalarRun)
+        bucket = &plan.packedScalarRunExchangesByProducerTarget[lookupKey];
+      else if (kind == CommunicationExchangeKind::IndexedBatchRun)
+        bucket = &plan.indexedBatchRunExchangesByProducerTarget[lookupKey];
+
+      if (!bucket)
+        continue;
+
+      if (!sourceInfo.isBatch) {
+        if (sourceInfo.cpus.empty()) {
+          emitPlanError(anchor, "communication plan source class has no CPU");
+          return failure();
+        }
+        int64_t sourceCore = static_cast<int64_t>(sourceInfo.cpus.front());
+        unsigned ordinal = 0;
+        for (auto [targetLane, targetCpu] : llvm::enumerate(targetInfo.cpus)) {
+          auto exchangeId = appendExchange(kind,
+                                           producer,
+                                           sourceClass,
+                                           targetClass,
+                                           payloadType,
+                                           sourceCore,
+                                           static_cast<int64_t>(targetCpu),
+                                           0,
+                                           static_cast<unsigned>(targetLane),
+                                           ordinal++);
+          if (failed(exchangeId)) {
+            emitPlanError(anchor, "communication plan failed to append scalar source exchange");
+            return failure();
+          }
+          bucket->push_back(*exchangeId);
+          if (!targetInfo.isBatch)
+            break;
+        }
+        continue;
+      }
+
+      if (sourceInfo.cpus.empty()) {
+        emitPlanError(anchor, "communication plan batch source class has no CPUs");
+        return failure();
+      }
+      if (targetInfo.isBatch && sourceInfo.cpus.size() != targetInfo.cpus.size()) {
+        emitPlanError(anchor, "communication plan cannot map batch classes of different sizes");
+        return failure();
+      }
+
+      for (auto [sourceLane, sourceCpu] : llvm::enumerate(sourceInfo.cpus)) {
+        CpuId targetCpu = targetInfo.isBatch ? targetInfo.cpus[sourceLane] : targetInfo.cpus.front();
+        auto exchangeId = appendExchange(kind,
+                                         producer,
+                                         sourceClass,
+                                         targetClass,
+                                         payloadType,
+                                         static_cast<int64_t>(sourceCpu),
+                                         static_cast<int64_t>(targetCpu),
+                                         static_cast<unsigned>(sourceLane),
+                                         targetInfo.isBatch ? static_cast<unsigned>(sourceLane) : 0,
+                                         static_cast<unsigned>(bucket->size()));
+        if (failed(exchangeId)) {
+          emitPlanError(anchor, "communication plan failed to append batch source exchange");
+          return failure();
+        }
+        bucket->push_back(*exchangeId);
+      }
+    }
+  }
+
+  for (const auto& entry : producerSourceClasses) {
+    ProducerKey producer = entry.first;
+    ClassId sourceClass = entry.second;
+    if (!producer.instance.op || producer.resultIndex >= producer.instance.op->getNumResults())
+      continue;
+
+    Value originalOutput = producer.instance.op->getResult(producer.resultIndex);
+    auto liveIt = liveExternalUseCache.find(originalOutput);
+    if (liveIt == liveExternalUseCache.end() || !liveIt->second)
+      continue;
+
+    auto ownerIt = hostOutputOwners.find(originalOutput);
+    if (ownerIt == hostOutputOwners.end())
+      continue;
+    ClassId ownerClass = ownerIt->second;
+    if (ownerClass == sourceClass)
+      continue;
+
+    auto sourceInfoIt = classById.find(sourceClass);
+    auto ownerInfoIt = classById.find(ownerClass);
+    if (sourceInfoIt == classById.end() || ownerInfoIt == classById.end()) {
+      emitPlanError(anchor, "communication plan references an unknown host-forward class");
+      return failure();
+    }
+    const CommunicationClassInfo& sourceInfo = sourceInfoIt->second;
+    const CommunicationClassInfo& ownerInfo = ownerInfoIt->second;
+    if (sourceInfo.isBatch != ownerInfo.isBatch) {
+      emitPlanError(anchor, "communication plan does not support mixed scalar/batch host forwarding");
+      return failure();
+    }
+    if (sourceInfo.isBatch && sourceInfo.cpus.size() != ownerInfo.cpus.size()) {
+      emitPlanError(anchor, "communication plan cannot host-forward between batch classes of different sizes");
+      return failure();
+    }
+
+    Type payloadType = getPayloadType(producer);
+    if (!payloadType) {
+      emitPlanError(anchor, "communication plan could not determine host-forward payload type");
+      return failure();
+    }
+
+    HostForwardKey lookupKey {originalOutput, sourceClass, ownerClass};
+    SmallVector<CommunicationExchangeId, 8>& bucket = plan.hostForwardExchangesByOutput[lookupKey];
+    if (!sourceInfo.isBatch) {
+      if (!bucket.empty())
+        continue;
+      auto exchangeId = appendExchange(CommunicationExchangeKind::HostForward,
+                                       producer,
+                                       sourceClass,
+                                       ownerClass,
+                                       payloadType,
+                                       static_cast<int64_t>(sourceInfo.cpus.front()),
+                                       static_cast<int64_t>(ownerInfo.cpus.front()),
+                                       0,
+                                       0,
+                                       0,
+                                       originalOutput);
+      if (failed(exchangeId)) {
+        emitPlanError(anchor, "communication plan failed to append scalar host-forward exchange");
+        return failure();
+      }
+      bucket.push_back(*exchangeId);
+      continue;
+    }
+
+    if (bucket.size() >= sourceInfo.cpus.size())
+      continue;
+    unsigned ordinal = static_cast<unsigned>(bucket.size());
+    if (producer.instance.laneStart >= sourceInfo.cpus.size()) {
+      emitPlanError(anchor, "communication plan host-forward producer lane is out of range");
+      return failure();
+    }
+    unsigned lane = producer.instance.laneStart;
+    auto exchangeId = appendExchange(CommunicationExchangeKind::HostForward,
+                                     producer,
+                                     sourceClass,
+                                     ownerClass,
+                                     payloadType,
+                                     static_cast<int64_t>(sourceInfo.cpus[lane]),
+                                     static_cast<int64_t>(ownerInfo.cpus[lane]),
+                                     lane,
+                                     lane,
+                                     ordinal,
+                                     originalOutput);
+    if (failed(exchangeId)) {
+      emitPlanError(anchor, "communication plan failed to append batch host-forward exchange");
+      return failure();
+    }
+    bucket.push_back(*exchangeId);
+  }
+
+  DenseSet<const ProjectedInputTransferFragment*> seenFragments;
+  for (auto& extractEntry : projectedPlans) {
+    for (auto& classEntry : extractEntry.second) {
+      ClassId targetClass = classEntry.first;
+      ProjectedInputTransferPlan& inputPlan = classEntry.second;
+
+      auto targetInfoIt = classById.find(targetClass);
+      if (targetInfoIt == classById.end()) {
+        emitPlanError(anchor, "communication plan references an unknown target class");
+        return failure();
+      }
+      const CommunicationClassInfo& targetInfo = targetInfoIt->second;
+
+      for (ProjectedInputTransferFragment& fragment : inputPlan.fragments) {
         if (!seenFragments.insert(&fragment).second) {
           emitPlanError(anchor, "communication plan found the same projected fragment twice");
           return failure();
@@ -111,129 +463,193 @@ CommunicationPlan::buildProjectedInputPlan(Operation* anchor,
         }
         const CommunicationClassInfo& sourceInfo = sourceInfoIt->second;
 
-        PlaceholderId placeholder {nextPlaceholder++};
-        plan.demands.push_back(InputDemand {placeholder, fragment.producer, targetClass, inputPlan.layout.fragmentType});
-        plan.publications.push_back(
-          ProducerPublication {fragment.producer, sourceClass, inputPlan.layout.fragmentType});
+        plan.demands.push_back(InputDemand {{static_cast<unsigned>(plan.demands.size())},
+                                            fragment.producer,
+                                            targetClass,
+                                            inputPlan.layout.fragmentType});
+        plan.publications.push_back(ProducerPublication {fragment.producer, sourceClass, inputPlan.layout.fragmentType});
 
-        if (fragment.sourceCoreId == fragment.targetCoreId) {
-          plan.phaseByFragment[&fragment] = CommunicationPhase::Local;
+        if (fragment.sourceCoreId == fragment.targetCoreId)
           continue;
-        }
 
         if (sourceInfo.rank == targetInfo.rank && sourceClass != targetClass) {
           emitPlanError(anchor, "communication plan found a same-rank inter-class projected exchange");
           return failure();
         }
 
-        if (!remoteChannels.insert(fragment.channelId).second) {
-          emitPlanError(anchor, "communication plan found a duplicate remote projected channel id");
-          return failure();
-        }
-
-        TransferKind kind = sourceClass == targetClass ? TransferKind::SameClassForward : TransferKind::RemoteChannel;
-        CommunicationPhase phase =
-          sourceInfo.rank == targetInfo.rank
-            ? (fragment.sourceCoreId < fragment.targetCoreId ? CommunicationPhase::LowToHigh
-                                                             : CommunicationPhase::HighToLow)
-            : (sourceInfo.rank < targetInfo.rank ? CommunicationPhase::LowToHigh : CommunicationPhase::HighToLow);
-        plan.phaseByFragment[&fragment] = phase;
-
-        size_t exchangeId = plan.exchanges.size();
-        plan.exchangeByFragment[&fragment] = exchangeId;
-        plan.exchanges.push_back(ExchangeDescriptor {
-          exchangeId,
-          kind,
-          placeholder,
+        FailureOr<CommunicationExchangeId> exchangeId = appendExchange(
+          CommunicationExchangeKind::ProjectedInput,
           fragment.producer,
           sourceClass,
           targetClass,
-          sourceInfo.rank,
-          targetInfo.rank,
           inputPlan.layout.fragmentType,
-          fragment.channelId,
           fragment.sourceCoreId,
           fragment.targetCoreId,
+          0,
           fragment.targetLane,
           fragment.ordinal,
-          phase,
-          &fragment,
-        });
-
-        GroupSignature signature {
-          phase,
-          sourceClass,
-          targetClass,
-          inputPlan.layout.fragmentType,
-          sourceInfo.isBatch,
-          targetInfo.isBatch,
-        };
-        auto groupIt = groupBySignature.find(signature);
-        if (groupIt == groupBySignature.end()) {
-          unsigned groupIndex = static_cast<unsigned>(plan.batchGroups.size());
-          groupBySignature[signature] = groupIndex;
-          plan.batchGroups.push_back(CommunicationBatchGroup {
-            phase,
-            sourceClass,
-            targetClass,
-            inputPlan.layout.fragmentType,
-            sourceInfo.isBatch,
-            targetInfo.isBatch,
-            {},
-          });
-          groupIt = groupBySignature.find(signature);
+          Value(),
+          &fragment);
+        if (failed(exchangeId)) {
+          emitPlanError(anchor, "communication plan failed to append projected input exchange");
+          return failure();
         }
-        plan.batchGroups[groupIt->second].exchanges.push_back(exchangeId);
+        fragment.channelId = plan.exchanges[exchangeId->value].channelId;
+        plan.projectedExchangesByFragment[&fragment].push_back(*exchangeId);
       }
     }
   }
 
-  for (const ExchangeDescriptor& exchange : plan.exchanges) {
-    if (!exchange.payloadType) {
-      emitPlanError(anchor, "communication plan has an exchange without payload type");
-      return failure();
-    }
-    if (exchange.phase == CommunicationPhase::Local) {
-      emitPlanError(anchor, "communication plan classified a remote exchange as local");
-      return failure();
-    }
-    if (exchange.kind == TransferKind::SameClassForward && exchange.sourceClass != exchange.targetClass) {
-      emitPlanError(anchor, "communication plan same-class exchange crosses class ids");
-      return failure();
-    }
-    if (exchange.kind == TransferKind::RemoteChannel && exchange.sourceRank == exchange.targetRank) {
-      emitPlanError(anchor, "communication plan remote exchange has equal source and target rank");
-      return failure();
-    }
-    if (exchange.sourceRank < exchange.targetRank && exchange.phase != CommunicationPhase::LowToHigh) {
-      emitPlanError(anchor, "communication plan phase does not match low-to-high rank direction");
-      return failure();
-    }
-    if (exchange.sourceRank > exchange.targetRank && exchange.phase != CommunicationPhase::HighToLow) {
-      emitPlanError(anchor, "communication plan phase does not match high-to-low rank direction");
-      return failure();
-    }
-    if (!exchange.fragment) {
-      emitPlanError(anchor, "communication plan exchange is not bound to a projected fragment");
-      return failure();
-    }
-  }
-
+  buildEventStreams(plan, plan.exchanges, plan.eventsByClass);
+  if (failed(plan.verify(anchor)))
+    return failure();
   return plan;
 }
 
-std::optional<CommunicationPhase>
-CommunicationPlan::getPhase(const ProjectedInputTransferFragment& fragment) const {
-  auto it = phaseByFragment.find(&fragment);
-  if (it == phaseByFragment.end())
-    return std::nullopt;
+LogicalResult CommunicationPlan::verify(Operation* anchor) const {
+  DenseSet<int64_t> channelIds;
+  SmallVector<const CommunicationExchange*, 8> malformed;
+  size_t malformedCount = 0;
+
+  auto recordMalformed = [&](const CommunicationExchange& exchange) {
+    ++malformedCount;
+    if (malformed.size() < 8)
+      malformed.push_back(&exchange);
+  };
+
+  for (auto [index, exchange] : llvm::enumerate(exchanges)) {
+    bool bad = false;
+    if (exchange.id.value != index)
+      bad = true;
+    if (exchange.sourceCore == exchange.targetCore)
+      bad = true;
+    if (!exchange.payloadType)
+      bad = true;
+    if (!channelIds.insert(exchange.channelId).second)
+      bad = true;
+    if (bad)
+      recordMalformed(exchange);
+  }
+
+  DenseMap<size_t, unsigned> sendCounts;
+  DenseMap<size_t, unsigned> receiveCounts;
+  for (const auto& classEvents : eventsByClass) {
+    std::optional<unsigned> previousOrder;
+    DenseSet<unsigned> classOrders;
+    for (const CommunicationEvent& event : classEvents.second) {
+      if (event.exchangeId.value >= exchanges.size()) {
+        if (anchor)
+          anchor->emitError("communication plan event references an invalid exchange id");
+        return failure();
+      }
+
+      const CommunicationExchange& exchange = exchanges[event.exchangeId.value];
+      if (previousOrder && event.order < *previousOrder) {
+        if (anchor)
+          anchor->emitError("communication plan event stream is not sorted");
+        return failure();
+      }
+      previousOrder = event.order;
+      if (!classOrders.insert(event.order).second) {
+        if (anchor)
+          anchor->emitError("communication plan event stream contains duplicate event order");
+        return failure();
+      }
+
+      if (event.kind == CommunicationEventKind::Send) {
+        ++sendCounts[event.exchangeId.value];
+        if (event.classId != exchange.sourceClass || event.core != exchange.sourceCore)
+          recordMalformed(exchange);
+      }
+      else {
+        ++receiveCounts[event.exchangeId.value];
+        if (event.classId != exchange.targetClass || event.core != exchange.targetCore)
+          recordMalformed(exchange);
+      }
+    }
+  }
+
+  for (const CommunicationExchange& exchange : exchanges) {
+    if (sendCounts[exchange.id.value] != 1 || receiveCounts[exchange.id.value] != 1)
+      recordMalformed(exchange);
+  }
+
+  auto verifyLookup = [&](ArrayRef<CommunicationExchangeId> ids) -> LogicalResult {
+    for (CommunicationExchangeId id : ids)
+      if (id.value >= exchanges.size())
+        return failure();
+    return success();
+  };
+
+  for (const auto& entry : projectedExchangesByFragment)
+    if (failed(verifyLookup(entry.second))) {
+      if (anchor)
+        anchor->emitError("communication plan projected lookup references an invalid exchange id");
+      return failure();
+    }
+
+  if (malformedCount != 0) {
+    if (anchor) {
+      InFlightDiagnostic diag = anchor->emitError("communication plan contains malformed exchanges");
+      for (const CommunicationExchange* exchange : malformed) {
+        diag << " [exchange=" << exchange->id.value << " kind=" << stringifyExchangeKind(exchange->kind)
+             << " sourceClass=" << exchange->sourceClass << " sourceCore=" << exchange->sourceCore
+             << " targetClass=" << exchange->targetClass << " targetCore=" << exchange->targetCore
+             << " channelId=" << exchange->channelId << " payloadType=" << exchange->payloadType << "]";
+      }
+      if (malformedCount > malformed.size())
+        diag << " additionalMalformed=" << (malformedCount - malformed.size());
+    }
+    return failure();
+  }
+
+  return success();
+}
+
+ArrayRef<CommunicationEvent> CommunicationPlan::getEventsForClass(ClassId classId) const {
+  auto it = eventsByClass.find(classId);
+  if (it == eventsByClass.end())
+    return {};
   return it->second;
 }
 
-std::optional<size_t> CommunicationPlan::getExchangeId(const ProjectedInputTransferFragment& fragment) const {
-  auto it = exchangeByFragment.find(&fragment);
-  if (it == exchangeByFragment.end())
-    return std::nullopt;
+ArrayRef<CommunicationExchangeId>
+CommunicationPlan::getOrdinaryValueExchanges(ProducerKey producer, ClassId targetClass) const {
+  auto it = ordinaryExchangesByProducerTarget.find({producer, targetClass});
+  if (it == ordinaryExchangesByProducerTarget.end())
+    return {};
+  return it->second;
+}
+
+ArrayRef<CommunicationExchangeId>
+CommunicationPlan::getProjectedInputExchanges(const ProjectedInputTransferFragment& fragment) const {
+  auto it = projectedExchangesByFragment.find(&fragment);
+  if (it == projectedExchangesByFragment.end())
+    return {};
+  return it->second;
+}
+
+ArrayRef<CommunicationExchangeId>
+CommunicationPlan::getPackedScalarRunExchanges(ProducerKey producer, ClassId targetClass) const {
+  auto it = packedScalarRunExchangesByProducerTarget.find({producer, targetClass});
+  if (it == packedScalarRunExchangesByProducerTarget.end())
+    return {};
+  return it->second;
+}
+
+ArrayRef<CommunicationExchangeId>
+CommunicationPlan::getIndexedBatchRunExchanges(ProducerKey producer, ClassId targetClass) const {
+  auto it = indexedBatchRunExchangesByProducerTarget.find({producer, targetClass});
+  if (it == indexedBatchRunExchangesByProducerTarget.end())
+    return {};
+  return it->second;
+}
+
+ArrayRef<CommunicationExchangeId>
+CommunicationPlan::getHostForwardExchanges(Value originalOutput, ClassId sourceClass, ClassId ownerClass) const {
+  auto it = hostForwardExchangesByOutput.find({originalOutput, sourceClass, ownerClass});
+  if (it == hostForwardExchangesByOutput.end())
+    return {};
   return it->second;
 }
 

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

@@ -5,6 +5,7 @@
 #include "mlir/Support/LLVM.h"
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SmallVector.h"
 
 #include <cstddef>
@@ -24,10 +25,31 @@ enum class TransferKind {
   WholeTensorBarrier
 };
 
-enum class CommunicationPhase {
-  Local,
-  LowToHigh,
-  HighToLow
+enum class CommunicationExchangeKind {
+  OrdinaryValue,
+  ProjectedInput,
+  PackedScalarRun,
+  IndexedBatchRun,
+  HostForward
+};
+
+enum class CommunicationEventKind {
+  Send,
+  Receive
+};
+
+struct CommunicationExchangeId {
+  size_t value = 0;
+
+  bool operator==(const CommunicationExchangeId& other) const { return value == other.value; }
+  operator size_t() const { return value; }
+};
+
+struct CommunicationExchangeIdInfo {
+  static CommunicationExchangeId getEmptyKey() { return {llvm::DenseMapInfo<size_t>::getEmptyKey()}; }
+  static CommunicationExchangeId getTombstoneKey() { return {llvm::DenseMapInfo<size_t>::getTombstoneKey()}; }
+  static unsigned getHashValue(CommunicationExchangeId id) { return llvm::hash_value(id.value); }
+  static bool isEqual(CommunicationExchangeId lhs, CommunicationExchangeId rhs) { return lhs == rhs; }
 };
 
 struct PlaceholderId {
@@ -61,61 +83,147 @@ struct CommunicationClassInfo {
   size_t rank = 0;
   bool isBatch = false;
   mlir::Operation* op = nullptr;
+  llvm::SmallVector<CpuId, 8> cpus;
 };
 
-struct ExchangeDescriptor {
-  size_t id = 0;
-  TransferKind kind = TransferKind::RemoteChannel;
+struct CommunicationExchange {
+  CommunicationExchangeId id;
+  CommunicationExchangeKind kind = CommunicationExchangeKind::OrdinaryValue;
+
   PlaceholderId placeholder;
   ProducerKey producer;
   ClassId sourceClass = 0;
   ClassId targetClass = 0;
+
   size_t sourceRank = 0;
   size_t targetRank = 0;
   mlir::Type payloadType;
   int64_t channelId = 0;
   int64_t sourceCore = 0;
   int64_t targetCore = 0;
+
+  unsigned sourceLane = 0;
   unsigned targetLane = 0;
   unsigned ordinal = 0;
-  CommunicationPhase phase = CommunicationPhase::Local;
+  uint64_t producerOrder = 0;
+
+  mlir::Value originalHostOutput;
   const ProjectedInputTransferFragment* fragment = nullptr;
 };
 
-struct CommunicationBatchGroup {
-  CommunicationPhase phase = CommunicationPhase::Local;
-  ClassId sourceClass = 0;
-  ClassId targetClass = 0;
-  mlir::Type payloadType;
-  bool sourceIsBatch = false;
-  bool targetIsBatch = false;
-  llvm::SmallVector<size_t, 8> exchanges;
+using ExchangeDescriptor = CommunicationExchange;
+
+struct CommunicationEvent {
+  CommunicationExchangeId exchangeId;
+  CommunicationEventKind kind = CommunicationEventKind::Send;
+  ClassId classId = 0;
+  int64_t core = 0;
+  unsigned order = 0;
 };
 
 class CommunicationPlan {
 public:
   using ProjectedInputPlanMap =
     llvm::DenseMap<mlir::Operation*, llvm::DenseMap<ClassId, ProjectedInputTransferPlan>>;
+  using ProducerSourceClassMap = llvm::DenseMap<ProducerKey, ClassId, ProducerKeyInfo>;
 
   static mlir::FailureOr<CommunicationPlan>
-  buildProjectedInputPlan(mlir::Operation* anchor,
-                          llvm::ArrayRef<CommunicationClassInfo> classes,
-                          const llvm::DenseMap<CpuId, ClassId>& cpuToClass,
-                          const ProjectedInputPlanMap& projectedPlans);
+  build(mlir::Operation* anchor,
+        llvm::ArrayRef<CommunicationClassInfo> classes,
+        const llvm::DenseMap<CpuId, ClassId>& cpuToClass,
+        const ProducerSourceClassMap& producerSourceClasses,
+        const llvm::DenseMap<ProducerKey, mlir::Type, ProducerKeyInfo>& producerPayloadTypes,
+        const llvm::DenseMap<ProducerKey, llvm::SmallVector<ClassId, 4>, ProducerKeyInfo>& producerDestClasses,
+        const llvm::DenseMap<ProducerKey, llvm::SmallVector<ClassId, 4>, ProducerKeyInfo>& ordinaryProducerDestClasses,
+        ProjectedInputPlanMap& projectedPlans,
+        const llvm::DenseMap<mlir::Value, ClassId>& hostOutputOwners,
+        const llvm::DenseMap<mlir::Value, bool>& liveExternalUseCache,
+        int64_t& nextChannelId);
 
-  std::optional<CommunicationPhase> getPhase(const ProjectedInputTransferFragment& fragment) const;
-  std::optional<size_t> getExchangeId(const ProjectedInputTransferFragment& fragment) const;
-  const ExchangeDescriptor& getExchange(size_t exchangeId) const { return exchanges[exchangeId]; }
-  llvm::ArrayRef<ExchangeDescriptor> getExchanges() const { return exchanges; }
-  llvm::ArrayRef<CommunicationBatchGroup> getBatchGroups() const { return batchGroups; }
+  mlir::LogicalResult verify(mlir::Operation* anchor) const;
+
+  llvm::ArrayRef<CommunicationEvent> getEventsForClass(ClassId classId) const;
+  llvm::ArrayRef<CommunicationExchange> getExchanges() const { return exchanges; }
+  const CommunicationExchange& getExchange(CommunicationExchangeId id) const { return exchanges[id.value]; }
+  const CommunicationExchange& getExchange(size_t id) const { return exchanges[id]; }
+
+  llvm::ArrayRef<CommunicationExchangeId>
+  getOrdinaryValueExchanges(ProducerKey producer, ClassId targetClass) const;
+  llvm::ArrayRef<CommunicationExchangeId>
+  getProjectedInputExchanges(const ProjectedInputTransferFragment& fragment) const;
+  llvm::ArrayRef<CommunicationExchangeId>
+  getPackedScalarRunExchanges(ProducerKey producer, ClassId targetClass) const;
+  llvm::ArrayRef<CommunicationExchangeId>
+  getIndexedBatchRunExchanges(ProducerKey producer, ClassId targetClass) const;
+  llvm::ArrayRef<CommunicationExchangeId>
+  getHostForwardExchanges(mlir::Value originalOutput, ClassId sourceClass, ClassId ownerClass) const;
 
 private:
-  llvm::DenseMap<const ProjectedInputTransferFragment*, CommunicationPhase> phaseByFragment;
-  llvm::DenseMap<const ProjectedInputTransferFragment*, size_t> exchangeByFragment;
+  struct ProducerTargetKey {
+    ProducerKey producer;
+    ClassId targetClass = 0;
+
+    bool operator==(const ProducerTargetKey& other) const {
+      return producer == other.producer && targetClass == other.targetClass;
+    }
+  };
+
+  struct ProducerTargetKeyInfo {
+    static ProducerTargetKey getEmptyKey() {
+      return {ProducerKeyInfo::getEmptyKey(), llvm::DenseMapInfo<ClassId>::getEmptyKey()};
+    }
+    static ProducerTargetKey getTombstoneKey() {
+      return {ProducerKeyInfo::getTombstoneKey(), llvm::DenseMapInfo<ClassId>::getTombstoneKey()};
+    }
+    static unsigned getHashValue(const ProducerTargetKey& key) {
+      return llvm::hash_combine(ProducerKeyInfo::getHashValue(key.producer), key.targetClass);
+    }
+    static bool isEqual(const ProducerTargetKey& lhs, const ProducerTargetKey& rhs) { return lhs == rhs; }
+  };
+
+  struct HostForwardKey {
+    mlir::Value originalOutput;
+    ClassId sourceClass = 0;
+    ClassId ownerClass = 0;
+
+    bool operator==(const HostForwardKey& other) const {
+      return originalOutput == other.originalOutput && sourceClass == other.sourceClass
+          && ownerClass == other.ownerClass;
+    }
+  };
+
+  struct HostForwardKeyInfo {
+    static HostForwardKey getEmptyKey() {
+      return {llvm::DenseMapInfo<mlir::Value>::getEmptyKey(),
+              llvm::DenseMapInfo<ClassId>::getEmptyKey(),
+              llvm::DenseMapInfo<ClassId>::getEmptyKey()};
+    }
+    static HostForwardKey getTombstoneKey() {
+      return {llvm::DenseMapInfo<mlir::Value>::getTombstoneKey(),
+              llvm::DenseMapInfo<ClassId>::getTombstoneKey(),
+              llvm::DenseMapInfo<ClassId>::getTombstoneKey()};
+    }
+    static unsigned getHashValue(const HostForwardKey& key) {
+      return llvm::hash_combine(key.originalOutput, key.sourceClass, key.ownerClass);
+    }
+    static bool isEqual(const HostForwardKey& lhs, const HostForwardKey& rhs) { return lhs == rhs; }
+  };
+
   llvm::SmallVector<InputDemand, 16> demands;
   llvm::SmallVector<ProducerPublication, 16> publications;
-  llvm::SmallVector<ExchangeDescriptor, 16> exchanges;
-  llvm::SmallVector<CommunicationBatchGroup, 8> batchGroups;
+  llvm::SmallVector<CommunicationExchange, 16> exchanges;
+  llvm::DenseMap<ClassId, llvm::SmallVector<CommunicationEvent, 16>> eventsByClass;
+
+  llvm::DenseMap<const ProjectedInputTransferFragment*, llvm::SmallVector<CommunicationExchangeId, 1>>
+    projectedExchangesByFragment;
+  llvm::DenseMap<ProducerTargetKey, llvm::SmallVector<CommunicationExchangeId, 8>, ProducerTargetKeyInfo>
+    ordinaryExchangesByProducerTarget;
+  llvm::DenseMap<ProducerTargetKey, llvm::SmallVector<CommunicationExchangeId, 8>, ProducerTargetKeyInfo>
+    packedScalarRunExchangesByProducerTarget;
+  llvm::DenseMap<ProducerTargetKey, llvm::SmallVector<CommunicationExchangeId, 8>, ProducerTargetKeyInfo>
+    indexedBatchRunExchangesByProducerTarget;
+  llvm::DenseMap<HostForwardKey, llvm::SmallVector<CommunicationExchangeId, 8>, HostForwardKeyInfo>
+    hostForwardExchangesByOutput;
 };
 
 } // namespace onnx_mlir::spatial

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

@@ -58,7 +58,7 @@ struct PackedScalarRunValue {
 
   mlir::RankedTensorType fragmentType;
   llvm::SmallVector<PackedScalarRunSlot, 8> slots;
-  MessageVector messages;
+  llvm::SmallVector<CommunicationExchangeId, 8> exchangeIds;
 };
 
 struct IndexedBatchRunValue {
@@ -68,7 +68,7 @@ struct IndexedBatchRunValue {
   mlir::Value packed;
   mlir::RankedTensorType fragmentType;
   llvm::SmallVector<PackedScalarRunSlot, 8> slots;
-  MessageVector messages;
+  llvm::SmallVector<CommunicationExchangeId, 8> exchangeIds;
 };
 
 struct LogicalSlotRange {
@@ -90,12 +90,13 @@ struct OutputDestinationGroup {
 struct BatchRunSendPlan {
   size_t resultIndex = 0;
   ClassId destinationClass = 0;
-  MessageVector messages;
+  llvm::SmallVector<CommunicationExchangeId, 8> exchangeIds;
 };
 
 enum class TensorDemandActionKind {
   DestinationFanout,
   SameClassIndexedFragment,
+  ProjectedInputPublication,
   TerminalBlueprintPublication,
   WholeTensorBarrier
 };
@@ -122,44 +123,6 @@ struct CompactRunPlan {
   llvm::SmallVector<RunOutputDemand, 4> outputs;
 };
 
-struct ScalarPeerEdgeKey {
-  int64_t sourceCore = 0;
-  int64_t targetCore = 0;
-  mlir::Type payloadType;
-};
-
-struct ScalarPeerReceiveKey {
-  int64_t sourceCore = 0;
-  int64_t targetCore = 0;
-  mlir::Type payloadType;
-  std::optional<int64_t> channelId;
-};
-
-struct PendingScalarSend {
-  ClassId sourceClass = 0;
-  int64_t sourceCore = 0;
-  int64_t targetCore = 0;
-  mlir::Type payloadType;
-  ScalarPeerEdgeKey waitForReceive;
-  mlir::Operation* payloadAnchor = nullptr;
-  mlir::Value payload;
-  MessageVector messages;
-  mlir::Location loc;
-};
-
-struct PendingProjectedScalarSend {
-  ClassId sourceClass = 0;
-  int64_t sourceCore = 0;
-  int64_t targetCore = 0;
-  mlir::Type payloadType;
-  ScalarPeerEdgeKey waitForReceive;
-  mlir::Operation* payloadAnchor = nullptr;
-  mlir::Value payload;
-  MessageVector messages;
-  ProjectedTransferDescriptor descriptor;
-  mlir::Location loc;
-};
-
 enum class BatchInputDemandKind {
   LaneFragment,
   ProjectedFragment,
@@ -171,11 +134,17 @@ struct BatchInputDemand {
   std::optional<ProducerKey> wholeTensorProducer;
 };
 
-struct ProjectedExchangeEmissionState {
+struct CommunicationExchangeEmissionState {
   bool producerBound = false;
   bool sendEmitted = false;
   bool receiveEmitted = false;
+
   mlir::Value producerPayload;
+  mlir::Value receiveValue;
+
+  mlir::Operation* producerAnchor = nullptr;
+  mlir::Operation* sendAnchor = nullptr;
+  mlir::Operation* receiveAnchor = nullptr;
   std::optional<mlir::Location> loc;
 };
 
@@ -220,6 +189,24 @@ public:
   }
 
   void recordIndexedBatchRun(IndexedBatchRunValue run) { indexedBatchRuns.push_back(std::move(run)); }
+  void recordDeferredExact(ProducerKey key, ClassId classId, llvm::ArrayRef<CommunicationExchangeId> exchangeIds) {
+    llvm::SmallVector<CommunicationExchangeId, 4>& bucket = exactExchangeIds[key][classId];
+    for (CommunicationExchangeId id : exchangeIds)
+      if (!llvm::is_contained(bucket, id))
+        bucket.push_back(id);
+  }
+  llvm::ArrayRef<CommunicationExchangeId> getDeferredExactExchangeIds(ProducerKey key, ClassId classId) const {
+    auto producerIt = exactExchangeIds.find(key);
+    if (producerIt == exactExchangeIds.end())
+      return {};
+    auto classIt = producerIt->second.find(classId);
+    if (classIt == producerIt->second.end())
+      return {};
+    return classIt->second;
+  }
+  bool hasDeferredExact(ProducerKey key, ClassId classId) const {
+    return !getDeferredExactExchangeIds(key, classId).empty();
+  }
 
   std::optional<mlir::Value> lookupExact(ProducerKey key, ClassId classId) const;
   std::optional<mlir::Value> lookup(MaterializerState& state, ProducerKey key, ClassId classId);
@@ -245,6 +232,10 @@ private:
   std::optional<mlir::Value> lookupPackedRun(MaterializerState& state, ProducerKey key, ClassId classId);
 
   llvm::DenseMap<ProducerKey, llvm::DenseMap<ClassId, mlir::Value>, ProducerKeyInfo> exactValues;
+  llvm::DenseMap<ProducerKey,
+                 llvm::DenseMap<ClassId, llvm::SmallVector<CommunicationExchangeId, 4>>,
+                 ProducerKeyInfo>
+    exactExchangeIds;
   llvm::SmallVector<PackedScalarRunValue, 8> packedScalarRuns;
   llvm::SmallVector<IndexedBatchRunValue, 8> indexedBatchRuns;
   llvm::DenseMap<WholeBatchAssemblyLookupKey,
@@ -285,19 +276,13 @@ struct MaterializerState {
     projectedExtractReplacements;
   llvm::DenseMap<mlir::Operation*, llvm::DenseMap<ClassId, ProjectedInputTransferPlan>>
     projectedInputTransferPlans;
-  CommunicationPlan projectedInputCommunicationPlan;
-  llvm::SmallVector<ProjectedExchangeEmissionState, 16> projectedExchangeStates;
-  CommunicationPhase currentProjectedCommunicationPhase = CommunicationPhase::LowToHigh;
-  llvm::SmallVector<size_t, 8> deferredProjectedExchanges;
+  CommunicationPlan communicationPlan;
+  llvm::SmallVector<CommunicationExchangeEmissionState, 16> communicationExchangeStates;
   llvm::SmallVector<ProducerKey, 8> deferredProducerValues;
   AvailableValueStore availableValues;
   llvm::DenseMap<mlir::Value, mlir::Value> hostReplacements;
   llvm::DenseMap<mlir::Value, ClassId> hostOutputOwners;
   llvm::SmallVector<PendingProjectedHostOutputFragment, 32> pendingProjectedHostOutputFragments;
-  llvm::SmallVector<ScalarPeerEdgeKey, 16> plannedScalarPeerReceives;
-  llvm::SmallVector<ScalarPeerEdgeKey, 8> materializedScalarPeerReceives;
-  llvm::SmallVector<PendingScalarSend, 8> pendingScalarSends;
-  llvm::SmallVector<PendingProjectedScalarSend, 8> pendingProjectedScalarSends;
   llvm::DenseSet<mlir::Operation*> oldComputeOps;
 
   MaterializerState(mlir::func::FuncOp func, const MergeScheduleResult& schedule, int64_t& nextChannelId)