CommunicationPlan

src/PIM/Dialect/Spatial/CMakeLists.txt

@@ -8,6 +8,7 @@ add_pim_library(SpatialOps
   SpatialOpsCanonicalization.cpp
   ${PIM_SRC_ROOT}/Conversion/ONNXToSpatial/CompileTime.cpp
   Transforms/MergeComputeNodes/MergeComputeNodesPass.cpp
+  Transforms/MergeComputeNodes/CommunicationPlan.cpp
   Transforms/MergeComputeNodes/HostOutputFinalization.cpp
   Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp
   Transforms/MergeComputeNodes/ProjectedFragments.cpp

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

@@ -0,0 +1,240 @@
+#include "CommunicationPlan.hpp"
+
+#include "mlir/IR/Diagnostics.h"
+
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/TypeSwitch.h"
+
+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;
+  }
+};
+
+struct GroupSignatureInfo {
+  static GroupSignature getEmptyKey() {
+    return {CommunicationPhase::Local,
+            llvm::DenseMapInfo<ClassId>::getEmptyKey(),
+            llvm::DenseMapInfo<ClassId>::getEmptyKey(),
+            llvm::DenseMapInfo<Type>::getEmptyKey(),
+            false,
+            false};
+  }
+
+  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; }
+};
+
+void emitPlanError(Operation* anchor, StringRef message) {
+  if (anchor)
+    anchor->emitError(message);
+}
+
+} // namespace
+
+FailureOr<CommunicationPlan>
+CommunicationPlan::buildProjectedInputPlan(Operation* anchor,
+                                           ArrayRef<CommunicationClassInfo> classes,
+                                           const DenseMap<CpuId, ClassId>& cpuToClass,
+                                           const ProjectedInputPlanMap& projectedPlans) {
+  CommunicationPlan plan;
+
+  DenseMap<ClassId, CommunicationClassInfo> classById;
+  for (CommunicationClassInfo klass : classes)
+    classById[klass.classId] = klass;
+
+  DenseSet<int64_t> remoteChannels;
+  DenseSet<const ProjectedInputTransferFragment*> seenFragments;
+  DenseMap<GroupSignature, unsigned, GroupSignatureInfo> groupBySignature;
+  unsigned nextPlaceholder = 0;
+
+  for (const auto& extractEntry : projectedPlans) {
+    for (const auto& classEntry : extractEntry.second) {
+      ClassId targetClass = classEntry.first;
+      const 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 (const ProjectedInputTransferFragment& fragment : inputPlan.fragments) {
+        if (!seenFragments.insert(&fragment).second) {
+          emitPlanError(anchor, "communication plan found the same projected fragment twice");
+          return failure();
+        }
+
+        auto sourceClassIt = cpuToClass.find(fragment.sourceCoreId);
+        if (sourceClassIt == cpuToClass.end()) {
+          emitPlanError(anchor, "communication plan could not map projected source core to a class");
+          return failure();
+        }
+
+        ClassId sourceClass = sourceClassIt->second;
+        auto sourceInfoIt = classById.find(sourceClass);
+        if (sourceInfoIt == classById.end()) {
+          emitPlanError(anchor, "communication plan references an unknown source class");
+          return failure();
+        }
+        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});
+
+        if (fragment.sourceCoreId == fragment.targetCoreId) {
+          plan.phaseByFragment[&fragment] = CommunicationPhase::Local;
+          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,
+          fragment.producer,
+          sourceClass,
+          targetClass,
+          sourceInfo.rank,
+          targetInfo.rank,
+          inputPlan.layout.fragmentType,
+          fragment.channelId,
+          fragment.sourceCoreId,
+          fragment.targetCoreId,
+          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);
+        }
+        plan.batchGroups[groupIt->second].exchanges.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();
+    }
+  }
+
+  return plan;
+}
+
+std::optional<CommunicationPhase>
+CommunicationPlan::getPhase(const ProjectedInputTransferFragment& fragment) const {
+  auto it = phaseByFragment.find(&fragment);
+  if (it == phaseByFragment.end())
+    return std::nullopt;
+  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;
+  return it->second;
+}
+
+} // namespace onnx_mlir::spatial

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

