From 47f6715296fc53c25f591c17179aab7ad885c77b Mon Sep 17 00:00:00 2001 From: ilgeco Date: Mon, 6 Jul 2026 17:25:31 +0200 Subject: [PATCH] CommunicationPlan --- src/PIM/Dialect/Spatial/CMakeLists.txt | 1 + .../MergeComputeNodes/CommunicationPlan.cpp | 240 +++ .../MergeComputeNodes/CommunicationPlan.hpp | 121 ++ .../MaterializeMergeSchedule.cpp | 1387 ++++++++++++----- .../MaterializedClassState.hpp | 25 +- 5 files changed, 1360 insertions(+), 414 deletions(-) create mode 100644 src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.cpp create mode 100644 src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.hpp diff --git a/src/PIM/Dialect/Spatial/CMakeLists.txt b/src/PIM/Dialect/Spatial/CMakeLists.txt index 4de8cce..0ca7852 100644 --- a/src/PIM/Dialect/Spatial/CMakeLists.txt +++ b/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 diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.cpp new file mode 100644 index 0000000..88fba20 --- /dev/null +++ b/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::getEmptyKey(), + llvm::DenseMapInfo::getEmptyKey(), + llvm::DenseMapInfo::getEmptyKey(), + false, + false}; + } + + static GroupSignature getTombstoneKey() { + return {CommunicationPhase::Local, + llvm::DenseMapInfo::getTombstoneKey(), + llvm::DenseMapInfo::getTombstoneKey(), + llvm::DenseMapInfo::getTombstoneKey(), + false, + false}; + } + + static unsigned getHashValue(const GroupSignature& key) { + return llvm::hash_combine(static_cast(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::buildProjectedInputPlan(Operation* anchor, + ArrayRef classes, + const DenseMap& cpuToClass, + const ProjectedInputPlanMap& projectedPlans) { + CommunicationPlan plan; + + DenseMap classById; + for (CommunicationClassInfo klass : classes) + classById[klass.classId] = klass; + + DenseSet remoteChannels; + DenseSet seenFragments; + DenseMap 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(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 +CommunicationPlan::getPhase(const ProjectedInputTransferFragment& fragment) const { + auto it = phaseByFragment.find(&fragment); + if (it == phaseByFragment.end()) + return std::nullopt; + return it->second; +} + +std::optional 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 diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.hpp new file mode 100644 index 0000000..509dedb --- /dev/null +++ b/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 +#include +#include + +#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(-1)}; } + static PlaceholderId getTombstoneKey() { return {static_cast(-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 exchanges; +}; + +class CommunicationPlan { +public: + using ProjectedInputPlanMap = + llvm::DenseMap>; + + static mlir::FailureOr + buildProjectedInputPlan(mlir::Operation* anchor, + llvm::ArrayRef classes, + const llvm::DenseMap& cpuToClass, + const ProjectedInputPlanMap& projectedPlans); + + std::optional getPhase(const ProjectedInputTransferFragment& fragment) const; + std::optional getExchangeId(const ProjectedInputTransferFragment& fragment) const; + const ExchangeDescriptor& getExchange(size_t exchangeId) const { return exchanges[exchangeId]; } + llvm::ArrayRef getExchanges() const { return exchanges; } + llvm::ArrayRef getBatchGroups() const { return batchGroups; } + +private: + llvm::DenseMap phaseByFragment; + llvm::DenseMap exchangeByFragment; + llvm::SmallVector demands; + llvm::SmallVector publications; + llvm::SmallVector exchanges; + llvm::SmallVector batchGroups; +}; + +} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp index f710b66..bac6c85 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp @@ -117,7 +117,8 @@ LogicalResult emitPlannedProjectedInputSendsForKeys(MaterializerState& state, MaterializedClass& sourceClass, ArrayRef keys, Value payload, - Location loc); + Location loc); +static void tagProjectedInputTransfer(Operation* op); std::optional getProjectedInputSliceMatch(MaterializerState& state, SpatComputeBatch batch, unsigned inputIndex); std::optional @@ -464,6 +465,29 @@ bool canUseProjectedLaneInput(MaterializerState& state, return !collectProjectedFragmentDemandsForBatchInput(state, consumerBatch, inputIndex, input, logicalConsumer).empty(); } +static bool hasUnreplacedBatchInputUse(MaterializerState& state, + SpatComputeBatch consumerBatch, + unsigned inputIndex, + ClassId classId) { + auto inputArg = consumerBatch.getInputArgument(inputIndex); + if (!inputArg) + return true; + + for (OpOperand& use : inputArg->getUses()) { + auto extract = dyn_cast(use.getOwner()); + if (!extract) + return true; + + auto planIt = state.projectedInputTransferPlans.find(extract.getOperation()); + if (planIt != state.projectedInputTransferPlans.end() && planIt->second.find(classId) != planIt->second.end()) + continue; + if (getProjectedWholeBatchReplacementProducer(state, extract)) + continue; + return true; + } + return false; +} + FailureOr classifyComputeBatchInputDemand(MaterializerState& state, MaterializedClass& targetClass, SpatComputeBatch consumerBatch, @@ -471,6 +495,10 @@ FailureOr classifyComputeBatchInputDemand(MaterializerState& s Value input, ComputeInstance logicalConsumer) { if (std::optional wholeBatchProducer = getWholeBatchProducerKeyForDirectBatchResult(input)) { + if (hasUnreplacedBatchInputUse(state, consumerBatch, inputIndex, targetClass.id)) + return BatchInputDemand {.kind = BatchInputDemandKind::WholeTensorBarrier, + .wholeTensorProducer = wholeBatchProducer}; + if (canUseProjectedLaneInput(state, consumerBatch, inputIndex, input, logicalConsumer)) return BatchInputDemand {.kind = BatchInputDemandKind::ProjectedFragment, .wholeTensorProducer = std::nullopt}; @@ -606,14 +634,6 @@ static int64_t getProjectedInputExchangeDirection(int64_t sourceCoreId, int64_t return sourceCoreId < targetCoreId ? 1 : 2; } -static bool isLowerToHigherProjectedInputFragment(const ProjectedInputTransferFragment& fragment) { - return fragment.sourceCoreId < fragment.targetCoreId; -} - -static bool isHigherToLowerProjectedInputFragment(const ProjectedInputTransferFragment& fragment) { - return fragment.sourceCoreId > fragment.targetCoreId; -} - static LogicalResult appendProjectedBatchingSignatureForCpu(MaterializerState& state, CpuId cpu, MaterializedBatchingSignature& signature) { @@ -670,9 +690,9 @@ static LogicalResult appendOrdinaryInputBatchingSignatureForCpu(MaterializerStat for (auto [inputIndex, input] : llvm::enumerate(consumerInputs)) { SmallVector producerKeys; if (auto batchConsumer = dyn_cast(logicalConsumer.op)) { - SmallVector projectedDemands = collectProjectedFragmentDemandsForBatchInput( - state, batchConsumer, static_cast(inputIndex), input, logicalConsumer); - if (!projectedDemands.empty()) + if (canUseProjectedLaneInput(state, batchConsumer, static_cast(inputIndex), input, logicalConsumer) + && !hasUnreplacedBatchInputUse(state, batchConsumer, static_cast(inputIndex), + state.cpuToClass.lookup(cpu))) continue; producerKeys = collectProducerKeysForBatchInputDestinations( state, batchConsumer, static_cast(inputIndex), input, logicalConsumer); @@ -2825,10 +2845,16 @@ bool isTerminalHostBatchOutput(Value output, const DenseSet& oldComp return !hasRealComputeConsumer(output, oldComputeOps); } -void appendDestinationClass(MaterializerState& state, ProducerKey key, ClassId classId) { +void appendDestinationClass(MaterializerState& state, ProducerKey key, ClassId classId, bool ordinary = true) { SmallVector& destinations = state.producerDestClasses[key]; if (!llvm::is_contained(destinations, classId)) destinations.push_back(classId); + + if (!ordinary) + return; + SmallVector& ordinaryDestinations = state.ordinaryProducerDestClasses[key]; + if (!llvm::is_contained(ordinaryDestinations, classId)) + ordinaryDestinations.push_back(classId); } void replaceLiveExternalUses(Value oldValue, Value replacement, const DenseSet& oldComputeOps) { @@ -2856,7 +2882,9 @@ LogicalResult collectProducerDestinations(MaterializerState& state) { if (auto batchConsumer = dyn_cast(logicalConsumer.op)) { SmallVector projectedKeys = collectProjectedProducerKeysForBatchInput( state, batchConsumer, static_cast(inputIndex), input, logicalConsumer); - projectedPlanOwnsFanout = !state.classes[targetClass].isBatch && !projectedKeys.empty(); + projectedPlanOwnsFanout = + !state.classes[targetClass].isBatch && !projectedKeys.empty() + && !hasUnreplacedBatchInputUse(state, batchConsumer, static_cast(inputIndex), targetClass); producerKeys = projectedPlanOwnsFanout ? std::move(projectedKeys) : collectProducerKeysForBatchInputDestinations( @@ -3674,9 +3702,10 @@ LogicalResult collectProjectedTransfers(MaterializerState& state) { } for (auto [inputIndex, input] : llvm::enumerate(batch.getInputs())) { + ProjectedBatchInputKey currentInputKey {batch.getOperation(), static_cast(inputIndex)}; SmallVector fragmentDemands = collectProjectedFragmentDemandsForBatchInput( - state, batch, static_cast(inputIndex), input, logicalConsumer); + state, batch, static_cast(inputIndex), input, logicalConsumer); SmallVector matches = collectAffineProjectedInputSlices(batch, static_cast(inputIndex)); @@ -3689,7 +3718,6 @@ LogicalResult collectProjectedTransfers(MaterializerState& state) { ProjectedInputTransferPlan& plan = state.projectedInputTransferPlans[match.extract.getOperation()][targetClassId]; - ProjectedBatchInputKey currentInputKey {batch.getOperation(), static_cast(inputIndex)}; if (plan.fragments.empty()) { plan.inputKey = currentInputKey; plan.extractOp = match.extract.getOperation(); @@ -3737,10 +3765,6 @@ LogicalResult collectProjectedTransfers(MaterializerState& state) { fragment.channelId = state.nextChannelId++; fragment.sourceCoreId = *sourceCpu; fragment.targetCoreId = *targetCpu; - if (fragment.sourceCoreId > fragment.targetCoreId) { - state.projectedInputPhaseBarrierClasses.insert(targetClassId); - ++state.pendingProjectedHighToLowReceives[targetClassId]; - } plan.fragments.push_back(std::move(fragment)); } } @@ -3795,7 +3819,6 @@ LogicalResult collectProjectedTransfers(MaterializerState& state) { } } - ProjectedBatchInputKey currentInputKey {batch.getOperation(), static_cast(inputIndex)}; if (!(descriptor.inputKey == currentInputKey) || descriptor.extractOp != match->extract.getOperation() || descriptor.sourceType != match->sourceType || descriptor.fragmentType != match->fragmentType || descriptor.fragmentShape != match->fragmentShape @@ -3895,6 +3918,28 @@ LogicalResult collectProjectedTransfers(MaterializerState& state) { return success(); } +LogicalResult buildProjectedInputCommunicationPlan(MaterializerState& state) { + SmallVector classes; + classes.reserve(state.classes.size()); + for (const MaterializedClass& klass : state.classes) { + classes.push_back(CommunicationClassInfo { + .classId = klass.id, + .rank = klass.id, + .isBatch = klass.isBatch, + .op = klass.op, + }); + } + + FailureOr plan = CommunicationPlan::buildProjectedInputPlan( + state.func, classes, state.cpuToClass, state.projectedInputTransferPlans); + if (failed(plan)) + return failure(); + state.projectedInputCommunicationPlan = std::move(*plan); + state.projectedExchangeStates.assign(state.projectedInputCommunicationPlan.getExchanges().size(), + ProjectedExchangeEmissionState {}); + return success(); +} + static std::optional collectScalarTargetProjectedDescriptor(MaterializerState& state, MaterializedClass& targetClass, @@ -3963,6 +4008,19 @@ ArrayRef getDestinationClasses(MaterializerState& state, ProducerKey ke return it->second; } +static bool hasOrdinaryDestinationClass(MaterializerState& state, ProducerKey key, ClassId classId) { + auto it = state.ordinaryProducerDestClasses.find(key); + return it != state.ordinaryProducerDestClasses.end() && llvm::is_contained(it->second, classId); +} + +static bool anyKeyRequiresOrdinaryDestination(MaterializerState& state, + ArrayRef keys, + ClassId classId) { + return llvm::any_of(keys, [&](ProducerKey key) { + return hasOrdinaryDestinationClass(state, key, classId); + }); +} + SmallVector filterProducerKeysForDestination(MaterializerState& state, ArrayRef keys, ClassId destinationClass) { @@ -4483,6 +4541,8 @@ bool peerExchangeOrderLess(PeerExchangeOrderKey lhs, PeerExchangeOrderKey rhs) { return lhs.channelId < rhs.channelId; } +static void tagProjectedInputTransfer(Operation* op, std::optional key); + void sortMessageIndicesByPeerExchangeOrder(const MessageVector& messages, SmallVectorImpl& indices) { llvm::stable_sort(indices, [&](size_t lhs, size_t rhs) { return peerExchangeOrderLess(makePeerExchangeOrderKey(messages, lhs), makePeerExchangeOrderKey(messages, rhs)); @@ -4670,7 +4730,40 @@ FailureOr buildScalarSourceFanoutPlan(MaterializerState& receivePlan.messages.channelIds, receivePlan.messages.sourceCoreIds, receivePlan.messages.targetCoreIds); } + bool addOrdinaryPayloadForProjectedDestination = + *descriptor && anyKeyRequiresOrdinaryDestination(state, keys, destinationClass); fanoutPlan.receivePlans.push_back(std::move(receivePlan)); + + if (addOrdinaryPayloadForProjectedDestination) { + ScalarSourceReceivePlan ordinaryPlan; + ordinaryPlan.targetClass = destinationClass; + ordinaryPlan.receiveType = payload.getType(); + + auto appendOrdinaryMessage = [&](CpuId targetCpu) -> LogicalResult { + auto checkedTargetCpu = getCheckedCoreId(targetClass.op, targetCpu, "scalar ordinary target core id"); + if (failed(checkedTargetCpu)) + return failure(); + ordinaryPlan.messages.append(state.nextChannelId++, *sourceCpu, *checkedTargetCpu); + return success(); + }; + + if (!targetClass.isBatch) { + if (failed(appendOrdinaryMessage(targetClass.cpus.front()))) + return failure(); + } + else { + for (CpuId targetCpu : targetClass.cpus) + if (failed(appendOrdinaryMessage(targetCpu))) + return failure(); + } + + if (!fanoutPlan.ordinaryMessages) + fanoutPlan.ordinaryMessages = MessageVector {}; + fanoutPlan.ordinaryMessages->append(ordinaryPlan.messages.channelIds, + ordinaryPlan.messages.sourceCoreIds, + ordinaryPlan.messages.targetCoreIds); + fanoutPlan.receivePlans.push_back(std::move(ordinaryPlan)); + } } for (ProjectedScalarSendGroup& group : fanoutPlan.projectedSendGroups) { @@ -5127,36 +5220,15 @@ LogicalResult appendConditionalBatchProjectedInputSendsAtCurrentInsertionPoint(M Value channelId = createIndexedChannelId(state, sourceClass.op, messages, index, loopLoc); Value sourceCoreId = createIndexedSourceCoreId(state, sourceClass.op, messages, index, loopLoc); Value targetCoreId = createIndexedTargetCoreId(state, sourceClass.op, messages, index, loopLoc); - SpatChannelSendOp::create(state.rewriter, loopLoc, channelId, sourceCoreId, targetCoreId, payload); + auto send = SpatChannelSendOp::create(state.rewriter, loopLoc, channelId, sourceCoreId, targetCoreId, payload); + tagProjectedInputTransfer(send.getOperation()); scf::YieldOp::create(state.rewriter, loopLoc); return success(); }); - return failed(sendLoop) ? failure() : success(); -} - -LogicalResult appendConditionalBatchProjectedInputSends(MaterializerState& state, - MaterializedClass& sourceClass, - Value payload, - ArrayRef fragments, - Location loc) { - state.rewriter.setInsertionPoint(sourceClass.body->getTerminator()); - return appendConditionalBatchProjectedInputSendsAtCurrentInsertionPoint(state, sourceClass, payload, fragments, loc); -} - -LogicalResult appendScalarProjectedInputSendNow(MaterializerState& state, - MaterializedClass& sourceClass, - Value payload, - ProjectedInputTransferFragment& fragment, - Location loc) { - assert(!sourceClass.isBatch && "scalar projected input send requires a scalar source class"); - if (sourceClass.cpus.empty()) - return sourceClass.op->emitError("scalar projected input send requires one source core"); - int64_t sourceCore = sourceClass.cpus.front(); - if (sourceCore != fragment.sourceCoreId) - return sourceClass.op->emitError("projected input send source core does not match materialized scalar class") - << " expected=" << sourceCore << " actual=" << fragment.sourceCoreId; - - appendScalarSend(state, sourceClass, payload, fragment.channelId, fragment.sourceCoreId, fragment.targetCoreId, loc); + if (failed(sendLoop)) + return failure(); + if (sendLoop->loop) + tagProjectedInputTransfer(sendLoop->loop.getOperation(), getBestPeerExchangeOrderKey(messages)); return success(); } @@ -5173,8 +5245,14 @@ LogicalResult appendScalarProjectedInputSendAtCurrentInsertionPoint(Materializer return sourceClass.op->emitError("projected input send source core does not match materialized scalar class") << " expected=" << sourceCore << " actual=" << fragment.sourceCoreId; - appendScalarSendAtCurrentInsertionPoint( - state, sourceClass, payload, fragment.channelId, fragment.sourceCoreId, fragment.targetCoreId, loc); + Value channelId = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, fragment.channelId); + Value sourceCoreId = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, fragment.sourceCoreId); + Value targetCoreId = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, fragment.targetCoreId); + auto send = SpatChannelSendOp::create(state.rewriter, loc, channelId, sourceCoreId, targetCoreId, payload); + tagProjectedInputTransfer( + send.getOperation(), + std::optional( + makePeerExchangeOrderKey(fragment.sourceCoreId, fragment.targetCoreId, fragment.channelId))); return success(); } @@ -5195,82 +5273,79 @@ LogicalResult emitProjectedInputSendsAtCurrentInsertionPoint(MaterializerState& return success(); } -LogicalResult emitProjectedInputSendsNow(MaterializerState& state, - MaterializedClass& sourceClass, - Value payload, - ArrayRef fragments, - Location loc) { - if (fragments.empty()) - return success(); - if (sourceClass.isBatch) - return appendConditionalBatchProjectedInputSends(state, sourceClass, payload, fragments, loc); - - for (ProjectedInputTransferFragment* fragment : fragments) - if (failed(appendScalarProjectedInputSendNow(state, sourceClass, payload, *fragment, loc))) - return failure(); - return success(); +static std::optional getProjectedExchangeId(MaterializerState& state, + const ProjectedInputTransferFragment& fragment, + Operation* anchor) { + std::optional exchangeId = state.projectedInputCommunicationPlan.getExchangeId(fragment); + if (!exchangeId && anchor) + anchor->emitError("projected remote fragment is missing from communication plan") + << " sourceCore=" << fragment.sourceCoreId << " targetCore=" << fragment.targetCoreId + << " channelId=" << fragment.channelId; + return exchangeId; } -void enqueuePendingProjectedInputSends(MaterializerState& state, - MaterializedClass& sourceClass, - Value payload, - ArrayRef fragments, - Location loc) { - if (fragments.empty()) - return; - state.pendingProjectedInputSends.push_back(PendingProjectedInputSend { - .sourceClass = sourceClass.id, - .payload = payload, - .fragments = SmallVector(fragments.begin(), fragments.end()), - .loc = loc, - }); +static void recordProjectedInputDeferral(MaterializerState& state, size_t exchangeId) { + if (!llvm::is_contained(state.deferredProjectedExchanges, exchangeId)) + state.deferredProjectedExchanges.push_back(exchangeId); } -LogicalResult flushPendingProjectedInputSendsForClass(MaterializerState& state, MaterializedClass& sourceClass) { - for (size_t index = 0; index < state.pendingProjectedInputSends.size();) { - PendingProjectedInputSend& pending = state.pendingProjectedInputSends[index]; - if (pending.sourceClass != sourceClass.id) { - ++index; +static bool isProjectedRemoteFragmentReadyForReceive(MaterializerState& state, + const ProjectedInputTransferFragment& fragment, + Operation* anchor) { + std::optional exchangeId = getProjectedExchangeId(state, fragment, anchor); + if (!exchangeId) + return false; + const ExchangeDescriptor& exchange = state.projectedInputCommunicationPlan.getExchange(*exchangeId); + ProjectedExchangeEmissionState& emission = state.projectedExchangeStates[*exchangeId]; + if (exchange.phase != state.currentProjectedCommunicationPhase || !emission.sendEmitted) { + recordProjectedInputDeferral(state, *exchangeId); + return false; + } + return true; +} + +static LogicalResult preflightProjectedProducerReceives(MaterializerState& state, + MaterializedClass& targetClass, + ProducerKey producer) { + for (auto& extractEntry : state.projectedInputTransferPlans) { + auto classIt = extractEntry.second.find(targetClass.id); + if (classIt == extractEntry.second.end()) continue; - } - if (failed(emitProjectedInputSendsAtCurrentInsertionPoint( - state, sourceClass, pending.payload, pending.fragments, pending.loc))) - return failure(); - state.pendingProjectedInputSends.erase(state.pendingProjectedInputSends.begin() + index); + for (const ProjectedInputTransferFragment& fragment : classIt->second.fragments) { + if (!(fragment.producer == producer) || fragment.sourceCoreId == fragment.targetCoreId) + continue; + + std::optional exchangeId = + getProjectedExchangeId(state, fragment, targetClass.op); + if (!exchangeId) + return failure(); + + const ExchangeDescriptor& exchange = state.projectedInputCommunicationPlan.getExchange(*exchangeId); + const ProjectedExchangeEmissionState& emission = state.projectedExchangeStates[*exchangeId]; + if (exchange.phase != state.currentProjectedCommunicationPhase || !emission.sendEmitted) + recordProjectedInputDeferral(state, *exchangeId); + } } return success(); } -LogicalResult flushProjectedInputPhaseBSends(MaterializerState& state, MaterializedClass& sourceClass) { - auto countIt = state.pendingProjectedHighToLowReceives.find(sourceClass.id); - if (countIt != state.pendingProjectedHighToLowReceives.end() && countIt->second != 0) - return sourceClass.op->emitError("projected input low-to-high phase started before high-to-low receives completed") - << " classId=" << sourceClass.id << " remainingHighToLowReceives=" << countIt->second, - failure(); - - return flushPendingProjectedInputSendsForClass(state, sourceClass); -} - -LogicalResult releaseProjectedInputPhaseBarrierReceives(MaterializerState& state, - MaterializedClass& targetClass, - unsigned receiveCount) { - if (receiveCount == 0) +static LogicalResult preflightProjectedWholeBatchProducerReceives(MaterializerState& state, + MaterializedClass& targetClass, + ProducerKey producer) { + auto batch = dyn_cast_or_null(producer.instance.op); + if (!batch) return success(); - auto countIt = state.pendingProjectedHighToLowReceives.find(targetClass.id); - if (countIt == state.pendingProjectedHighToLowReceives.end() || countIt->second < receiveCount) - return targetClass.op->emitError("projected input high-to-low receive count underflow") - << " classId=" << targetClass.id << " receiveCount=" << receiveCount, - failure(); + uint32_t laneEnd = producer.instance.laneStart + producer.instance.laneCount; + if (laneEnd > static_cast(batch.getLaneCount())) + return failure(); - countIt->second -= receiveCount; - if (countIt->second != 0) - return success(); - - state.pendingProjectedHighToLowReceives.erase(countIt); - state.projectedInputPhaseBarrierClasses.erase(targetClass.id); - return flushProjectedInputPhaseBSends(state, targetClass); + for (uint32_t lane = producer.instance.laneStart; lane < laneEnd; ++lane) + if (failed(preflightProjectedProducerReceives( + state, targetClass, getBatchLaneProducerKey(batch, lane, 1, producer.resultIndex)))) + return failure(); + return success(); } LogicalResult emitPlannedProjectedInputSendsForKeys(MaterializerState& state, @@ -5285,39 +5360,97 @@ LogicalResult emitPlannedProjectedInputSendsForKeys(MaterializerState& state, for (ProducerKey key : keys) keySet.insert(key); - SmallVector readyFragments; - SmallVector deferredFragments; + SmallVector lowToHighFragments; + SmallVector highToLowFragments; for (auto& extractEntry : state.projectedInputTransferPlans) { for (auto& classEntry : extractEntry.second) { ProjectedInputTransferPlan& plan = classEntry.second; for (ProjectedInputTransferFragment& fragment : plan.fragments) { - if (fragment.sendEmitted || !keySet.contains(fragment.producer)) + if (!keySet.contains(fragment.producer)) continue; - if (fragment.sourceCoreId == fragment.targetCoreId) { - fragment.sendEmitted = true; - continue; + std::optional phase = state.projectedInputCommunicationPlan.getPhase(fragment); + if (!phase) { + return sourceClass.op->emitError("projected input send fragment is missing from communication plan") + << " sourceClass=" << sourceClass.id << " sourceCore=" << fragment.sourceCoreId + << " targetCore=" << fragment.targetCoreId, + failure(); } - if (!isHigherToLowerProjectedInputFragment(fragment) - && state.projectedInputPhaseBarrierClasses.contains(sourceClass.id)) { - deferredFragments.push_back(&fragment); - fragment.sendEmitted = true; - continue; + std::optional exchangeId = state.projectedInputCommunicationPlan.getExchangeId(fragment); + if (exchangeId) { + ProjectedExchangeEmissionState& emission = state.projectedExchangeStates[*exchangeId]; + emission.producerBound = true; + emission.producerPayload = payload; + emission.loc = loc; + if (emission.sendEmitted) + continue; + if (*phase != state.currentProjectedCommunicationPhase) + continue; } - readyFragments.push_back(&fragment); - fragment.sendEmitted = true; + switch (*phase) { + case CommunicationPhase::Local: + fragment.sendEmitted = true; + break; + case CommunicationPhase::LowToHigh: + lowToHighFragments.push_back(&fragment); + fragment.sendEmitted = true; + break; + case CommunicationPhase::HighToLow: + highToLowFragments.push_back(&fragment); + fragment.sendEmitted = true; + break; + } } } } - if (failed(emitProjectedInputSendsNow(state, sourceClass, payload, readyFragments, loc))) + if (failed(emitProjectedInputSendsAtCurrentInsertionPoint(state, sourceClass, payload, lowToHighFragments, loc))) return failure(); - enqueuePendingProjectedInputSends(state, sourceClass, payload, deferredFragments, loc); + if (failed(emitProjectedInputSendsAtCurrentInsertionPoint(state, sourceClass, payload, highToLowFragments, loc))) + return failure(); + for (ProjectedInputTransferFragment* fragment : lowToHighFragments) { + if (std::optional exchangeId = state.projectedInputCommunicationPlan.getExchangeId(*fragment)) + state.projectedExchangeStates[*exchangeId].sendEmitted = true; + fragment->sendEmitted = true; + } + for (ProjectedInputTransferFragment* fragment : highToLowFragments) { + if (std::optional exchangeId = state.projectedInputCommunicationPlan.getExchangeId(*fragment)) + state.projectedExchangeStates[*exchangeId].sendEmitted = true; + fragment->sendEmitted = true; + } return success(); } +static FailureOr emitBoundProjectedInputSendsForCurrentPhase(MaterializerState& state) { + bool emittedAny = false; + for (const ExchangeDescriptor& exchange : state.projectedInputCommunicationPlan.getExchanges()) { + if (exchange.phase != state.currentProjectedCommunicationPhase) + continue; + ProjectedExchangeEmissionState& emission = state.projectedExchangeStates[exchange.id]; + if (!emission.producerBound || emission.sendEmitted) + continue; + if (!exchange.fragment) + return state.func.emitError("projected communication exchange is missing its fragment"); + if (!emission.loc) + return state.func.emitError("projected communication exchange has no producer location"); + + MaterializedClass& sourceClass = state.classes[exchange.sourceClass]; + OpBuilder::InsertionGuard guard(state.rewriter); + state.rewriter.setInsertionPoint(sourceClass.body->getTerminator()); + ProjectedInputTransferFragment* fragment = const_cast(exchange.fragment); + SmallVector fragments {fragment}; + if (failed(emitProjectedInputSendsAtCurrentInsertionPoint( + state, sourceClass, emission.producerPayload, fragments, *emission.loc))) + return failure(); + emission.sendEmitted = true; + fragment->sendEmitted = true; + emittedAny = true; + } + return emittedAny; +} + static FailureOr> collectOrderedBatchProjectedInputFragments(MaterializedClass& targetClass, @@ -5387,29 +5520,98 @@ enum class ProjectedInputExchangeDirection { HigherToLower }; -static ProjectedInputExchangeDirection getProjectedInputExchangeDirectionKind( - const ProjectedInputTransferFragment& fragment) { - if (isLocalProjectedInputFragment(fragment)) +static constexpr StringLiteral projectedInputTransferAttr = "raptor.projected_input_transfer"; +static constexpr StringLiteral projectedInputExchangeMaxCoreAttr = "raptor.projected_input_exchange_max_core"; +static constexpr StringLiteral projectedInputExchangeSourceCoreAttr = "raptor.projected_input_exchange_source_core"; +static constexpr StringLiteral projectedInputExchangeTargetCoreAttr = "raptor.projected_input_exchange_target_core"; +static constexpr StringLiteral projectedInputExchangeChannelAttr = "raptor.projected_input_exchange_channel"; +static constexpr StringLiteral projectedInputDirectionAttr = "raptor.projected_input_direction"; +static constexpr StringLiteral lowerToHigherProjectedInputAttrValue = "lower_to_higher"; +static constexpr StringLiteral higherToLowerProjectedInputAttrValue = "higher_to_lower"; + +static void tagProjectedInputTransfer(Operation* op, std::optional key) { + if (!op) + return; + op->setAttr(projectedInputTransferAttr, UnitAttr::get(op->getContext())); + if (!key) + return; + MLIRContext* context = op->getContext(); + op->setAttr(projectedInputExchangeMaxCoreAttr, IntegerAttr::get(IndexType::get(context), key->maxCore)); + op->setAttr(projectedInputExchangeSourceCoreAttr, IntegerAttr::get(IndexType::get(context), key->sourceCore)); + op->setAttr(projectedInputExchangeTargetCoreAttr, IntegerAttr::get(IndexType::get(context), key->targetCore)); + op->setAttr(projectedInputExchangeChannelAttr, IntegerAttr::get(IndexType::get(context), key->channelId)); +} + +static void tagProjectedInputTransfer(Operation* op) { + tagProjectedInputTransfer(op, std::nullopt); +} + +static FailureOr +getProjectedInputExchangeDirectionKind(MaterializerState& state, + MaterializedClass& targetClass, + const ProjectedInputTransferFragment& fragment) { + std::optional phase = state.projectedInputCommunicationPlan.getPhase(fragment); + if (!phase) + return targetClass.op->emitError("projected input receive fragment is missing from communication plan") + << " targetClass=" << targetClass.id << " sourceCore=" << fragment.sourceCoreId + << " targetCore=" << fragment.targetCoreId, + failure(); + + switch (*phase) { + case CommunicationPhase::Local: return ProjectedInputExchangeDirection::Local; - return isLowerToHigherProjectedInputFragment(fragment) ? ProjectedInputExchangeDirection::LowerToHigher - : ProjectedInputExchangeDirection::HigherToLower; + case CommunicationPhase::LowToHigh: + return ProjectedInputExchangeDirection::LowerToHigher; + case CommunicationPhase::HighToLow: + return ProjectedInputExchangeDirection::HigherToLower; + } + llvm_unreachable("unknown communication phase"); +} + +static StringRef getProjectedInputDirectionAttrValue(ProjectedInputExchangeDirection direction) { + switch (direction) { + case ProjectedInputExchangeDirection::Local: + return StringRef(); + case ProjectedInputExchangeDirection::LowerToHigher: + return lowerToHigherProjectedInputAttrValue; + case ProjectedInputExchangeDirection::HigherToLower: + return higherToLowerProjectedInputAttrValue; + } + llvm_unreachable("unknown projected input direction"); +} + +static void tagProjectedInputReceive(Value value, ProjectedInputExchangeDirection direction) { + StringRef attrValue = getProjectedInputDirectionAttrValue(direction); + if (attrValue.empty()) + return; + + auto receive = value.getDefiningOp(); + if (!receive) + return; + + tagProjectedInputTransfer(receive.getOperation()); + MLIRContext* context = receive->getContext(); + receive->setAttr(projectedInputDirectionAttr, StringAttr::get(context, attrValue)); } static FailureOr getUniformProjectedInputSlotDirection( - MaterializedClass& targetClass, ArrayRef fragments) { + MaterializerState& state, MaterializedClass& targetClass, ArrayRef fragments) { if (fragments.empty()) return targetClass.op->emitError("projected input batch transfer has an empty slot"), failure(); - ProjectedInputExchangeDirection direction = getProjectedInputExchangeDirectionKind(*fragments.front()); + FailureOr direction = + getProjectedInputExchangeDirectionKind(state, targetClass, *fragments.front()); + if (failed(direction)) + return failure(); for (const ProjectedInputTransferFragment* fragment : fragments.drop_front()) { - if (getProjectedInputExchangeDirectionKind(*fragment) == direction) + FailureOr fragmentDirection = + getProjectedInputExchangeDirectionKind(state, targetClass, *fragment); + if (failed(fragmentDirection)) + return failure(); + if (*fragmentDirection == *direction) continue; - return targetClass.op->emitError( - "projected input batch transfer slot mixes communication directions; split lanes before materialization") - << " fragmentCount=" << fragments.size(), - failure(); } - return direction; + return *direction; } static FailureOr resolveUniformLocalBatchProjectedInputSlot(MaterializerState& state, @@ -5459,152 +5661,366 @@ static FailureOr materializeRemoteBatchProjectedInputSlotReceive( .getOutput(); } -static FailureOr materializeBatchProjectedInputPayloadSlot( +struct ProjectedInputPlanMaterialization { + struct BatchSlotGroup { + unsigned slot = 0; + ProjectedInputExchangeDirection direction = ProjectedInputExchangeDirection::Local; + SmallVector lanes; + }; + + Operation* extractOp = nullptr; + const ProjectedInputTransferPlan* plan = nullptr; + SmallVector fragments; + Value payload; + Value laneArg; + SmallVector batchSlotGroups; + SmallVector localScalarFragments; + SmallVector lowerToHigherScalarFragments; + SmallVector higherToLowerScalarFragments; +}; + +static FailureOr +prepareProjectedInputPlanMaterialization(MaterializerState& state, + MaterializedClass& targetClass, + Operation* extractOp, + const ProjectedInputTransferPlan& plan, + Location loc) { + if (failed(verifyProjectedFragmentLayout(targetClass.op, plan.layout))) + return failure(); + + ProjectedInputPlanMaterialization work; + work.extractOp = extractOp; + work.plan = &plan; + + if (targetClass.isBatch) { + auto batch = dyn_cast(targetClass.op); + if (!batch) + return targetClass.op->emitError("projected input transfer batch target must be a scheduled compute_batch"), + failure(); + std::optional laneArg = batch.getLaneArgument(); + if (!laneArg) + return targetClass.op->emitError("projected input transfer batch target is missing its lane argument"), + failure(); + + const unsigned fragmentsPerLane = plan.layout.fragmentsPerLogicalSlot; + if (fragmentsPerLane == 0 || plan.layout.payloadFragmentCount != fragmentsPerLane) + return targetClass.op->emitError("projected input batch transfer requires one logical payload per target lane"), + failure(); + + FailureOr> maybeFragments = + collectOrderedBatchProjectedInputFragments(targetClass, plan, fragmentsPerLane); + if (failed(maybeFragments)) + return failure(); + work.fragments = std::move(*maybeFragments); + work.laneArg = *laneArg; + + FailureOr payloadType = + getProjectedPayloadType(targetClass.op, plan.layout.fragmentType, fragmentsPerLane); + if (failed(payloadType)) + return failure(); + work.payload = + tensor::EmptyOp::create(state.rewriter, loc, payloadType->getShape(), payloadType->getElementType()).getResult(); + + for (unsigned slot = 0; slot < fragmentsPerLane; ++slot) { + SmallVector slotFragments = + getBatchProjectedInputPayloadSlot(work.fragments, targetClass.cpus.size(), fragmentsPerLane, slot); + FailureOr direction = + getUniformProjectedInputSlotDirection(state, targetClass, slotFragments); + if (failed(direction)) + return failure(); + ProjectedInputPlanMaterialization::BatchSlotGroup group; + group.slot = slot; + group.direction = *direction; + for (unsigned lane = 0; lane < targetClass.cpus.size(); ++lane) + group.lanes.push_back(lane); + work.batchSlotGroups.push_back(std::move(group)); + } + return work; + } + + work.fragments.reserve(plan.fragments.size()); + for (const ProjectedInputTransferFragment& fragment : plan.fragments) + work.fragments.push_back(&fragment); + llvm::sort(work.fragments, [](const ProjectedInputTransferFragment* lhs, const ProjectedInputTransferFragment* rhs) { + return lhs->ordinal < rhs->ordinal; + }); + + if (work.fragments.size() != plan.layout.payloadFragmentCount) + return targetClass.op->emitError("projected input transfer did not collect one complete logical payload") + << " expected=" << plan.layout.payloadFragmentCount << " actual=" << work.fragments.size(), + failure(); + for (auto [index, fragment] : llvm::enumerate(work.fragments)) + if (fragment->ordinal != index) + return targetClass.op->emitError("projected input transfer fragments are not ordinal-contiguous"), failure(); + + if (plan.layout.payloadFragmentCount != 1) { + FailureOr payloadType = + getProjectedPayloadType(targetClass.op, plan.layout.fragmentType, plan.layout.payloadFragmentCount); + if (failed(payloadType)) + return failure(); + work.payload = + tensor::EmptyOp::create(state.rewriter, loc, payloadType->getShape(), payloadType->getElementType()).getResult(); + } + + for (auto [index, fragment] : llvm::enumerate(work.fragments)) { + FailureOr direction = + getProjectedInputExchangeDirectionKind(state, targetClass, *fragment); + if (failed(direction)) + return failure(); + if (*direction == ProjectedInputExchangeDirection::HigherToLower) + work.higherToLowerScalarFragments.push_back(index); + else if (*direction == ProjectedInputExchangeDirection::LowerToHigher) + work.lowerToHigherScalarFragments.push_back(index); + else + work.localScalarFragments.push_back(index); + } + + return work; +} + +static FailureOr materializeScalarProjectedInputFragment(MaterializerState& state, + MaterializedClass& targetClass, + const ProjectedInputTransferPlan& plan, + const ProjectedInputTransferFragment& fragment, + Location loc) { + FailureOr direction = + getProjectedInputExchangeDirectionKind(state, targetClass, fragment); + if (failed(direction)) + return failure(); + + if (*direction == ProjectedInputExchangeDirection::Local) { + std::optional local = state.availableValues.lookup(state, fragment.producer, targetClass.id); + if (!local) + return targetClass.op->emitError("projected input transfer could not resolve local fragment"), failure(); + return *local; + } + + if (!isProjectedRemoteFragmentReadyForReceive(state, fragment, targetClass.op)) + return failure(); + + FailureOr received = + createScalarPeerReceiveAndFlush(state, + targetClass, + plan.layout.fragmentType, + getOrCreateIndexConstant(state.constantFolder, targetClass.op, fragment.channelId), + getOrCreateIndexConstant(state.constantFolder, targetClass.op, fragment.sourceCoreId), + getOrCreateIndexConstant(state.constantFolder, targetClass.op, fragment.targetCoreId), + makeScalarPeerReceiveKey(fragment.sourceCoreId, + fragment.targetCoreId, + plan.layout.fragmentType, + fragment.channelId), + loc); + if (failed(received)) + return failure(); + tagProjectedInputReceive(*received, *direction); + if (std::optional exchangeId = state.projectedInputCommunicationPlan.getExchangeId(fragment)) + state.projectedExchangeStates[*exchangeId].receiveEmitted = true; + return *received; +} + +static FailureOr materializeBatchProjectedInputPayloadSlotGroupValue( MaterializerState& state, MaterializedClass& targetClass, const ProjectedInputTransferPlan& plan, - ArrayRef fragments, + ArrayRef slotFragments, + ArrayRef lanes, Value laneArg, Location loc) { - FailureOr direction = getUniformProjectedInputSlotDirection(targetClass, fragments); + SmallVector groupFragments; + groupFragments.reserve(lanes.size()); + SmallVector laneToGroupIndex(targetClass.cpus.size(), 0); + for (auto [groupIndex, lane] : llvm::enumerate(lanes)) { + groupFragments.push_back(slotFragments[lane]); + laneToGroupIndex[lane] = static_cast(groupIndex); + } + + FailureOr direction = + getUniformProjectedInputSlotDirection(state, targetClass, groupFragments); if (failed(direction)) return failure(); if (*direction == ProjectedInputExchangeDirection::Local) - return resolveUniformLocalBatchProjectedInputSlot(state, targetClass, fragments); + return resolveUniformLocalBatchProjectedInputSlot(state, targetClass, groupFragments); - return materializeRemoteBatchProjectedInputSlotReceive( - state, targetClass, plan.layout.fragmentType, fragments, laneArg, loc); + for (const ProjectedInputTransferFragment* fragment : groupFragments) + if (!isProjectedRemoteFragmentReadyForReceive(state, *fragment, targetClass.op)) + return failure(); + + Value groupIndex = createIndexedIndexValue(state, targetClass.op, laneToGroupIndex, laneArg, loc); + FailureOr received = materializeRemoteBatchProjectedInputSlotReceive( + state, targetClass, plan.layout.fragmentType, groupFragments, groupIndex, loc); + if (failed(received)) + return failure(); + tagProjectedInputReceive(*received, *direction); + for (const ProjectedInputTransferFragment* fragment : groupFragments) + if (std::optional exchangeId = state.projectedInputCommunicationPlan.getExchangeId(*fragment)) + state.projectedExchangeStates[*exchangeId].receiveEmitted = true; + return *received; } -static LogicalResult materializeBatchProjectedInputPayloadSlotsInOrder( +static FailureOr insertProjectedBatchSlotGroup(MaterializerState& state, + MaterializedClass& targetClass, + ProjectedInputPlanMaterialization& work, + ArrayRef slotFragments, + ArrayRef lanes, + unsigned slot, + Location loc) { + const ProjectedInputTransferPlan& plan = *work.plan; + FailureOr fragmentValue = materializeBatchProjectedInputPayloadSlotGroupValue( + state, targetClass, plan, slotFragments, lanes, work.laneArg, loc); + if (failed(fragmentValue)) + return failure(); + + Value payloadOrdinal = getOrCreateIndexConstant(state.constantFolder, targetClass.op, static_cast(slot)); + FailureOr offset = + scaleIndexByDim0SizeInClass(state, targetClass, payloadOrdinal, plan.layout.fragmentType.getDimSize(0), loc); + if (failed(offset)) + return failure(); + return createDim0InsertSliceInClass(state, targetClass, loc, *fragmentValue, work.payload, *offset); +} + +static LogicalResult materializeProjectedInputBatchSlotGroup( MaterializerState& state, MaterializedClass& targetClass, - const ProjectedInputTransferPlan& plan, - ArrayRef fragments, - unsigned fragmentsPerLane, - Value laneArg, - Location loc, - Value& packed) { - SmallVector localSlots; - SmallVector lowerToHigherSlots; - SmallVector higherToLowerSlots; + ProjectedInputPlanMaterialization& work, + const ProjectedInputPlanMaterialization::BatchSlotGroup& group, + Location loc) { + const ProjectedInputTransferPlan& plan = *work.plan; + SmallVector slotFragments = + getBatchProjectedInputPayloadSlot( + work.fragments, targetClass.cpus.size(), plan.layout.fragmentsPerLogicalSlot, group.slot); - for (unsigned slot = 0; slot < fragmentsPerLane; ++slot) { - SmallVector slotFragments = - getBatchProjectedInputPayloadSlot(fragments, targetClass.cpus.size(), fragmentsPerLane, slot); - FailureOr direction = - getUniformProjectedInputSlotDirection(targetClass, slotFragments); - if (failed(direction)) - return failure(); + if (group.lanes.size() != targetClass.cpus.size()) + return targetClass.op->emitError("projected input batch transfer produced a partial lane group"), failure(); - switch (*direction) { - case ProjectedInputExchangeDirection::Local: - localSlots.push_back(slot); - break; - case ProjectedInputExchangeDirection::LowerToHigher: - lowerToHigherSlots.push_back(slot); - break; - case ProjectedInputExchangeDirection::HigherToLower: - higherToLowerSlots.push_back(slot); - break; - } - } - - const auto materializeSlot = [&](unsigned slot) -> LogicalResult { - SmallVector slotFragments = - getBatchProjectedInputPayloadSlot(fragments, targetClass.cpus.size(), fragmentsPerLane, slot); - FailureOr fragmentValue = - materializeBatchProjectedInputPayloadSlot(state, targetClass, plan, slotFragments, laneArg, loc); - if (failed(fragmentValue)) - return failure(); - - Value payloadOrdinal = getOrCreateIndexConstant(state.constantFolder, targetClass.op, static_cast(slot)); - FailureOr offset = scaleIndexByDim0SizeInClass( - state, targetClass, payloadOrdinal, plan.layout.fragmentType.getDimSize(0), loc); - if (failed(offset)) - return failure(); - FailureOr updated = createDim0InsertSliceInClass(state, targetClass, loc, *fragmentValue, packed, *offset); - if (failed(updated)) - return failure(); - packed = *updated; - return success(); - }; - - for (unsigned slot : localSlots) - if (failed(materializeSlot(slot))) - return failure(); - unsigned highToLowReceiveCount = 0; - for (unsigned slot : higherToLowerSlots) { - if (failed(materializeSlot(slot))) - return failure(); - highToLowReceiveCount += static_cast(targetClass.cpus.size()); - } - if (failed(releaseProjectedInputPhaseBarrierReceives(state, targetClass, highToLowReceiveCount))) + FailureOr updated = + insertProjectedBatchSlotGroup(state, targetClass, work, slotFragments, group.lanes, group.slot, loc); + if (failed(updated)) return failure(); - for (unsigned slot : lowerToHigherSlots) - if (failed(materializeSlot(slot))) - return failure(); + work.payload = *updated; return success(); } - -FailureOr materializeBatchProjectedInputTransferPlanReceives(MaterializerState& state, - MaterializedClass& targetClass, - const ProjectedInputTransferPlan& plan, - Location loc, - Operation* insertionAnchor) { - auto batch = dyn_cast(targetClass.op); - if (!batch) - return targetClass.op->emitError("projected input transfer batch target must be a scheduled compute_batch"), - failure(); - std::optional laneArg = batch.getLaneArgument(); - if (!laneArg) - return targetClass.op->emitError("projected input transfer batch target is missing its lane argument"), - failure(); - - const unsigned fragmentsPerLane = plan.layout.fragmentsPerLogicalSlot; - if (fragmentsPerLane == 0 || plan.layout.payloadFragmentCount != fragmentsPerLane) - return targetClass.op->emitError("projected input batch transfer requires one logical payload per target lane"), - failure(); - - FailureOr> maybeFragments = - collectOrderedBatchProjectedInputFragments(targetClass, plan, fragmentsPerLane); - if (failed(maybeFragments)) - return failure(); - ArrayRef fragments = *maybeFragments; - - FailureOr payloadType = - getProjectedPayloadType(targetClass.op, plan.layout.fragmentType, fragmentsPerLane); - if (failed(payloadType)) +static LogicalResult materializeProjectedInputScalarFragment(MaterializerState& state, + MaterializedClass& targetClass, + ProjectedInputPlanMaterialization& work, + size_t index, + Location loc) { + const ProjectedInputTransferPlan& plan = *work.plan; + const ProjectedInputTransferFragment& fragment = *work.fragments[index]; + FailureOr received = materializeScalarProjectedInputFragment(state, targetClass, plan, fragment, loc); + if (failed(received)) return failure(); - OpBuilder::InsertionGuard guard(state.rewriter); - if (insertionAnchor) { - if (insertionAnchor->getBlock() != targetClass.body) - return targetClass.op->emitError("projected input transfer insertion anchor is outside the target class body"), - failure(); - state.rewriter.setInsertionPoint(insertionAnchor); + if (plan.layout.payloadFragmentCount == 1) { + work.payload = *received; + return success(); } - Value packed = tensor::EmptyOp::create(state.rewriter, loc, payloadType->getShape(), payloadType->getElementType()).getResult(); - if (failed(materializeBatchProjectedInputPayloadSlotsInOrder( - state, targetClass, plan, fragments, fragmentsPerLane, *laneArg, loc, packed))) + Value offset = getOrCreateIndexConstant( + state.constantFolder, targetClass.op, static_cast(index) * plan.layout.fragmentType.getDimSize(0)); + FailureOr next = createDim0InsertSliceInClass(state, targetClass, loc, *received, work.payload, offset); + if (failed(next)) { + targetClass.op->emitError("failed to pack projected input transfer fragment") + << " targetClass=" << targetClass.id << " ordinal=" << fragment.ordinal + << " sourceCore=" << fragment.sourceCoreId << " targetCore=" << fragment.targetCoreId + << " fragmentType=" << plan.layout.fragmentType; return failure(); - - return packed; + } + work.payload = *next; + return success(); } -FailureOr materializeProjectedInputTransferPlanReceives(MaterializerState& state, +static bool hasRemoteProjectedInputFragment(const ProjectedInputTransferPlan& plan) { + return llvm::any_of(plan.fragments, [](const ProjectedInputTransferFragment& fragment) { + return !isLocalProjectedInputFragment(fragment); + }); +} + +static LogicalResult preflightProjectedInputPlanReceives(MaterializerState& state, + MaterializedClass& targetClass, + const ProjectedInputTransferPlan& plan) { + for (const ProjectedInputTransferFragment& fragment : plan.fragments) { + if (isLocalProjectedInputFragment(fragment)) + continue; + std::optional exchangeId = state.projectedInputCommunicationPlan.getExchangeId(fragment); + if (!exchangeId) + return targetClass.op->emitError("projected input preflight found a remote fragment missing from communication plan") + << " sourceCore=" << fragment.sourceCoreId << " targetCore=" << fragment.targetCoreId + << " channelId=" << fragment.channelId, + failure(); + const ExchangeDescriptor& exchange = state.projectedInputCommunicationPlan.getExchange(*exchangeId); + const ProjectedExchangeEmissionState& emission = state.projectedExchangeStates[*exchangeId]; + if (exchange.phase != state.currentProjectedCommunicationPhase || !emission.sendEmitted) + recordProjectedInputDeferral(state, *exchangeId); + } + return success(); +} + +static LogicalResult preflightProjectedInputsForInstance(MaterializerState& state, + MaterializedClass& targetClass, + const ComputeInstance& instance) { + SmallVector requestedExtracts; + instance.op->walk([&](tensor::ExtractSliceOp extract) { + auto planIt = state.projectedInputTransferPlans.find(extract.getOperation()); + if (planIt == state.projectedInputTransferPlans.end()) + return; + if (planIt->second.find(targetClass.id) != planIt->second.end()) + requestedExtracts.push_back(extract.getOperation()); + }); + + for (Operation* requestedExtract : requestedExtracts) { + for (auto& extractEntry : state.projectedInputTransferPlans) { + auto classIt = extractEntry.second.find(targetClass.id); + if (classIt == extractEntry.second.end()) + continue; + + auto replacementIt = state.projectedExtractReplacements.find(extractEntry.first); + if (replacementIt != state.projectedExtractReplacements.end() + && replacementIt->second.find(targetClass.id) != replacementIt->second.end()) + continue; + if (extractEntry.first != requestedExtract && !hasRemoteProjectedInputFragment(classIt->second)) + continue; + + if (failed(preflightProjectedInputPlanReceives(state, targetClass, classIt->second))) + return failure(); + } + } + return success(); +} + +static FailureOr> +collectProjectedInputPlanMaterializationsForClass(MaterializerState& state, + MaterializedClass& targetClass, + Operation* requestedExtract, + Location loc) { + SmallVector works; + for (auto& extractEntry : state.projectedInputTransferPlans) { + auto classIt = extractEntry.second.find(targetClass.id); + if (classIt == extractEntry.second.end()) + continue; + + auto replacementIt = state.projectedExtractReplacements.find(extractEntry.first); + if (replacementIt != state.projectedExtractReplacements.end() + && replacementIt->second.find(targetClass.id) != replacementIt->second.end()) + continue; + if (extractEntry.first != requestedExtract && !hasRemoteProjectedInputFragment(classIt->second)) + continue; + + FailureOr work = + prepareProjectedInputPlanMaterialization(state, targetClass, extractEntry.first, classIt->second, loc); + if (failed(work)) + return failure(); + works.push_back(std::move(*work)); + } + return works; +} + +static LogicalResult materializeProjectedInputReceivesForClass(MaterializerState& state, MaterializedClass& targetClass, - const ProjectedInputTransferPlan& plan, + Operation* requestedExtract, Location loc, Operation* insertionAnchor = nullptr) { - if (failed(verifyProjectedFragmentLayout(targetClass.op, plan.layout))) - return failure(); - if (targetClass.isBatch) - return materializeBatchProjectedInputTransferPlanReceives(state, targetClass, plan, loc, insertionAnchor); - OpBuilder::InsertionGuard guard(state.rewriter); if (insertionAnchor) { if (insertionAnchor->getBlock() != targetClass.body) @@ -5613,135 +6029,76 @@ FailureOr materializeProjectedInputTransferPlanReceives(MaterializerState state.rewriter.setInsertionPoint(insertionAnchor); } - SmallVector fragments; - fragments.reserve(plan.fragments.size()); - for (const ProjectedInputTransferFragment& fragment : plan.fragments) - fragments.push_back(&fragment); - llvm::sort(fragments, [](const ProjectedInputTransferFragment* lhs, const ProjectedInputTransferFragment* rhs) { - return lhs->ordinal < rhs->ordinal; - }); - - if (fragments.size() != plan.layout.payloadFragmentCount) - return targetClass.op->emitError("projected input transfer did not collect one complete logical payload") - << " expected=" << plan.layout.payloadFragmentCount << " actual=" << fragments.size(); - for (auto [index, fragment] : llvm::enumerate(fragments)) - if (fragment->ordinal != index) - return targetClass.op->emitError("projected input transfer fragments are not ordinal-contiguous"); - - const auto materializeOneFragment = [&](const ProjectedInputTransferFragment& fragment) -> FailureOr { - if (fragment.sourceCoreId == fragment.targetCoreId) { - std::optional local = state.availableValues.lookup(state, fragment.producer, targetClass.id); - if (!local) - return targetClass.op->emitError("projected input transfer could not resolve local fragment"), failure(); - return *local; - } - - return createScalarPeerReceiveAndFlush(state, - targetClass, - plan.layout.fragmentType, - getOrCreateIndexConstant(state.constantFolder, targetClass.op, fragment.channelId), - getOrCreateIndexConstant(state.constantFolder, targetClass.op, fragment.sourceCoreId), - getOrCreateIndexConstant(state.constantFolder, targetClass.op, fragment.targetCoreId), - makeScalarPeerReceiveKey(fragment.sourceCoreId, - fragment.targetCoreId, - plan.layout.fragmentType, - fragment.channelId), - loc); - }; - - if (fragments.size() == 1) { - const ProjectedInputTransferFragment& fragment = *fragments.front(); - if (isLowerToHigherProjectedInputFragment(fragment)) - if (failed(flushProjectedInputPhaseBSends(state, targetClass))) - return failure(); - FailureOr value = materializeOneFragment(fragment); - if (failed(value)) - return failure(); - if (isHigherToLowerProjectedInputFragment(fragment)) - if (failed(releaseProjectedInputPhaseBarrierReceives(state, targetClass, 1))) - return failure(); - return *value; - } - - FailureOr payloadType = - getProjectedPayloadType(targetClass.op, plan.layout.fragmentType, plan.layout.payloadFragmentCount); - if (failed(payloadType)) + FailureOr> maybeWorks = + collectProjectedInputPlanMaterializationsForClass(state, targetClass, requestedExtract, loc); + if (failed(maybeWorks)) return failure(); - Value packed = tensor::EmptyOp::create(state.rewriter, loc, payloadType->getShape(), payloadType->getElementType()).getResult(); - - SmallVector local; - SmallVector lowerToHigher; - SmallVector higherToLower; - for (auto [index, fragment] : llvm::enumerate(fragments)) { - if (isHigherToLowerProjectedInputFragment(*fragment)) - higherToLower.push_back(index); - else if (isLowerToHigherProjectedInputFragment(*fragment)) - lowerToHigher.push_back(index); - else - local.push_back(index); - } - - const auto materializeAndInsert = [&](size_t index) -> LogicalResult { - const ProjectedInputTransferFragment& fragment = *fragments[index]; - FailureOr received = materializeOneFragment(fragment); - if (failed(received)) - return failure(); - - Value offset = getOrCreateIndexConstant( - state.constantFolder, targetClass.op, static_cast(index) * plan.layout.fragmentType.getDimSize(0)); - FailureOr next = createDim0InsertSliceInClass(state, targetClass, loc, *received, packed, offset); - if (failed(next)) { - targetClass.op->emitError("failed to pack projected input transfer fragment") - << " targetClass=" << targetClass.id << " ordinal=" << fragment.ordinal - << " sourceCore=" << fragment.sourceCoreId << " targetCore=" << fragment.targetCoreId - << " fragmentType=" << plan.layout.fragmentType << " payloadType=" << *payloadType; - return failure(); - } - packed = *next; + SmallVector& works = *maybeWorks; + if (works.empty()) return success(); - }; - for (size_t index : local) - if (failed(materializeAndInsert(index))) - return failure(); + for (ProjectedInputPlanMaterialization& work : works) { + for (const auto& group : work.batchSlotGroups) + if (group.direction == ProjectedInputExchangeDirection::Local + && failed(materializeProjectedInputBatchSlotGroup(state, targetClass, work, group, loc))) + return failure(); + for (size_t index : work.localScalarFragments) + if (failed(materializeProjectedInputScalarFragment(state, targetClass, work, index, loc))) + return failure(); + } - for (size_t index : higherToLower) - if (failed(materializeAndInsert(index))) - return failure(); + for (ProjectedInputPlanMaterialization& work : works) { + for (const auto& group : work.batchSlotGroups) + if (group.direction == ProjectedInputExchangeDirection::LowerToHigher + && failed(materializeProjectedInputBatchSlotGroup(state, targetClass, work, group, loc))) + return failure(); + for (size_t index : work.lowerToHigherScalarFragments) + if (failed(materializeProjectedInputScalarFragment(state, targetClass, work, index, loc))) + return failure(); + } - if (failed(releaseProjectedInputPhaseBarrierReceives( - state, targetClass, static_cast(higherToLower.size())))) - return failure(); + for (ProjectedInputPlanMaterialization& work : works) { + for (const auto& group : work.batchSlotGroups) + if (group.direction == ProjectedInputExchangeDirection::HigherToLower + && failed(materializeProjectedInputBatchSlotGroup(state, targetClass, work, group, loc))) + return failure(); + for (size_t index : work.higherToLowerScalarFragments) + if (failed(materializeProjectedInputScalarFragment(state, targetClass, work, index, loc))) + return failure(); + } - for (size_t index : lowerToHigher) - if (failed(materializeAndInsert(index))) - return failure(); - - return packed; + for (const ProjectedInputPlanMaterialization& work : works) { + if (!work.payload) + return targetClass.op->emitError("projected input transfer did not produce a replacement payload") + << " targetClass=" << targetClass.id, + failure(); + state.projectedExtractReplacements[work.extractOp][targetClass.id] = + ProjectedExtractReplacement {work.payload, work.plan->layout}; + } + return success(); } LogicalResult verifyNoPendingProjectedInputSends(MaterializerState& state) { - if (!state.pendingProjectedHighToLowReceives.empty()) { - auto countIt = state.pendingProjectedHighToLowReceives.begin(); - MaterializedClass& targetClass = state.classes[countIt->first]; - return targetClass.op->emitError("projected input high-to-low receive phase was not fully materialized") - << " classId=" << countIt->first << " remainingReceives=" << countIt->second, - failure(); + for (const ExchangeDescriptor& exchange : state.projectedInputCommunicationPlan.getExchanges()) { + const ProjectedExchangeEmissionState& emission = state.projectedExchangeStates[exchange.id]; + if (!emission.sendEmitted) { + MaterializedClass& targetClass = state.classes[exchange.targetClass]; + return targetClass.op->emitError("projected input communication plan left a remote send unemitted") + << " exchange=" << exchange.id << " targetClass=" << exchange.targetClass + << " sourceCore=" << exchange.sourceCore << " targetCore=" << exchange.targetCore + << " channelId=" << exchange.channelId, + failure(); + } + if (!emission.receiveEmitted) { + MaterializedClass& targetClass = state.classes[exchange.targetClass]; + return targetClass.op->emitError("projected input communication plan left a remote receive unemitted") + << " exchange=" << exchange.id << " targetClass=" << exchange.targetClass + << " sourceCore=" << exchange.sourceCore << " targetCore=" << exchange.targetCore + << " channelId=" << exchange.channelId, + failure(); + } } - - if (state.pendingProjectedInputSends.empty()) - return success(); - - PendingProjectedInputSend& pending = state.pendingProjectedInputSends.front(); - MaterializedClass& sourceClass = state.classes[pending.sourceClass]; - InFlightDiagnostic diag = sourceClass.op->emitError("projected input send phase barrier was never released") - << " sourceClass=" << pending.sourceClass - << " fragmentCount=" << pending.fragments.size(); - if (!pending.fragments.empty()) { - ProjectedInputTransferFragment* fragment = pending.fragments.front(); - diag << " firstSourceCore=" << fragment->sourceCoreId << " firstTargetCore=" << fragment->targetCoreId; - } - return failure(); + return success(); } MessageVector filterMessageVector(const MessageVector& messages, ArrayRef indices) { @@ -7093,6 +7450,33 @@ FailureOr emitProjectedWholeBatchFragmentInsertLoop(MaterializerState& st return loop->results.front(); } +static bool isProducerValueMaterialized(MaterializerState& state, ProducerKey key) { + auto cpuIt = state.schedule.computeToCpuMap.find(key.instance); + if (cpuIt == state.schedule.computeToCpuMap.end()) + return false; + auto classIt = state.cpuToClass.find(cpuIt->second); + if (classIt == state.cpuToClass.end()) + return false; + auto rangeIt = state.scheduledInstanceToLogicalSlots.find(key.instance); + if (rangeIt == state.scheduledInstanceToLogicalSlots.end()) + return false; + + LogicalSlotRange range = rangeIt->second; + for (SlotId slot = range.start; slot < range.start + range.count; ++slot) + if (!state.materializedLogicalSlots.contains({classIt->second, slot})) + return false; + return true; +} + +static void recordDeferredProducerValue(MaterializerState& state, ProducerKey key) { + if (!llvm::is_contained(state.deferredProducerValues, key)) + state.deferredProducerValues.push_back(key); +} + +static bool hasDeferredMaterializationDependency(const MaterializerState& state) { + return !state.deferredProjectedExchanges.empty() || !state.deferredProducerValues.empty(); +} + std::optional getStaticProjectedPackedFragmentIndex(tensor::ExtractSliceOp extract) { auto sourceType = dyn_cast(extract.getSource().getType()); auto resultType = dyn_cast(extract.getResult().getType()); @@ -7213,6 +7597,10 @@ FailureOr materializeProjectedWholeBatchInputFromFragments( uint32_t laneEnd = key.instance.laneStart + key.instance.laneCount; if (laneEnd > static_cast(batch.getLaneCount())) return failure(); + if (failed(preflightProjectedWholeBatchProducerReceives(state, targetClass, key))) + return failure(); + if (!state.deferredProjectedExchanges.empty()) + return failure(); if (targetClass.isBatch) { state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); @@ -7223,8 +7611,15 @@ FailureOr materializeProjectedWholeBatchInputFromFragments( for (uint32_t lane = key.instance.laneStart; lane < laneEnd; ++lane) { ProducerKey laneKey = getBatchLaneProducerKey(batch, lane, 1, key.resultIndex); std::optional fragment = state.availableValues.lookup(state, laneKey, targetClass.id); - if (!fragment) + if (!fragment && !isProducerValueMaterialized(state, laneKey)) { + recordDeferredProducerValue(state, laneKey); return failure(); + } + if (!fragment) + return targetClass.op->emitError("projected whole-batch assembly is missing a lane fragment") + << " targetClass=" << targetClass.id << " lane=" << lane + << " sourceOp='" << batch->getName() << "' resultIndex=" << key.resultIndex, + failure(); FailureOr> offsets = evaluateStaticProjectionIndices(projection->getMixedOffsets(), *laneArg, lane); @@ -7343,8 +7738,15 @@ FailureOr materializeProjectedWholeBatchInputFromFragments( continue; std::optional fragment = state.availableValues.lookup(state, laneKey, targetClass.id); - if (!fragment) + if (!fragment && !isProducerValueMaterialized(state, laneKey)) { + recordDeferredProducerValue(state, laneKey); return failure(); + } + if (!fragment) + return targetClass.op->emitError("projected whole-batch assembly is missing a lane fragment") + << " targetClass=" << targetClass.id << " lane=" << lane + << " sourceOp='" << batch->getName() << "' resultIndex=" << key.resultIndex, + failure(); auto fragmentType = dyn_cast(fragment->getType()); if (!fragmentType) @@ -7791,6 +8193,10 @@ FailureOr resolveInputValue(MaterializerState& state, if (isWholeBatchProducerKey(*producer)) { if (!allowWholeBatchFallback) { + if (!isProducerValueMaterialized(state, *producer)) { + recordDeferredProducerValue(state, *producer); + return failure(); + } consumerInstance.op->emitError("failed to resolve compute_batch input without a compact fragment path") << " from '" << producer->instance.op->getName() << "' laneStart=" << producer->instance.laneStart << " laneCount=" << producer->instance.laneCount << " resultIndex=" << producer->resultIndex; @@ -7798,7 +8204,7 @@ FailureOr resolveInputValue(MaterializerState& state, } FailureOr wholeBatch = materializeWholeBatchInput(state, targetClass, *producer, input.getType(), consumerInstance.op->getLoc()); - if (failed(wholeBatch)) + if (failed(wholeBatch) && !hasDeferredMaterializationDependency(state)) consumerInstance.op->emitError("failed to materialize whole-batch input") << " from '" << producer->instance.op->getName() << "' laneStart=" << producer->instance.laneStart << " laneCount=" << producer->instance.laneCount << " resultIndex=" << producer->resultIndex; @@ -7807,9 +8213,22 @@ FailureOr resolveInputValue(MaterializerState& state, return rejectNonLocalResolvedValue(*wholeBatch); } + if (!isProducerValueMaterialized(state, *producer)) { + recordDeferredProducerValue(state, *producer); + return failure(); + } + + std::optional producerClass; + if (auto cpuIt = state.schedule.computeToCpuMap.find(producer->instance); + cpuIt != state.schedule.computeToCpuMap.end()) { + if (auto classIt = state.cpuToClass.find(cpuIt->second); classIt != state.cpuToClass.end()) + producerClass = classIt->second; + } consumerInstance.op->emitError("failed to resolve producer value") << " from op '" << producer->instance.op->getName() << "' laneStart=" << producer->instance.laneStart - << " laneCount=" << producer->instance.laneCount << " resultIndex=" << producer->resultIndex; + << " laneCount=" << producer->instance.laneCount << " resultIndex=" << producer->resultIndex + << " targetClass=" << targetClass.id + << " producerClass=" << (producerClass ? static_cast(*producerClass) : -1); return failure(); } @@ -7945,6 +8364,8 @@ LogicalResult mapInputs(MaterializerState& state, for (auto [index, input] : llvm::enumerate(compute.getInputs())) { FailureOr mapped = resolveInputValue(state, targetClass, input, instance, indexing); if (failed(mapped)) { + if (hasDeferredMaterializationDependency(state)) + return failure(); std::optional producer = getInputRequestProducerKey(input, instance); auto diagnostic = compute.emitOpError("failed to resolve materialized compute input") << " #" << index; if (producer) { @@ -7976,7 +8397,8 @@ LogicalResult mapInputs(MaterializerState& state, auto batch = cast(op); for (auto [index, input] : llvm::enumerate(batch.getInputs())) { if (hasProjectedInputReplacement( - state, batch, static_cast(index), input, instance, targetClass.id)) + state, batch, static_cast(index), input, instance, targetClass.id) + && !hasUnreplacedBatchInputUse(state, batch, static_cast(index), targetClass.id)) continue; FailureOr demand = @@ -7989,17 +8411,21 @@ LogicalResult mapInputs(MaterializerState& state, assert(demand->wholeTensorProducer && "whole-tensor input demand must carry a producer"); mapped = materializeWholeBatchInput( state, targetClass, *demand->wholeTensorProducer, input.getType(), batch.getOperation()->getLoc()); - if (failed(mapped)) + if (failed(mapped) && !hasDeferredMaterializationDependency(state)) return batch.emitOpError("failed to materialize whole-batch compute_batch input") << " #" << index << " from '" << demand->wholeTensorProducer->instance.op->getName() << "' laneStart=" << demand->wholeTensorProducer->instance.laneStart << " laneCount=" << demand->wholeTensorProducer->instance.laneCount << " resultIndex=" << demand->wholeTensorProducer->resultIndex; + if (failed(mapped)) + return failure(); } else { mapped = resolveInputValue(state, targetClass, input, instance, indexing, /*allowWholeBatchFallback=*/false); - if (failed(mapped)) + if (failed(mapped) && !hasDeferredMaterializationDependency(state)) return batch.emitOpError("failed to resolve materialized compute_batch input"); + if (failed(mapped)) + return failure(); } auto inputArg = batch.getInputArgument(index); @@ -8126,18 +8552,21 @@ FailureOr> getOrMaterializeProjectedE if (classIt == planIt->second.end()) return std::optional {}; - FailureOr payload = materializeProjectedInputTransferPlanReceives( - state, targetClass, classIt->second, extract.getLoc(), insertionAnchor); - if (failed(payload)) { - targetClass.op->emitError("failed to materialize projected input transfer replacement") - << " targetClass=" << targetClass.id << " fragmentCount=" << classIt->second.fragments.size() - << " payloadFragmentCount=" << classIt->second.layout.payloadFragmentCount; + if (failed(materializeProjectedInputReceivesForClass( + state, targetClass, extract.getOperation(), extract.getLoc(), insertionAnchor))) { + targetClass.op->emitError("failed to materialize projected input transfer replacements for class") + << " targetClass=" << targetClass.id; return failure(); } - ProjectedExtractReplacement replacement {*payload, classIt->second.layout}; - state.projectedExtractReplacements[extract.getOperation()][targetClass.id] = replacement; - return std::optional {replacement}; + if (std::optional replacement = + lookupProjectedExtractReplacement(state, targetClass, extract)) + return replacement; + + targetClass.op->emitError("projected input transfer class materialization did not produce requested replacement") + << " targetClass=" << targetClass.id << " fragmentCount=" << classIt->second.fragments.size() + << " payloadFragmentCount=" << classIt->second.layout.payloadFragmentCount; + return failure(); } Operation* getTopLevelMaterializedClassBodyOp(Operation* nestedOp, MaterializedClass& targetClass) { @@ -8767,6 +9196,8 @@ FailureOr> cloneInstanceBody(MaterializerState& state, mapWeights(state, targetClass, instance, mapper); if (failed(mapInputs(state, targetClass, instance, mapper, indexing))) { + if (hasDeferredMaterializationDependency(state)) + return failure(); sourceOp->emitError("failed to map inputs while cloning compute/compute_batch body") << " laneStart=" << instance.laneStart << " laneCount=" << instance.laneCount << " targetClass=" << targetClass.id; @@ -9453,6 +9884,14 @@ void markMaterializationRunSlots(MaterializerState& state, state.materializedLogicalSlots.insert({classId, startSlot + static_cast(slotIndex)}); } +void unmarkMaterializationRunSlots(MaterializerState& state, + ClassId classId, + SlotId startSlot, + ArrayRef run) { + for (auto slotIndex : llvm::seq(0, run.size())) + state.materializedLogicalSlots.erase({classId, startSlot + static_cast(slotIndex)}); +} + LogicalResult materializeScalarBatchRun(MaterializerState& state, MaterializedClass& targetClass, SlotId startSlot, @@ -9487,6 +9926,10 @@ LogicalResult materializeScalarBatchRun(MaterializerState& state, FailureOr> packedOutputs = materializeBatchOutputGroupLoop(state, targetClass, run, group); if (failed(packedOutputs)) { + if (hasDeferredMaterializationDependency(state)) { + unmarkMaterializationRunSlots(state, targetClass.id, startSlot, run); + return failure(); + } InFlightDiagnostic diagnostic = sourceBatch.emitOpError("failed to materialize scalar batch output group") << " targetClass=" << targetClass.id << " runSize=" << run.size(); if (!run.empty() && !run.front().peers.empty()) @@ -10028,9 +10471,6 @@ LogicalResult materializeInstanceSlot(MaterializerState& state, const ComputeIns } } - if (!state.materializedLogicalSlots.insert({classId, startLogicalSlot}).second) - return success(); - FailureOr> peers = getMaterializationRunSlotPeers(state, targetClass, startLogicalSlot); if (failed(peers)) @@ -10060,6 +10500,149 @@ LogicalResult materializeInstanceSlot(MaterializerState& state, const ComputeIns return failure(); } + state.materializedLogicalSlots.insert({classId, startLogicalSlot}); + return success(); +} + +enum class MaterializationAttemptKind { + Success, + DeferredProjectedInput, + Failure +}; + +static LogicalResult preflightProjectedWholeBatchInputsForInstance(MaterializerState& state, + MaterializedClass& targetClass, + const ComputeInstance& instance) { + if (auto compute = dyn_cast(instance.op)) { + for (Value input : compute.getInputs()) { + std::optional producer = getInputRequestProducerKey(input, instance); + if (producer && isWholeBatchProducerKey(*producer) + && failed(preflightProjectedWholeBatchProducerReceives(state, targetClass, *producer))) + return failure(); + } + return success(); + } + + auto batch = cast(instance.op); + for (auto [index, input] : llvm::enumerate(batch.getInputs())) { + if (hasProjectedInputReplacement(state, batch, static_cast(index), input, instance, targetClass.id) + && !hasUnreplacedBatchInputUse(state, batch, static_cast(index), targetClass.id)) + continue; + + FailureOr demand = + classifyComputeBatchInputDemand(state, targetClass, batch, static_cast(index), input, instance); + if (failed(demand)) + return failure(); + if (demand->kind == BatchInputDemandKind::WholeTensorBarrier && demand->wholeTensorProducer + && failed(preflightProjectedWholeBatchProducerReceives(state, targetClass, *demand->wholeTensorProducer))) + return failure(); + } + return success(); +} + +static MaterializationAttemptKind attemptMaterializeInstanceSlot(MaterializerState& state, + const ComputeInstance& instance) { + state.deferredProjectedExchanges.clear(); + state.deferredProducerValues.clear(); + + auto cpuIt = state.schedule.computeToCpuMap.find(instance); + if (cpuIt == state.schedule.computeToCpuMap.end()) + return MaterializationAttemptKind::Failure; + ClassId classId = state.cpuToClass.lookup(cpuIt->second); + MaterializedClass& targetClass = state.classes[classId]; + if (failed(preflightProjectedInputsForInstance(state, targetClass, instance))) + return MaterializationAttemptKind::Failure; + if (failed(preflightProjectedWholeBatchInputsForInstance(state, targetClass, instance))) + return MaterializationAttemptKind::Failure; + if (!state.deferredProjectedExchanges.empty() || !state.deferredProducerValues.empty()) + return MaterializationAttemptKind::DeferredProjectedInput; + + if (succeeded(materializeInstanceSlot(state, instance))) + return MaterializationAttemptKind::Success; + if (!state.deferredProjectedExchanges.empty() || !state.deferredProducerValues.empty()) + return MaterializationAttemptKind::DeferredProjectedInput; + return MaterializationAttemptKind::Failure; +} + +static bool areProjectedPhaseSendsComplete(MaterializerState& state, CommunicationPhase phase) { + for (const ExchangeDescriptor& exchange : state.projectedInputCommunicationPlan.getExchanges()) { + if (exchange.phase != phase) + continue; + if (!state.projectedExchangeStates[exchange.id].sendEmitted) + return false; + } + return true; +} + +static LogicalResult emitProjectedDeferralDiagnostic(MaterializerState& state, + const ComputeInstance& instance, + size_t remainingCount) { + InFlightDiagnostic diag = instance.op->emitError("projected input communication made no materialization progress") + << " remainingInstances=" << remainingCount + << " phase=" + << (state.currentProjectedCommunicationPhase == CommunicationPhase::LowToHigh + ? "LowToHigh" + : "HighToLow"); + unsigned reported = 0; + for (size_t exchangeId : state.deferredProjectedExchanges) { + if (reported++ == 4) + break; + const ExchangeDescriptor& exchange = state.projectedInputCommunicationPlan.getExchange(exchangeId); + const ProjectedExchangeEmissionState& emission = state.projectedExchangeStates[exchangeId]; + diag << " [exchange=" << exchangeId << " sourceCore=" << exchange.sourceCore + << " targetCore=" << exchange.targetCore << " channel=" << exchange.channelId + << " producerBound=" << emission.producerBound << " sendEmitted=" << emission.sendEmitted + << " receiveEmitted=" << emission.receiveEmitted << "]"; + } + for (ProducerKey producer : state.deferredProducerValues) { + if (reported++ == 4) + break; + diag << " [producerOp='" << producer.instance.op->getName() << "' laneStart=" << producer.instance.laneStart + << " laneCount=" << producer.instance.laneCount << " resultIndex=" << producer.resultIndex << "]"; + } + return failure(); +} + +static LogicalResult materializeInstanceSlotsWithProjectedDeferral(MaterializerState& state) { + SmallVector remaining(state.schedule.dominanceOrderCompute.begin(), + state.schedule.dominanceOrderCompute.end()); + + while (!remaining.empty()) { + bool progress = false; + for (size_t index = 0; index < remaining.size();) { + const ComputeInstance& instance = remaining[index]; + MaterializationAttemptKind attempt = attemptMaterializeInstanceSlot(state, instance); + if (attempt == MaterializationAttemptKind::Success) { + remaining.erase(remaining.begin() + index); + progress = true; + continue; + } + if (attempt == MaterializationAttemptKind::Failure) { + InFlightDiagnostic diag = instance.op->emitError( + "merge schedule materialization failed while materializing scheduled compute instance"); + diag << " laneStart=" << instance.laneStart << " laneCount=" << instance.laneCount; + return failure(); + } + ++index; + } + + FailureOr emitted = emitBoundProjectedInputSendsForCurrentPhase(state); + if (failed(emitted)) + return failure(); + progress |= *emitted; + + if (!progress && state.currentProjectedCommunicationPhase == CommunicationPhase::LowToHigh + && areProjectedPhaseSendsComplete(state, CommunicationPhase::LowToHigh)) { + state.currentProjectedCommunicationPhase = CommunicationPhase::HighToLow; + FailureOr highToLowEmitted = emitBoundProjectedInputSendsForCurrentPhase(state); + if (failed(highToLowEmitted)) + return failure(); + progress = true; + } + + if (!progress) + return emitProjectedDeferralDiagnostic(state, remaining.front(), remaining.size()); + } return success(); } @@ -10114,21 +10697,17 @@ MergeScheduleMaterializer::run(func::FuncOp func, const MergeScheduleResult& sch return failure(); if (failed(runStage("projected-transfer collection", [&]() { return collectProjectedTransfers(state); }))) return failure(); + if (failed(runStage("projected input communication planning", [&]() { + return buildProjectedInputCommunicationPlan(state); + }))) + return failure(); if (failed(runStage("planned scalar-peer receive collection", [&]() { return collectPlannedScalarPeerReceives(state); }))) return failure(); if (failed(runStage("instance-slot materialization", [&]() -> LogicalResult { - for (const ComputeInstance& instance : schedule.dominanceOrderCompute) { - if (succeeded(materializeInstanceSlot(state, instance))) - continue; - InFlightDiagnostic diag = instance.op->emitError( - "merge schedule materialization failed while materializing scheduled compute instance"); - diag << " laneStart=" << instance.laneStart << " laneCount=" << instance.laneCount; - return failure(); - } - return success(); + return materializeInstanceSlotsWithProjectedDeferral(state); }))) return failure(); diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializedClassState.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializedClassState.hpp index 5f41098..06c6168 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializedClassState.hpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializedClassState.hpp @@ -12,6 +12,7 @@ #include #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 fragments; - mlir::Location loc; -}; - enum class BatchInputDemandKind { LaneFragment, ProjectedFragment, @@ -177,6 +171,14 @@ struct BatchInputDemand { std::optional wholeTensorProducer; }; +struct ProjectedExchangeEmissionState { + bool producerBound = false; + bool sendEmitted = false; + bool receiveEmitted = false; + mlir::Value producerPayload; + std::optional loc; +}; + struct CloneIndexingContext { std::optional runSlotIndex; std::optional projectionSlotIndex; @@ -267,6 +269,7 @@ struct MaterializerState { llvm::DenseSet materializedLogicalSlots; llvm::DenseMap, ProducerKeyInfo> producerDestClasses; + llvm::DenseMap, ProducerKeyInfo> ordinaryProducerDestClasses; llvm::DenseMap, SameClassConsumerLookupKeyInfo> sameClassConsumerIndex; llvm::DenseMap @@ -282,6 +285,11 @@ struct MaterializerState { projectedExtractReplacements; llvm::DenseMap> projectedInputTransferPlans; + CommunicationPlan projectedInputCommunicationPlan; + llvm::SmallVector projectedExchangeStates; + CommunicationPhase currentProjectedCommunicationPhase = CommunicationPhase::LowToHigh; + llvm::SmallVector deferredProjectedExchanges; + llvm::SmallVector deferredProducerValues; AvailableValueStore availableValues; llvm::DenseMap hostReplacements; llvm::DenseMap hostOutputOwners; @@ -290,9 +298,6 @@ struct MaterializerState { llvm::SmallVector materializedScalarPeerReceives; llvm::SmallVector pendingScalarSends; llvm::SmallVector pendingProjectedScalarSends; - llvm::SmallVector pendingProjectedInputSends; - llvm::DenseSet projectedInputPhaseBarrierClasses; - llvm::DenseMap pendingProjectedHighToLowReceives; llvm::DenseSet oldComputeOps; MaterializerState(mlir::func::FuncOp func, const MergeScheduleResult& schedule, int64_t& nextChannelId)