@@ -0,0 +1,121 @@
+#pragma once
+
+#include "mlir/IR/Operation.h"
+#include "mlir/IR/Types.h"
+#include "mlir/Support/LLVM.h"
+
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallVector.h"
+
+#include <cstddef>
+#include <cstdint>
+#include <optional>
+
+#include "MergeScheduleKeys.hpp"
+#include "ProjectedFragments.hpp"
+
+namespace onnx_mlir::spatial {
+
+enum class TransferKind {
+  DirectValue,
+  LocalFragment,
+  SameClassForward,
+  RemoteChannel,
+  WholeTensorBarrier
+};
+
+enum class CommunicationPhase {
+  Local,
+  LowToHigh,
+  HighToLow
+};
+
+struct PlaceholderId {
+  unsigned value = 0;
+
+  bool operator==(const PlaceholderId& other) const { return value == other.value; }
+};
+
+struct PlaceholderIdInfo {
+  static PlaceholderId getEmptyKey() { return {static_cast<unsigned>(-1)}; }
+  static PlaceholderId getTombstoneKey() { return {static_cast<unsigned>(-2)}; }
+  static unsigned getHashValue(PlaceholderId id) { return llvm::hash_value(id.value); }
+  static bool isEqual(PlaceholderId lhs, PlaceholderId rhs) { return lhs == rhs; }
+};
+
+struct InputDemand {
+  PlaceholderId placeholder;
+  ProducerKey producer;
+  ClassId targetClass = 0;
+  mlir::Type requiredType;
+};
+
+struct ProducerPublication {
+  ProducerKey producer;
+  ClassId sourceClass = 0;
+  mlir::Type payloadType;
+};
+
+struct CommunicationClassInfo {
+  ClassId classId = 0;
+  size_t rank = 0;
+  bool isBatch = false;
+  mlir::Operation* op = nullptr;
+};
+
+struct ExchangeDescriptor {
+  size_t id = 0;
+  TransferKind kind = TransferKind::RemoteChannel;
+  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 targetLane = 0;
+  unsigned ordinal = 0;
+  CommunicationPhase phase = CommunicationPhase::Local;
+  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;
+};
+
+class CommunicationPlan {
+public:
+  using ProjectedInputPlanMap =
+    llvm::DenseMap<mlir::Operation*, llvm::DenseMap<ClassId, ProjectedInputTransferPlan>>;
+
+  static mlir::FailureOr<CommunicationPlan>
+  buildProjectedInputPlan(mlir::Operation* anchor,
+                          llvm::ArrayRef<CommunicationClassInfo> classes,
+                          const llvm::DenseMap<CpuId, ClassId>& cpuToClass,
+                          const ProjectedInputPlanMap& projectedPlans);
+
+  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; }
+
+private:
+  llvm::DenseMap<const ProjectedInputTransferFragment*, CommunicationPhase> phaseByFragment;
+  llvm::DenseMap<const ProjectedInputTransferFragment*, size_t> exchangeByFragment;
+  llvm::SmallVector<InputDemand, 16> demands;
+  llvm::SmallVector<ProducerPublication, 16> publications;
+  llvm::SmallVector<ExchangeDescriptor, 16> exchanges;
+  llvm::SmallVector<CommunicationBatchGroup, 8> batchGroups;
+};
+
+} // namespace onnx_mlir::spatial

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

@@ -12,6 +12,7 @@
 #include <optional>
 
 #include "MaterializeMergeSchedule.hpp"
+#include "CommunicationPlan.hpp"
 #include "MergeMessages.hpp"
 #include "MergeScheduleKeys.hpp"
 #include "ProjectedFragments.hpp"
@@ -159,13 +160,6 @@ struct PendingProjectedScalarSend {
   mlir::Location loc;
 };
 
-struct PendingProjectedInputSend {
-  ClassId sourceClass = 0;
-  mlir::Value payload;
-  llvm::SmallVector<ProjectedInputTransferFragment*, 4> fragments;
-  mlir::Location loc;
-};
-
 enum class BatchInputDemandKind {
   LaneFragment,
   ProjectedFragment,
@@ -177,6 +171,14 @@ struct BatchInputDemand {
   std::optional<ProducerKey> wholeTensorProducer;
 };
 
+struct ProjectedExchangeEmissionState {
+  bool producerBound = false;
+  bool sendEmitted = false;
+  bool receiveEmitted = false;
+  mlir::Value producerPayload;
+  std::optional<mlir::Location> loc;
+};
+
 struct CloneIndexingContext {
   std::optional<mlir::Value> runSlotIndex;
   std::optional<mlir::Value> projectionSlotIndex;
@@ -267,6 +269,7 @@ struct MaterializerState {
   llvm::DenseSet<ClassSlotKey> materializedLogicalSlots;
 
   llvm::DenseMap<ProducerKey, llvm::SmallVector<ClassId, 4>, ProducerKeyInfo> producerDestClasses;
+  llvm::DenseMap<ProducerKey, llvm::SmallVector<ClassId, 4>, ProducerKeyInfo> ordinaryProducerDestClasses;
   llvm::DenseMap<SameClassConsumerLookupKey, llvm::SmallVector<ProducerKey, 4>, SameClassConsumerLookupKeyInfo>
     sameClassConsumerIndex;
   llvm::DenseMap<ProjectedBatchInputKey, AffineProjectedInputSliceMatch, ProjectedBatchInputKeyInfo>
@@ -282,6 +285,11 @@ 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;
+  llvm::SmallVector<ProducerKey, 8> deferredProducerValues;
   AvailableValueStore availableValues;
   llvm::DenseMap<mlir::Value, mlir::Value> hostReplacements;
   llvm::DenseMap<mlir::Value, ClassId> hostOutputOwners;
@@ -290,9 +298,6 @@ struct MaterializerState {
   llvm::SmallVector<ScalarPeerEdgeKey, 8> materializedScalarPeerReceives;
   llvm::SmallVector<PendingScalarSend, 8> pendingScalarSends;
   llvm::SmallVector<PendingProjectedScalarSend, 8> pendingProjectedScalarSends;
-  llvm::SmallVector<PendingProjectedInputSend, 16> pendingProjectedInputSends;
-  llvm::DenseSet<ClassId> projectedInputPhaseBarrierClasses;
-  llvm::DenseMap<ClassId, unsigned> pendingProjectedHighToLowReceives;
   llvm::DenseSet<mlir::Operation*> oldComputeOps;
 
   MaterializerState(mlir::func::FuncOp func, const MergeScheduleResult& schedule, int64_t& nextChannelId)