From 8338caf3f3dfef3cd306f5599bd66f3ac7f31ef0 Mon Sep 17 00:00:00 2001 From: NiccoloN Date: Tue, 7 Jul 2026 12:54:34 +0200 Subject: [PATCH] cose brutte --- .../MergeComputeNodes/CommunicationPlan.cpp | 722 +++- .../MergeComputeNodes/CommunicationPlan.hpp | 166 +- .../MaterializeMergeSchedule.cpp | 3679 ++++++++--------- .../MaterializedClassState.hpp | 85 +- 4 files changed, 2482 insertions(+), 2170 deletions(-) diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.cpp index 88fba20..9b46571 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.cpp @@ -3,86 +3,255 @@ #include "mlir/IR/Diagnostics.h" #include "llvm/ADT/DenseSet.h" -#include "llvm/ADT/TypeSwitch.h" +#include "llvm/ADT/STLExtras.h" + +#include using namespace mlir; namespace onnx_mlir::spatial { namespace { -struct GroupSignature { - CommunicationPhase phase = CommunicationPhase::Local; - ClassId sourceClass = 0; - ClassId targetClass = 0; - Type payloadType; - bool sourceIsBatch = false; - bool targetIsBatch = false; - - bool operator==(const GroupSignature& other) const { - return phase == other.phase && sourceClass == other.sourceClass && targetClass == other.targetClass - && payloadType == other.payloadType && sourceIsBatch == other.sourceIsBatch - && targetIsBatch == other.targetIsBatch; +StringRef stringifyExchangeKind(CommunicationExchangeKind kind) { + switch (kind) { + case CommunicationExchangeKind::HostForward: return "HostForward"; + case CommunicationExchangeKind::OrdinaryValue: return "OrdinaryValue"; + case CommunicationExchangeKind::PackedScalarRun: return "PackedScalarRun"; + case CommunicationExchangeKind::IndexedBatchRun: return "IndexedBatchRun"; + case CommunicationExchangeKind::ProjectedInput: return "ProjectedInput"; } -}; + llvm_unreachable("unknown communication exchange kind"); +} -struct GroupSignatureInfo { - static GroupSignature getEmptyKey() { - return {CommunicationPhase::Local, - llvm::DenseMapInfo::getEmptyKey(), - llvm::DenseMapInfo::getEmptyKey(), - llvm::DenseMapInfo::getEmptyKey(), - false, - false}; +unsigned getExchangeKindOrder(CommunicationExchangeKind kind) { + switch (kind) { + case CommunicationExchangeKind::HostForward: return 0; + case CommunicationExchangeKind::OrdinaryValue: return 1; + case CommunicationExchangeKind::PackedScalarRun: return 2; + case CommunicationExchangeKind::IndexedBatchRun: return 3; + case CommunicationExchangeKind::ProjectedInput: return 4; } - - static GroupSignature getTombstoneKey() { - return {CommunicationPhase::Local, - llvm::DenseMapInfo::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; } -}; + llvm_unreachable("unknown communication exchange kind"); +} void emitPlanError(Operation* anchor, StringRef message) { if (anchor) anchor->emitError(message); } +struct PeerPair { + int64_t minCore = 0; + int64_t maxCore = 0; + + bool operator==(const PeerPair& other) const { + return minCore == other.minCore && maxCore == other.maxCore; + } +}; + +struct PeerPairInfo { + static PeerPair getEmptyKey() { + return {llvm::DenseMapInfo::getEmptyKey(), llvm::DenseMapInfo::getEmptyKey()}; + } + static PeerPair getTombstoneKey() { + return {llvm::DenseMapInfo::getTombstoneKey(), llvm::DenseMapInfo::getTombstoneKey()}; + } + static unsigned getHashValue(PeerPair key) { return llvm::hash_combine(key.minCore, key.maxCore); } + static bool isEqual(PeerPair lhs, PeerPair rhs) { return lhs == rhs; } +}; + +PeerPair getPeerPair(const CommunicationExchange& exchange) { + return {std::min(exchange.sourceCore, exchange.targetCore), std::max(exchange.sourceCore, exchange.targetCore)}; +} + +bool exchangeOrderLess(const CommunicationExchange& lhs, const CommunicationExchange& rhs) { + if (getExchangeKindOrder(lhs.kind) != getExchangeKindOrder(rhs.kind)) + return getExchangeKindOrder(lhs.kind) < getExchangeKindOrder(rhs.kind); + if (lhs.sourceClass != rhs.sourceClass) + return lhs.sourceClass < rhs.sourceClass; + if (lhs.targetClass != rhs.targetClass) + return lhs.targetClass < rhs.targetClass; + if (lhs.producerOrder != rhs.producerOrder) + return lhs.producerOrder < rhs.producerOrder; + if (lhs.producer.resultIndex != rhs.producer.resultIndex) + return lhs.producer.resultIndex < rhs.producer.resultIndex; + if (lhs.producer.instance.laneStart != rhs.producer.instance.laneStart) + return lhs.producer.instance.laneStart < rhs.producer.instance.laneStart; + if (lhs.producer.instance.laneCount != rhs.producer.instance.laneCount) + return lhs.producer.instance.laneCount < rhs.producer.instance.laneCount; + if (lhs.sourceLane != rhs.sourceLane) + return lhs.sourceLane < rhs.sourceLane; + if (lhs.targetLane != rhs.targetLane) + return lhs.targetLane < rhs.targetLane; + if (lhs.ordinal != rhs.ordinal) + return lhs.ordinal < rhs.ordinal; + return lhs.channelId < rhs.channelId; +} + +void sortExchangeIds(ArrayRef exchanges, + SmallVectorImpl& ids) { + llvm::stable_sort(ids, [&](CommunicationExchangeId lhs, CommunicationExchangeId rhs) { + return exchangeOrderLess(exchanges[lhs.value], exchanges[rhs.value]); + }); +} + +void addEvent(CommunicationPlan& plan, + DenseMap>& eventsByClass, + CommunicationExchangeId id, + CommunicationEventKind kind, + ClassId classId, + int64_t core, + unsigned& order) { + eventsByClass[classId].push_back(CommunicationEvent {id, kind, classId, core, order++}); + (void)plan; +} + +void buildEventStreams(CommunicationPlan& plan, + SmallVectorImpl& exchanges, + DenseMap>& eventsByClass) { + DenseMap, PeerPairInfo> exchangesByPair; + for (const CommunicationExchange& exchange : exchanges) + exchangesByPair[getPeerPair(exchange)].push_back(exchange.id); + + SmallVector pairs; + pairs.reserve(exchangesByPair.size()); + for (const auto& entry : exchangesByPair) + pairs.push_back(entry.first); + llvm::sort(pairs, [](PeerPair lhs, PeerPair rhs) { + if (lhs.minCore != rhs.minCore) + return lhs.minCore < rhs.minCore; + return lhs.maxCore < rhs.maxCore; + }); + + unsigned order = 0; + for (PeerPair pair : pairs) { + SmallVector hiToLo; + SmallVector loToHi; + for (CommunicationExchangeId id : exchangesByPair.lookup(pair)) { + const CommunicationExchange& exchange = exchanges[id.value]; + if (exchange.sourceCore == pair.maxCore && exchange.targetCore == pair.minCore) + hiToLo.push_back(id); + else + loToHi.push_back(id); + } + + sortExchangeIds(exchanges, hiToLo); + sortExchangeIds(exchanges, loToHi); + + for (CommunicationExchangeId id : hiToLo) { + const CommunicationExchange& exchange = exchanges[id.value]; + addEvent(plan, eventsByClass, id, CommunicationEventKind::Receive, exchange.targetClass, exchange.targetCore, order); + } + for (CommunicationExchangeId id : hiToLo) { + const CommunicationExchange& exchange = exchanges[id.value]; + addEvent(plan, eventsByClass, id, CommunicationEventKind::Send, exchange.sourceClass, exchange.sourceCore, order); + } + for (CommunicationExchangeId id : loToHi) { + const CommunicationExchange& exchange = exchanges[id.value]; + addEvent(plan, eventsByClass, id, CommunicationEventKind::Send, exchange.sourceClass, exchange.sourceCore, order); + } + for (CommunicationExchangeId id : loToHi) { + const CommunicationExchange& exchange = exchanges[id.value]; + addEvent(plan, eventsByClass, id, CommunicationEventKind::Receive, exchange.targetClass, exchange.targetCore, order); + } + } +} + } // namespace FailureOr -CommunicationPlan::buildProjectedInputPlan(Operation* anchor, - ArrayRef classes, - const DenseMap& cpuToClass, - const ProjectedInputPlanMap& projectedPlans) { +CommunicationPlan::build(Operation* anchor, + ArrayRef classes, + const DenseMap& cpuToClass, + const ProducerSourceClassMap& producerSourceClasses, + const DenseMap& producerPayloadTypes, + const DenseMap, ProducerKeyInfo>& producerDestClasses, + const DenseMap, ProducerKeyInfo>& ordinaryProducerDestClasses, + ProjectedInputPlanMap& projectedPlans, + const DenseMap& hostOutputOwners, + const DenseMap& liveExternalUseCache, + int64_t& nextChannelId) { CommunicationPlan plan; DenseMap classById; for (CommunicationClassInfo klass : classes) - classById[klass.classId] = klass; + classById[klass.classId] = std::move(klass); - DenseSet remoteChannels; - DenseSet seenFragments; - DenseMap groupBySignature; - unsigned nextPlaceholder = 0; + auto appendExchange = [&](CommunicationExchangeKind kind, + ProducerKey producer, + ClassId sourceClass, + ClassId targetClass, + Type payloadType, + int64_t sourceCore, + int64_t targetCore, + unsigned sourceLane, + unsigned targetLane, + unsigned ordinal, + Value originalHostOutput = Value(), + const ProjectedInputTransferFragment* fragment = nullptr) -> FailureOr { + if (sourceCore == targetCore) + return failure(); - for (const auto& extractEntry : projectedPlans) { - for (const auto& classEntry : extractEntry.second) { - ClassId targetClass = classEntry.first; - const ProjectedInputTransferPlan& inputPlan = classEntry.second; + auto sourceInfoIt = classById.find(sourceClass); + auto targetInfoIt = classById.find(targetClass); + if (sourceInfoIt == classById.end() || targetInfoIt == classById.end()) + return failure(); + + CommunicationExchangeId exchangeId {plan.exchanges.size()}; + plan.exchanges.push_back(CommunicationExchange { + exchangeId, + kind, + PlaceholderId {static_cast(plan.demands.size())}, + producer, + sourceClass, + targetClass, + sourceInfoIt->second.rank, + targetInfoIt->second.rank, + payloadType, + nextChannelId++, + sourceCore, + targetCore, + sourceLane, + targetLane, + ordinal, + exchangeId.value, + originalHostOutput, + fragment, + }); + return exchangeId; + }; + + auto getPayloadType = [&](ProducerKey producer) -> Type { + auto typeIt = producerPayloadTypes.find(producer); + if (typeIt != producerPayloadTypes.end()) + return typeIt->second; + if (!producer.instance.op || producer.resultIndex >= producer.instance.op->getNumResults()) + return Type(); + return producer.instance.op->getResult(producer.resultIndex).getType(); + }; + + for (const auto& entry : producerDestClasses) { + ProducerKey producer = entry.first; + auto sourceClassIt = producerSourceClasses.find(producer); + if (sourceClassIt == producerSourceClasses.end()) { + emitPlanError(anchor, "communication plan could not find producer source class"); + return failure(); + } + ClassId sourceClass = sourceClassIt->second; + auto sourceInfoIt = classById.find(sourceClass); + if (sourceInfoIt == classById.end()) { + emitPlanError(anchor, "communication plan references an unknown producer source class"); + return failure(); + } + const CommunicationClassInfo& sourceInfo = sourceInfoIt->second; + Type payloadType = getPayloadType(producer); + if (!payloadType) { + emitPlanError(anchor, "communication plan could not determine producer payload type"); + return failure(); + } + + for (ClassId targetClass : entry.second) { + if (targetClass == sourceClass) + continue; auto targetInfoIt = classById.find(targetClass); if (targetInfoIt == classById.end()) { @@ -91,7 +260,190 @@ CommunicationPlan::buildProjectedInputPlan(Operation* anchor, } const CommunicationClassInfo& targetInfo = targetInfoIt->second; - for (const ProjectedInputTransferFragment& fragment : inputPlan.fragments) { + CommunicationExchangeKind kind = CommunicationExchangeKind::OrdinaryValue; + if (sourceInfo.isBatch && !targetInfo.isBatch) + kind = CommunicationExchangeKind::PackedScalarRun; + else if (sourceInfo.isBatch && targetInfo.isBatch) + kind = CommunicationExchangeKind::IndexedBatchRun; + else if (!llvm::is_contained(ordinaryProducerDestClasses.lookup(producer), targetClass)) + continue; + + SmallVector* bucket = nullptr; + ProducerTargetKey lookupKey {producer, targetClass}; + if (kind == CommunicationExchangeKind::OrdinaryValue) + bucket = &plan.ordinaryExchangesByProducerTarget[lookupKey]; + else if (kind == CommunicationExchangeKind::PackedScalarRun) + bucket = &plan.packedScalarRunExchangesByProducerTarget[lookupKey]; + else if (kind == CommunicationExchangeKind::IndexedBatchRun) + bucket = &plan.indexedBatchRunExchangesByProducerTarget[lookupKey]; + + if (!bucket) + continue; + + if (!sourceInfo.isBatch) { + if (sourceInfo.cpus.empty()) { + emitPlanError(anchor, "communication plan source class has no CPU"); + return failure(); + } + int64_t sourceCore = static_cast(sourceInfo.cpus.front()); + unsigned ordinal = 0; + for (auto [targetLane, targetCpu] : llvm::enumerate(targetInfo.cpus)) { + auto exchangeId = appendExchange(kind, + producer, + sourceClass, + targetClass, + payloadType, + sourceCore, + static_cast(targetCpu), + 0, + static_cast(targetLane), + ordinal++); + if (failed(exchangeId)) { + emitPlanError(anchor, "communication plan failed to append scalar source exchange"); + return failure(); + } + bucket->push_back(*exchangeId); + if (!targetInfo.isBatch) + break; + } + continue; + } + + if (sourceInfo.cpus.empty()) { + emitPlanError(anchor, "communication plan batch source class has no CPUs"); + return failure(); + } + if (targetInfo.isBatch && sourceInfo.cpus.size() != targetInfo.cpus.size()) { + emitPlanError(anchor, "communication plan cannot map batch classes of different sizes"); + return failure(); + } + + for (auto [sourceLane, sourceCpu] : llvm::enumerate(sourceInfo.cpus)) { + CpuId targetCpu = targetInfo.isBatch ? targetInfo.cpus[sourceLane] : targetInfo.cpus.front(); + auto exchangeId = appendExchange(kind, + producer, + sourceClass, + targetClass, + payloadType, + static_cast(sourceCpu), + static_cast(targetCpu), + static_cast(sourceLane), + targetInfo.isBatch ? static_cast(sourceLane) : 0, + static_cast(bucket->size())); + if (failed(exchangeId)) { + emitPlanError(anchor, "communication plan failed to append batch source exchange"); + return failure(); + } + bucket->push_back(*exchangeId); + } + } + } + + for (const auto& entry : producerSourceClasses) { + ProducerKey producer = entry.first; + ClassId sourceClass = entry.second; + if (!producer.instance.op || producer.resultIndex >= producer.instance.op->getNumResults()) + continue; + + Value originalOutput = producer.instance.op->getResult(producer.resultIndex); + auto liveIt = liveExternalUseCache.find(originalOutput); + if (liveIt == liveExternalUseCache.end() || !liveIt->second) + continue; + + auto ownerIt = hostOutputOwners.find(originalOutput); + if (ownerIt == hostOutputOwners.end()) + continue; + ClassId ownerClass = ownerIt->second; + if (ownerClass == sourceClass) + continue; + + auto sourceInfoIt = classById.find(sourceClass); + auto ownerInfoIt = classById.find(ownerClass); + if (sourceInfoIt == classById.end() || ownerInfoIt == classById.end()) { + emitPlanError(anchor, "communication plan references an unknown host-forward class"); + return failure(); + } + const CommunicationClassInfo& sourceInfo = sourceInfoIt->second; + const CommunicationClassInfo& ownerInfo = ownerInfoIt->second; + if (sourceInfo.isBatch != ownerInfo.isBatch) { + emitPlanError(anchor, "communication plan does not support mixed scalar/batch host forwarding"); + return failure(); + } + if (sourceInfo.isBatch && sourceInfo.cpus.size() != ownerInfo.cpus.size()) { + emitPlanError(anchor, "communication plan cannot host-forward between batch classes of different sizes"); + return failure(); + } + + Type payloadType = getPayloadType(producer); + if (!payloadType) { + emitPlanError(anchor, "communication plan could not determine host-forward payload type"); + return failure(); + } + + HostForwardKey lookupKey {originalOutput, sourceClass, ownerClass}; + SmallVector& bucket = plan.hostForwardExchangesByOutput[lookupKey]; + if (!sourceInfo.isBatch) { + if (!bucket.empty()) + continue; + auto exchangeId = appendExchange(CommunicationExchangeKind::HostForward, + producer, + sourceClass, + ownerClass, + payloadType, + static_cast(sourceInfo.cpus.front()), + static_cast(ownerInfo.cpus.front()), + 0, + 0, + 0, + originalOutput); + if (failed(exchangeId)) { + emitPlanError(anchor, "communication plan failed to append scalar host-forward exchange"); + return failure(); + } + bucket.push_back(*exchangeId); + continue; + } + + if (bucket.size() >= sourceInfo.cpus.size()) + continue; + unsigned ordinal = static_cast(bucket.size()); + if (producer.instance.laneStart >= sourceInfo.cpus.size()) { + emitPlanError(anchor, "communication plan host-forward producer lane is out of range"); + return failure(); + } + unsigned lane = producer.instance.laneStart; + auto exchangeId = appendExchange(CommunicationExchangeKind::HostForward, + producer, + sourceClass, + ownerClass, + payloadType, + static_cast(sourceInfo.cpus[lane]), + static_cast(ownerInfo.cpus[lane]), + lane, + lane, + ordinal, + originalOutput); + if (failed(exchangeId)) { + emitPlanError(anchor, "communication plan failed to append batch host-forward exchange"); + return failure(); + } + bucket.push_back(*exchangeId); + } + + DenseSet seenFragments; + for (auto& extractEntry : projectedPlans) { + for (auto& classEntry : extractEntry.second) { + ClassId targetClass = classEntry.first; + ProjectedInputTransferPlan& inputPlan = classEntry.second; + + auto targetInfoIt = classById.find(targetClass); + if (targetInfoIt == classById.end()) { + emitPlanError(anchor, "communication plan references an unknown target class"); + return failure(); + } + const CommunicationClassInfo& targetInfo = targetInfoIt->second; + + for (ProjectedInputTransferFragment& fragment : inputPlan.fragments) { if (!seenFragments.insert(&fragment).second) { emitPlanError(anchor, "communication plan found the same projected fragment twice"); return failure(); @@ -111,129 +463,193 @@ CommunicationPlan::buildProjectedInputPlan(Operation* anchor, } const CommunicationClassInfo& sourceInfo = sourceInfoIt->second; - PlaceholderId placeholder {nextPlaceholder++}; - plan.demands.push_back(InputDemand {placeholder, fragment.producer, targetClass, inputPlan.layout.fragmentType}); - plan.publications.push_back( - ProducerPublication {fragment.producer, sourceClass, inputPlan.layout.fragmentType}); + plan.demands.push_back(InputDemand {{static_cast(plan.demands.size())}, + fragment.producer, + targetClass, + inputPlan.layout.fragmentType}); + plan.publications.push_back(ProducerPublication {fragment.producer, sourceClass, inputPlan.layout.fragmentType}); - if (fragment.sourceCoreId == fragment.targetCoreId) { - plan.phaseByFragment[&fragment] = CommunicationPhase::Local; + if (fragment.sourceCoreId == fragment.targetCoreId) continue; - } if (sourceInfo.rank == targetInfo.rank && sourceClass != targetClass) { emitPlanError(anchor, "communication plan found a same-rank inter-class projected exchange"); return failure(); } - if (!remoteChannels.insert(fragment.channelId).second) { - emitPlanError(anchor, "communication plan found a duplicate remote projected channel id"); - return failure(); - } - - TransferKind kind = sourceClass == targetClass ? TransferKind::SameClassForward : TransferKind::RemoteChannel; - CommunicationPhase phase = - sourceInfo.rank == targetInfo.rank - ? (fragment.sourceCoreId < fragment.targetCoreId ? CommunicationPhase::LowToHigh - : CommunicationPhase::HighToLow) - : (sourceInfo.rank < targetInfo.rank ? CommunicationPhase::LowToHigh : CommunicationPhase::HighToLow); - plan.phaseByFragment[&fragment] = phase; - - size_t exchangeId = plan.exchanges.size(); - plan.exchangeByFragment[&fragment] = exchangeId; - plan.exchanges.push_back(ExchangeDescriptor { - exchangeId, - kind, - placeholder, + FailureOr exchangeId = appendExchange( + CommunicationExchangeKind::ProjectedInput, fragment.producer, sourceClass, targetClass, - sourceInfo.rank, - targetInfo.rank, inputPlan.layout.fragmentType, - fragment.channelId, fragment.sourceCoreId, fragment.targetCoreId, + 0, fragment.targetLane, fragment.ordinal, - phase, - &fragment, - }); - - GroupSignature signature { - phase, - sourceClass, - targetClass, - inputPlan.layout.fragmentType, - sourceInfo.isBatch, - targetInfo.isBatch, - }; - auto groupIt = groupBySignature.find(signature); - if (groupIt == groupBySignature.end()) { - unsigned groupIndex = static_cast(plan.batchGroups.size()); - groupBySignature[signature] = groupIndex; - plan.batchGroups.push_back(CommunicationBatchGroup { - phase, - sourceClass, - targetClass, - inputPlan.layout.fragmentType, - sourceInfo.isBatch, - targetInfo.isBatch, - {}, - }); - groupIt = groupBySignature.find(signature); + Value(), + &fragment); + if (failed(exchangeId)) { + emitPlanError(anchor, "communication plan failed to append projected input exchange"); + return failure(); } - plan.batchGroups[groupIt->second].exchanges.push_back(exchangeId); + fragment.channelId = plan.exchanges[exchangeId->value].channelId; + plan.projectedExchangesByFragment[&fragment].push_back(*exchangeId); } } } - for (const ExchangeDescriptor& exchange : plan.exchanges) { - if (!exchange.payloadType) { - emitPlanError(anchor, "communication plan has an exchange without payload type"); - return failure(); - } - if (exchange.phase == CommunicationPhase::Local) { - emitPlanError(anchor, "communication plan classified a remote exchange as local"); - return failure(); - } - if (exchange.kind == TransferKind::SameClassForward && exchange.sourceClass != exchange.targetClass) { - emitPlanError(anchor, "communication plan same-class exchange crosses class ids"); - return failure(); - } - if (exchange.kind == TransferKind::RemoteChannel && exchange.sourceRank == exchange.targetRank) { - emitPlanError(anchor, "communication plan remote exchange has equal source and target rank"); - return failure(); - } - if (exchange.sourceRank < exchange.targetRank && exchange.phase != CommunicationPhase::LowToHigh) { - emitPlanError(anchor, "communication plan phase does not match low-to-high rank direction"); - return failure(); - } - if (exchange.sourceRank > exchange.targetRank && exchange.phase != CommunicationPhase::HighToLow) { - emitPlanError(anchor, "communication plan phase does not match high-to-low rank direction"); - return failure(); - } - if (!exchange.fragment) { - emitPlanError(anchor, "communication plan exchange is not bound to a projected fragment"); - return failure(); - } - } - + buildEventStreams(plan, plan.exchanges, plan.eventsByClass); + if (failed(plan.verify(anchor))) + return failure(); return plan; } -std::optional -CommunicationPlan::getPhase(const ProjectedInputTransferFragment& fragment) const { - auto it = phaseByFragment.find(&fragment); - if (it == phaseByFragment.end()) - return std::nullopt; +LogicalResult CommunicationPlan::verify(Operation* anchor) const { + DenseSet channelIds; + SmallVector malformed; + size_t malformedCount = 0; + + auto recordMalformed = [&](const CommunicationExchange& exchange) { + ++malformedCount; + if (malformed.size() < 8) + malformed.push_back(&exchange); + }; + + for (auto [index, exchange] : llvm::enumerate(exchanges)) { + bool bad = false; + if (exchange.id.value != index) + bad = true; + if (exchange.sourceCore == exchange.targetCore) + bad = true; + if (!exchange.payloadType) + bad = true; + if (!channelIds.insert(exchange.channelId).second) + bad = true; + if (bad) + recordMalformed(exchange); + } + + DenseMap sendCounts; + DenseMap receiveCounts; + for (const auto& classEvents : eventsByClass) { + std::optional previousOrder; + DenseSet classOrders; + for (const CommunicationEvent& event : classEvents.second) { + if (event.exchangeId.value >= exchanges.size()) { + if (anchor) + anchor->emitError("communication plan event references an invalid exchange id"); + return failure(); + } + + const CommunicationExchange& exchange = exchanges[event.exchangeId.value]; + if (previousOrder && event.order < *previousOrder) { + if (anchor) + anchor->emitError("communication plan event stream is not sorted"); + return failure(); + } + previousOrder = event.order; + if (!classOrders.insert(event.order).second) { + if (anchor) + anchor->emitError("communication plan event stream contains duplicate event order"); + return failure(); + } + + if (event.kind == CommunicationEventKind::Send) { + ++sendCounts[event.exchangeId.value]; + if (event.classId != exchange.sourceClass || event.core != exchange.sourceCore) + recordMalformed(exchange); + } + else { + ++receiveCounts[event.exchangeId.value]; + if (event.classId != exchange.targetClass || event.core != exchange.targetCore) + recordMalformed(exchange); + } + } + } + + for (const CommunicationExchange& exchange : exchanges) { + if (sendCounts[exchange.id.value] != 1 || receiveCounts[exchange.id.value] != 1) + recordMalformed(exchange); + } + + auto verifyLookup = [&](ArrayRef ids) -> LogicalResult { + for (CommunicationExchangeId id : ids) + if (id.value >= exchanges.size()) + return failure(); + return success(); + }; + + for (const auto& entry : projectedExchangesByFragment) + if (failed(verifyLookup(entry.second))) { + if (anchor) + anchor->emitError("communication plan projected lookup references an invalid exchange id"); + return failure(); + } + + if (malformedCount != 0) { + if (anchor) { + InFlightDiagnostic diag = anchor->emitError("communication plan contains malformed exchanges"); + for (const CommunicationExchange* exchange : malformed) { + diag << " [exchange=" << exchange->id.value << " kind=" << stringifyExchangeKind(exchange->kind) + << " sourceClass=" << exchange->sourceClass << " sourceCore=" << exchange->sourceCore + << " targetClass=" << exchange->targetClass << " targetCore=" << exchange->targetCore + << " channelId=" << exchange->channelId << " payloadType=" << exchange->payloadType << "]"; + } + if (malformedCount > malformed.size()) + diag << " additionalMalformed=" << (malformedCount - malformed.size()); + } + return failure(); + } + + return success(); +} + +ArrayRef CommunicationPlan::getEventsForClass(ClassId classId) const { + auto it = eventsByClass.find(classId); + if (it == eventsByClass.end()) + return {}; return it->second; } -std::optional CommunicationPlan::getExchangeId(const ProjectedInputTransferFragment& fragment) const { - auto it = exchangeByFragment.find(&fragment); - if (it == exchangeByFragment.end()) - return std::nullopt; +ArrayRef +CommunicationPlan::getOrdinaryValueExchanges(ProducerKey producer, ClassId targetClass) const { + auto it = ordinaryExchangesByProducerTarget.find({producer, targetClass}); + if (it == ordinaryExchangesByProducerTarget.end()) + return {}; + return it->second; +} + +ArrayRef +CommunicationPlan::getProjectedInputExchanges(const ProjectedInputTransferFragment& fragment) const { + auto it = projectedExchangesByFragment.find(&fragment); + if (it == projectedExchangesByFragment.end()) + return {}; + return it->second; +} + +ArrayRef +CommunicationPlan::getPackedScalarRunExchanges(ProducerKey producer, ClassId targetClass) const { + auto it = packedScalarRunExchangesByProducerTarget.find({producer, targetClass}); + if (it == packedScalarRunExchangesByProducerTarget.end()) + return {}; + return it->second; +} + +ArrayRef +CommunicationPlan::getIndexedBatchRunExchanges(ProducerKey producer, ClassId targetClass) const { + auto it = indexedBatchRunExchangesByProducerTarget.find({producer, targetClass}); + if (it == indexedBatchRunExchangesByProducerTarget.end()) + return {}; + return it->second; +} + +ArrayRef +CommunicationPlan::getHostForwardExchanges(Value originalOutput, ClassId sourceClass, ClassId ownerClass) const { + auto it = hostForwardExchangesByOutput.find({originalOutput, sourceClass, ownerClass}); + if (it == hostForwardExchangesByOutput.end()) + return {}; return it->second; } diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.hpp index 509dedb..d109545 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.hpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/CommunicationPlan.hpp @@ -5,6 +5,7 @@ #include "mlir/Support/LLVM.h" #include "llvm/ADT/DenseMap.h" +#include "llvm/ADT/DenseSet.h" #include "llvm/ADT/SmallVector.h" #include @@ -24,10 +25,31 @@ enum class TransferKind { WholeTensorBarrier }; -enum class CommunicationPhase { - Local, - LowToHigh, - HighToLow +enum class CommunicationExchangeKind { + OrdinaryValue, + ProjectedInput, + PackedScalarRun, + IndexedBatchRun, + HostForward +}; + +enum class CommunicationEventKind { + Send, + Receive +}; + +struct CommunicationExchangeId { + size_t value = 0; + + bool operator==(const CommunicationExchangeId& other) const { return value == other.value; } + operator size_t() const { return value; } +}; + +struct CommunicationExchangeIdInfo { + static CommunicationExchangeId getEmptyKey() { return {llvm::DenseMapInfo::getEmptyKey()}; } + static CommunicationExchangeId getTombstoneKey() { return {llvm::DenseMapInfo::getTombstoneKey()}; } + static unsigned getHashValue(CommunicationExchangeId id) { return llvm::hash_value(id.value); } + static bool isEqual(CommunicationExchangeId lhs, CommunicationExchangeId rhs) { return lhs == rhs; } }; struct PlaceholderId { @@ -61,61 +83,147 @@ struct CommunicationClassInfo { size_t rank = 0; bool isBatch = false; mlir::Operation* op = nullptr; + llvm::SmallVector cpus; }; -struct ExchangeDescriptor { - size_t id = 0; - TransferKind kind = TransferKind::RemoteChannel; +struct CommunicationExchange { + CommunicationExchangeId id; + CommunicationExchangeKind kind = CommunicationExchangeKind::OrdinaryValue; + PlaceholderId placeholder; ProducerKey producer; ClassId sourceClass = 0; ClassId targetClass = 0; + size_t sourceRank = 0; size_t targetRank = 0; mlir::Type payloadType; int64_t channelId = 0; int64_t sourceCore = 0; int64_t targetCore = 0; + + unsigned sourceLane = 0; unsigned targetLane = 0; unsigned ordinal = 0; - CommunicationPhase phase = CommunicationPhase::Local; + uint64_t producerOrder = 0; + + mlir::Value originalHostOutput; const ProjectedInputTransferFragment* fragment = nullptr; }; -struct CommunicationBatchGroup { - CommunicationPhase phase = CommunicationPhase::Local; - ClassId sourceClass = 0; - ClassId targetClass = 0; - mlir::Type payloadType; - bool sourceIsBatch = false; - bool targetIsBatch = false; - llvm::SmallVector exchanges; +using ExchangeDescriptor = CommunicationExchange; + +struct CommunicationEvent { + CommunicationExchangeId exchangeId; + CommunicationEventKind kind = CommunicationEventKind::Send; + ClassId classId = 0; + int64_t core = 0; + unsigned order = 0; }; class CommunicationPlan { public: using ProjectedInputPlanMap = llvm::DenseMap>; + using ProducerSourceClassMap = llvm::DenseMap; static mlir::FailureOr - buildProjectedInputPlan(mlir::Operation* anchor, - llvm::ArrayRef classes, - const llvm::DenseMap& cpuToClass, - const ProjectedInputPlanMap& projectedPlans); + build(mlir::Operation* anchor, + llvm::ArrayRef classes, + const llvm::DenseMap& cpuToClass, + const ProducerSourceClassMap& producerSourceClasses, + const llvm::DenseMap& producerPayloadTypes, + const llvm::DenseMap, ProducerKeyInfo>& producerDestClasses, + const llvm::DenseMap, ProducerKeyInfo>& ordinaryProducerDestClasses, + ProjectedInputPlanMap& projectedPlans, + const llvm::DenseMap& hostOutputOwners, + const llvm::DenseMap& liveExternalUseCache, + int64_t& nextChannelId); - 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; } + mlir::LogicalResult verify(mlir::Operation* anchor) const; + + llvm::ArrayRef getEventsForClass(ClassId classId) const; + llvm::ArrayRef getExchanges() const { return exchanges; } + const CommunicationExchange& getExchange(CommunicationExchangeId id) const { return exchanges[id.value]; } + const CommunicationExchange& getExchange(size_t id) const { return exchanges[id]; } + + llvm::ArrayRef + getOrdinaryValueExchanges(ProducerKey producer, ClassId targetClass) const; + llvm::ArrayRef + getProjectedInputExchanges(const ProjectedInputTransferFragment& fragment) const; + llvm::ArrayRef + getPackedScalarRunExchanges(ProducerKey producer, ClassId targetClass) const; + llvm::ArrayRef + getIndexedBatchRunExchanges(ProducerKey producer, ClassId targetClass) const; + llvm::ArrayRef + getHostForwardExchanges(mlir::Value originalOutput, ClassId sourceClass, ClassId ownerClass) const; private: - llvm::DenseMap phaseByFragment; - llvm::DenseMap exchangeByFragment; + struct ProducerTargetKey { + ProducerKey producer; + ClassId targetClass = 0; + + bool operator==(const ProducerTargetKey& other) const { + return producer == other.producer && targetClass == other.targetClass; + } + }; + + struct ProducerTargetKeyInfo { + static ProducerTargetKey getEmptyKey() { + return {ProducerKeyInfo::getEmptyKey(), llvm::DenseMapInfo::getEmptyKey()}; + } + static ProducerTargetKey getTombstoneKey() { + return {ProducerKeyInfo::getTombstoneKey(), llvm::DenseMapInfo::getTombstoneKey()}; + } + static unsigned getHashValue(const ProducerTargetKey& key) { + return llvm::hash_combine(ProducerKeyInfo::getHashValue(key.producer), key.targetClass); + } + static bool isEqual(const ProducerTargetKey& lhs, const ProducerTargetKey& rhs) { return lhs == rhs; } + }; + + struct HostForwardKey { + mlir::Value originalOutput; + ClassId sourceClass = 0; + ClassId ownerClass = 0; + + bool operator==(const HostForwardKey& other) const { + return originalOutput == other.originalOutput && sourceClass == other.sourceClass + && ownerClass == other.ownerClass; + } + }; + + struct HostForwardKeyInfo { + static HostForwardKey getEmptyKey() { + return {llvm::DenseMapInfo::getEmptyKey(), + llvm::DenseMapInfo::getEmptyKey(), + llvm::DenseMapInfo::getEmptyKey()}; + } + static HostForwardKey getTombstoneKey() { + return {llvm::DenseMapInfo::getTombstoneKey(), + llvm::DenseMapInfo::getTombstoneKey(), + llvm::DenseMapInfo::getTombstoneKey()}; + } + static unsigned getHashValue(const HostForwardKey& key) { + return llvm::hash_combine(key.originalOutput, key.sourceClass, key.ownerClass); + } + static bool isEqual(const HostForwardKey& lhs, const HostForwardKey& rhs) { return lhs == rhs; } + }; + llvm::SmallVector demands; llvm::SmallVector publications; - llvm::SmallVector exchanges; - llvm::SmallVector batchGroups; + llvm::SmallVector exchanges; + llvm::DenseMap> eventsByClass; + + llvm::DenseMap> + projectedExchangesByFragment; + llvm::DenseMap, ProducerTargetKeyInfo> + ordinaryExchangesByProducerTarget; + llvm::DenseMap, ProducerTargetKeyInfo> + packedScalarRunExchangesByProducerTarget; + llvm::DenseMap, ProducerTargetKeyInfo> + indexedBatchRunExchangesByProducerTarget; + llvm::DenseMap, HostForwardKeyInfo> + hostForwardExchangesByOutput; }; } // 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 bac6c85..05dfec2 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializeMergeSchedule.cpp @@ -5,6 +5,7 @@ #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/BuiltinOps.h" +#include "mlir/IR/Dominance.h" #include "mlir/IR/IRMapping.h" #include "mlir/IR/Matchers.h" #include "mlir/IR/PatternMatch.h" @@ -81,8 +82,6 @@ FailureOr materializeTensorValueForMaterializedClassUse(MaterializerState IRMapping* mapper = nullptr); FailureOr materializeWholeBatchInput( MaterializerState& state, MaterializedClass& targetClass, ProducerKey key, Type resultType, Location loc); -ScalarPeerEdgeKey makeScalarPeerEdgeKey(int64_t sourceCore, int64_t targetCore, Type payloadType); -bool hasPlannedScalarPeerReceive(const MaterializerState& state, const ScalarPeerEdgeKey& key); FailureOr localizeMaterializedClassOperand(MaterializerState& state, MaterializedClass& targetClass, Value value, @@ -118,6 +117,15 @@ LogicalResult emitPlannedProjectedInputSendsForKeys(MaterializerState& state, ArrayRef keys, Value payload, Location loc); +static FailureOr emitReadyCommunicationEventsForClass(MaterializerState& state, + MaterializedClass& materializedClass, + Location loc); +static Operation* getProjectedReceiveInsertionAnchor(MaterializedClass& targetClass); +static LogicalResult materializeProjectedInputReceivesForClass(MaterializerState& state, + MaterializedClass& targetClass, + Operation* requestedExtract, + Location loc, + Operation* insertionAnchor); static void tagProjectedInputTransfer(Operation* op); std::optional getProjectedInputSliceMatch(MaterializerState& state, SpatComputeBatch batch, unsigned inputIndex); @@ -125,15 +133,36 @@ std::optional getProjectedWholeBatchReplacementProducer(MaterializerState& state, SpatComputeBatch batch, unsigned inputIndex); std::optional getProjectedWholeBatchReplacementProducer(MaterializerState& state, tensor::ExtractSliceOp extract); +bool hasProjectedInputReplacement(MaterializerState& state, + SpatComputeBatch batch, + unsigned inputIndex, + Value input, + ComputeInstance logicalConsumer, + ClassId classId); LogicalResult recordLocalPackedSlotValue(MaterializerState& state, MaterializedClass& materializedClass, ArrayRef keys, Value packed); +Operation* getTopLevelMaterializedClassBodyOp(Operation* nestedOp, MaterializedClass& targetClass); +LogicalResult registerPackedRunValue(MaterializerState& state, + MaterializedClass& materializedClass, + ArrayRef keys, + Value packed, + Type fragmentType, + Location loc); +FailureOr getPackedSlotFragmentType(Value packed, size_t keyCount); +SmallVector& getBatchOutputFragmentTypesCached(MaterializerState& state, SpatComputeBatch batch); +static bool isProducerValueMaterialized(MaterializerState& state, ProducerKey key); +static void recordDeferredProducerValue(MaterializerState& state, ProducerKey key); FailureOr materializeProjectedWholeBatchExtractReplacement(MaterializerState& state, MaterializedClass& targetClass, tensor::ExtractSliceOp extract, ProducerKey producer, IRMapping* mapper = nullptr); +FailureOr> materializeLocalProjectedInputExtractReplacement(MaterializerState& state, + MaterializedClass& targetClass, + tensor::ExtractSliceOp extract, + IRMapping* mapper = nullptr); enum class MaterializedLayoutKind { DenseLogical, RowPackedNCHWFromRows, @@ -275,12 +304,6 @@ FailureOr getPackedBatchTensorType(Type laneType, size_t laneC return RankedTensorType::get(shape, tensorType.getElementType(), tensorType.getEncoding()); } -LogicalResult verifyPackableFragmentType(Operation* anchor, Type fragmentType, size_t count, StringRef message) { - if (failed(getPackedBatchTensorType(fragmentType, count))) - return anchor->emitError(message); - return success(); -} - ComputeInstance getScheduledChunkForLogicalInstance(MaterializerState& state, ComputeInstance logicalInstance) { auto it = state.logicalInstanceToScheduledChunk.find(logicalInstance); if (it != state.logicalInstanceToScheduledChunk.end()) @@ -465,6 +488,61 @@ bool canUseProjectedLaneInput(MaterializerState& state, return !collectProjectedFragmentDemandsForBatchInput(state, consumerBatch, inputIndex, input, logicalConsumer).empty(); } +static bool isDebugPaddedVggInput(Value input) { + auto type = dyn_cast(input.getType()); + return type && type.hasStaticShape() && type.getRank() == 4 && type.getDimSize(0) == 1 + && type.getDimSize(1) == 3 && type.getDimSize(2) == 226 && type.getDimSize(3) == 226; +} + +static bool hasProjectedPlanForBatchInputKey(MaterializerState& state, + SpatComputeBatch consumerBatch, + unsigned inputIndex, + ClassId classId) { + ProjectedBatchInputKey inputKey {consumerBatch.getOperation(), inputIndex}; + for (const auto& extractEntry : state.projectedInputTransferPlans) { + auto classIt = extractEntry.second.find(classId); + if (classIt != extractEntry.second.end() && classIt->second.inputKey == inputKey) + return true; + } + + for (const auto& producerEntry : state.projectedTransfers) { + auto classIt = producerEntry.second.find(classId); + if (classIt != producerEntry.second.end() && classIt->second.inputKey == inputKey) + return true; + } + + return false; +} + +static SmallVector collectProjectedInputKeyProducersForClass(MaterializerState& state, + SpatComputeBatch consumerBatch, + unsigned inputIndex, + ClassId classId) { + ProjectedBatchInputKey inputKey {consumerBatch.getOperation(), inputIndex}; + SmallVector producers; + + auto appendProducer = [&](ProducerKey producer) { + if (!llvm::is_contained(producers, producer)) + producers.push_back(producer); + }; + + for (const auto& producerEntry : state.projectedTransfers) { + auto classIt = producerEntry.second.find(classId); + if (classIt != producerEntry.second.end() && classIt->second.inputKey == inputKey) + appendProducer(producerEntry.first); + } + + for (const auto& extractEntry : state.projectedInputTransferPlans) { + auto classIt = extractEntry.second.find(classId); + if (classIt == extractEntry.second.end() || !(classIt->second.inputKey == inputKey)) + continue; + for (const ProjectedInputTransferFragment& fragment : classIt->second.fragments) + appendProducer(fragment.producer); + } + + return producers; +} + static bool hasUnreplacedBatchInputUse(MaterializerState& state, SpatComputeBatch consumerBatch, unsigned inputIndex, @@ -481,13 +559,77 @@ static bool hasUnreplacedBatchInputUse(MaterializerState& state, auto planIt = state.projectedInputTransferPlans.find(extract.getOperation()); if (planIt != state.projectedInputTransferPlans.end() && planIt->second.find(classId) != planIt->second.end()) continue; + if (hasProjectedPlanForBatchInputKey(state, consumerBatch, inputIndex, classId)) + continue; if (getProjectedWholeBatchReplacementProducer(state, extract)) continue; + if (inputIndex == 0) { + static unsigned debugCount = 0; + if (debugCount++ < 16) { + llvm::errs() << "debug unreplaced batch input use classId=" << classId + << " consumer='" << consumerBatch->getName() << "'" + << " useOwner='" << use.getOwner()->getName() << "'" + << " isExtract=" << static_cast(extract) + << " hasPlan=" << (planIt != state.projectedInputTransferPlans.end()) + << " planHasClass=" + << (planIt != state.projectedInputTransferPlans.end() + && planIt->second.find(classId) != planIt->second.end()) + << " hasWholeReplacement=" + << getProjectedWholeBatchReplacementProducer(state, extract).has_value() + << "\n"; + } + } return true; } return false; } +static bool hasUnreplacedComputeInputUse(MaterializerState& state, + SpatCompute compute, + unsigned inputIndex, + ClassId classId) { + auto inputArg = compute.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; + return true; + } + return false; +} + +static bool hasProjectedComputeInputReplacement(MaterializerState& state, + SpatCompute compute, + unsigned inputIndex, + ClassId classId) { + auto inputArg = compute.getInputArgument(inputIndex); + if (!inputArg) + return false; + + for (OpOperand& use : inputArg->getUses()) { + auto extract = dyn_cast(use.getOwner()); + if (!extract) + continue; + + auto replacementIt = state.projectedExtractReplacements.find(extract.getOperation()); + if (replacementIt != state.projectedExtractReplacements.end() + && replacementIt->second.find(classId) != replacementIt->second.end()) + return true; + + auto planIt = state.projectedInputTransferPlans.find(extract.getOperation()); + if (planIt != state.projectedInputTransferPlans.end() && planIt->second.find(classId) != planIt->second.end()) + return true; + } + return false; +} + FailureOr classifyComputeBatchInputDemand(MaterializerState& state, MaterializedClass& targetClass, SpatComputeBatch consumerBatch, @@ -690,7 +832,8 @@ static LogicalResult appendOrdinaryInputBatchingSignatureForCpu(MaterializerStat for (auto [inputIndex, input] : llvm::enumerate(consumerInputs)) { SmallVector producerKeys; if (auto batchConsumer = dyn_cast(logicalConsumer.op)) { - if (canUseProjectedLaneInput(state, batchConsumer, static_cast(inputIndex), input, logicalConsumer) + if (hasProjectedInputReplacement( + state, batchConsumer, static_cast(inputIndex), input, logicalConsumer, state.cpuToClass.lookup(cpu)) && !hasUnreplacedBatchInputUse(state, batchConsumer, static_cast(inputIndex), state.cpuToClass.lookup(cpu))) continue; @@ -1908,56 +2051,6 @@ FailureOr createDim0ExtractSliceInClass(MaterializerState& state, return createDim0ExtractSlice(state, loc, *localizedSource, *localizedOffset, firstSize); } -Value createStaticExtractSlice(MaterializerState& state, - Location loc, - Value source, - ArrayRef sliceOffsets, - ArrayRef resultShape) { - auto sourceType = cast(source.getType()); - assert(sliceOffsets.size() == static_cast(sourceType.getRank()) && "offset rank mismatch"); - assert(resultShape.size() == static_cast(sourceType.getRank()) && "result rank mismatch"); - SmallVector offsets; - SmallVector sizes; - SmallVector strides; - offsets.reserve(sourceType.getRank()); - sizes.reserve(sourceType.getRank()); - strides.reserve(sourceType.getRank()); - - for (int64_t dim = 0; dim < sourceType.getRank(); ++dim) { - offsets.push_back(sliceOffsets[dim]); - sizes.push_back(state.rewriter.getIndexAttr(resultShape[dim])); - strides.push_back(state.rewriter.getIndexAttr(1)); - } - - return tensor::ExtractSliceOp::create(state.rewriter, loc, source, offsets, sizes, strides).getResult(); -} - -FailureOr createStaticExtractSliceInClass(MaterializerState& state, - MaterializedClass& targetClass, - Location loc, - Value source, - ArrayRef sliceOffsets, - ArrayRef resultShape) { - FailureOr localizedSource = materializeTensorValueForMaterializedClassUse( - state, - targetClass, - source, - targetClass.op, - "createStaticExtractSliceInClass tried to reuse a tensor from another materialized class"); - if (failed(localizedSource)) - return failure(); - - SmallVector localizedOffsets; - localizedOffsets.reserve(sliceOffsets.size()); - for (OpFoldResult offset : sliceOffsets) { - FailureOr localized = rematerializeIndexOpFoldResultInClass(state, targetClass, offset, loc); - if (failed(localized)) - return failure(); - localizedOffsets.push_back(*localized); - } - return createStaticExtractSlice(state, loc, *localizedSource, localizedOffsets, resultShape); -} - Value createIndexedIndexValue(MaterializerState& state, Operation* anchor, ArrayRef values, @@ -1966,28 +2059,6 @@ Value createIndexedIndexValue(MaterializerState& state, std::optional preferredPeriod = std::nullopt, bool allowExhaustiveTiledSearch = true); -FailureOr> -buildProjectedFragmentOffsetsInClass(MaterializerState& state, - MaterializedClass& targetClass, - const ProjectedTransferDescriptor& descriptor, - Value flatFragmentIndex, - Location loc) { - FailureOr localizedIndex = rematerializeIndexValueInClass(state, targetClass, flatFragmentIndex, loc); - if (failed(localizedIndex)) - return failure(); - SmallVector fragmentOffsets; - fragmentOffsets.reserve(descriptor.layout.fragmentShape.size()); - for (ArrayRef dimOffsets : descriptor.fragmentOffsetsByDim) - fragmentOffsets.push_back(createIndexedIndexValue(state, - targetClass.op, - dimOffsets, - *localizedIndex, - loc, - static_cast(descriptor.layout.payloadFragmentCount), - /*allowExhaustiveTiledSearch=*/false)); - return fragmentOffsets; -} - Value createDim0InsertSlice( MaterializerState& state, Location loc, Value fragment, Value destination, OpFoldResult firstOffset) { auto fragmentType = cast(fragment.getType()); @@ -2143,6 +2214,25 @@ Value createIndexedSourceCoreId( MaterializerState& state, Operation* anchor, const MessageVector& messages, Value index, Location loc); Value createIndexedTargetCoreId( MaterializerState& state, Operation* anchor, const MessageVector& messages, Value index, Location loc); +Value appendReceive( + MaterializerState& state, MaterializedClass& targetClass, Type type, const MessageVector& messages, Location loc); +FailureOr buildMessageVectorForExchangeIds(MaterializerState& state, + ArrayRef ids, + CommunicationExchangeKind expectedKind, + ClassId expectedSourceClass, + ClassId expectedTargetClass, + Type expectedPayloadType, + Operation* anchor); +LogicalResult bindCommunicationPayload(MaterializerState& state, + CommunicationExchangeId id, + MaterializedClass& sourceClass, + Value payload, + Operation* anchor, + Location loc); +FailureOr emitPlannedReceive(MaterializerState& state, + CommunicationExchangeId id, + MaterializedClass& targetClass, + Location loc); Value getPackedSliceForRunIndex(MaterializerState& state, Operation* anchor, @@ -2223,23 +2313,7 @@ struct ReceiveMessagePartition { std::optional getConstantIndexValue(Value value); -ReceiveMessagePartition partitionReceivesByPendingSendUnlocks(const MaterializerState& state, - const MaterializedClass& targetClass, - Type payloadType, - const MessageVector& messages, - bool preserveMessageOrder = false); -ScalarPeerReceiveKey makeScalarPeerReceiveKey(int64_t sourceCore, - int64_t targetCore, - Type payloadType, - std::optional channelId = std::nullopt); -FailureOr createScalarPeerReceiveAndFlush(MaterializerState& state, - MaterializedClass& targetClass, - Type payloadType, - Value channelId, - Value sourceCoreId, - Value targetCoreId, - std::optional staticKey, - Location loc); +ReceiveMessagePartition partitionReceiveMessagesInPlanOrder(const MessageVector& messages); MessageVector filterMessageVector(const MessageVector& messages, ArrayRef indices); bool isDeferredLocalPackedScalarRun(const PackedScalarRunValue& run) { @@ -2262,13 +2336,13 @@ LogicalResult validatePackedScalarRunMetadata(Operation* anchor, const PackedSca if (receiveCount == 0) return anchor->emitError("packed scalar run has no receives"); - if (failed(run.messages.verify(anchor))) - return failure(); + if (!run.exchangeIds.empty()) { + if (run.exchangeIds.size() != receiveCount) + return anchor->emitError("packed scalar run exchange id count is inconsistent"); + return success(); + } - if (run.messages.size() != receiveCount) - return anchor->emitError("packed scalar run receive metadata count is inconsistent"); - - return success(); + return anchor->emitError("packed scalar run deferred receive is missing planned exchange ids"); } FailureOr materializePackedScalarRunValue(MaterializerState& state, @@ -2348,6 +2422,18 @@ FailureOr materializePackedScalarRunValue(MaterializerState& state, if (failed(validatePackedScalarRunMetadata(targetClass.op, run))) return failure(); + FailureOr builtMessages = + buildMessageVectorForExchangeIds(state, + run.exchangeIds, + CommunicationExchangeKind::PackedScalarRun, + state.communicationPlan.getExchange(run.exchangeIds.front()).sourceClass, + targetClass.id, + run.fragmentType, + targetClass.op); + if (failed(builtMessages)) + return failure(); + const MessageVector& receiveMessages = *builtMessages; + FailureOr fullPackedType = getPackedBatchTensorType(run.fragmentType, getPackedScalarRunReceiveCount(run)); if (failed(fullPackedType)) @@ -2359,21 +2445,17 @@ FailureOr materializePackedScalarRunValue(MaterializerState& state, .getResult(); Value packedInit = init; ReceiveMessagePartition partition = - partitionReceivesByPendingSendUnlocks(state, targetClass, run.fragmentType, run.messages, /*preserveMessageOrder=*/true); + partitionReceiveMessagesInPlanOrder(receiveMessages); for (size_t index : partition.criticalStaticIndices) { state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - FailureOr received = createScalarPeerReceiveAndFlush( - state, - targetClass, - run.fragmentType, - getOrCreateIndexConstant(state.constantFolder, targetClass.op, run.messages.channelIds[index]), - getOrCreateIndexConstant(state.constantFolder, targetClass.op, run.messages.sourceCoreIds[index]), - getOrCreateIndexConstant(state.constantFolder, targetClass.op, run.messages.targetCoreIds[index]), - makeScalarPeerReceiveKey( - run.messages.sourceCoreIds[index], run.messages.targetCoreIds[index], run.fragmentType, run.messages.channelIds[index]), - loc); - if (failed(received)) - return failure(); + Value received = SpatChannelReceiveOp::create( + state.rewriter, + loc, + run.fragmentType, + getOrCreateIndexConstant(state.constantFolder, targetClass.op, receiveMessages.channelIds[index]), + getOrCreateIndexConstant(state.constantFolder, targetClass.op, receiveMessages.sourceCoreIds[index]), + getOrCreateIndexConstant(state.constantFolder, targetClass.op, receiveMessages.targetCoreIds[index])) + .getOutput(); FailureOr offset = scaleIndexByDim0SizeInClass( state, targetClass, @@ -2382,13 +2464,13 @@ FailureOr materializePackedScalarRunValue(MaterializerState& state, loc); if (failed(offset)) return failure(); - FailureOr updated = createDim0InsertSliceInClass(state, targetClass, loc, *received, packedInit, *offset); + FailureOr updated = createDim0InsertSliceInClass(state, targetClass, loc, received, packedInit, *offset); if (failed(updated)) return failure(); packedInit = *updated; } - auto remainingMessages = filterMessageVector(run.messages, partition.remainingIndices); + auto remainingMessages = filterMessageVector(receiveMessages, partition.remainingIndices); SmallVector remainingPackedRowOffsets; remainingPackedRowOffsets.reserve(partition.remainingIndices.size()); for (size_t originalIndex : partition.remainingIndices) @@ -2499,6 +2581,52 @@ std::optional AvailableValueStore::lookup(MaterializerState& state, Produ if (std::optional exact = lookupExact(key, classId)) return exact; + auto exchangeProducerIt = exactExchangeIds.find(key); + if (exchangeProducerIt != exactExchangeIds.end()) { + auto exchangeClassIt = exchangeProducerIt->second.find(classId); + if (exchangeClassIt != exchangeProducerIt->second.end() && !exchangeClassIt->second.empty()) { + ArrayRef ids = exchangeClassIt->second; + const CommunicationExchange& firstExchange = state.communicationPlan.getExchange(ids.front()); + bool ready = true; + for (CommunicationExchangeId id : ids) { + if (id.value >= state.communicationExchangeStates.size() + || !state.communicationExchangeStates[id.value].sendEmitted) { + ready = false; + break; + } + } + if (!ready) + return std::nullopt; + + MaterializedClass& materializedClass = state.classes[classId]; + FailureOr messages = buildMessageVectorForExchangeIds(state, + ids, + firstExchange.kind, + firstExchange.sourceClass, + firstExchange.targetClass, + Type(), + materializedClass.op); + if (failed(messages)) + return std::nullopt; + Value received = + appendReceive(state, materializedClass, firstExchange.payloadType, *messages, materializedClass.op->getLoc()); + if (!received) + return std::nullopt; + + for (CommunicationExchangeId id : ids) { + CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[id.value]; + emission.receiveEmitted = true; + emission.receiveValue = received; + if (Operation* receiveDef = received.getDefiningOp()) + emission.receiveAnchor = getTopLevelMaterializedClassBodyOp(receiveDef, materializedClass); + } + + record(key, classId, received); + exchangeProducerIt->second.erase(exchangeClassIt); + return received; + } + } + if (std::optional packedRunValue = lookupPackedRun(state, key, classId)) return packedRunValue; @@ -2880,23 +3008,45 @@ LogicalResult collectProducerDestinations(MaterializerState& state) { SmallVector producerKeys; bool projectedPlanOwnsFanout = false; if (auto batchConsumer = dyn_cast(logicalConsumer.op)) { - SmallVector projectedKeys = collectProjectedProducerKeysForBatchInput( - state, batchConsumer, static_cast(inputIndex), input, logicalConsumer); + bool hasProjectedReplacement = + hasProjectedInputReplacement( + state, batchConsumer, static_cast(inputIndex), input, logicalConsumer, targetClass); + bool hasUnreplacedUse = + hasUnreplacedBatchInputUse(state, batchConsumer, static_cast(inputIndex), targetClass); + if (isDebugPaddedVggInput(input)) { + static unsigned debugCount = 0; + if (debugCount++ < 12) { + llvm::errs() << "debug padded destination targetClass=" << targetClass + << " inputIndex=" << inputIndex + << " hasProjectedReplacement=" << hasProjectedReplacement + << " hasUnreplacedUse=" << hasUnreplacedUse + << " canUseProjectedLane=" + << canUseProjectedLaneInput( + state, batchConsumer, static_cast(inputIndex), input, logicalConsumer) + << " hasWholeBatchReplacement=" + << getProjectedWholeBatchReplacementProducer(state, batchConsumer, static_cast(inputIndex)) + .has_value() + << "\n"; + } + } projectedPlanOwnsFanout = - !state.classes[targetClass].isBatch && !projectedKeys.empty() - && !hasUnreplacedBatchInputUse(state, batchConsumer, static_cast(inputIndex), targetClass); + hasProjectedReplacement && !hasUnreplacedUse; producerKeys = projectedPlanOwnsFanout - ? std::move(projectedKeys) + ? collectProjectedInputKeyProducersForClass( + state, batchConsumer, static_cast(inputIndex), targetClass) : collectProducerKeysForBatchInputDestinations( state, batchConsumer, static_cast(inputIndex), input, logicalConsumer); } else { + if (auto computeConsumer = dyn_cast(logicalConsumer.op)) { + if (hasProjectedComputeInputReplacement( + state, computeConsumer, static_cast(inputIndex), targetClass) + && !hasUnreplacedComputeInputUse(state, computeConsumer, static_cast(inputIndex), targetClass)) + continue; + } producerKeys = collectProducerKeysForDestinations(input, logicalConsumer); } - if (projectedPlanOwnsFanout) - continue; - for (ProducerKey producerKey : producerKeys) { ComputeInstance scheduledProducer = getScheduledChunkForLogicalInstance(state, producerKey.instance); auto producerCpuIt = state.schedule.computeToCpuMap.find(scheduledProducer); @@ -2906,6 +3056,8 @@ LogicalResult collectProducerDestinations(MaterializerState& state) { ClassId sourceClass = state.cpuToClass.lookup(producerCpuIt->second); if (sourceClass == targetClass) { + if (projectedPlanOwnsFanout) + continue; SameClassConsumerLookupKey lookupKey {producerKey.instance.op, producerKey.resultIndex, targetClass}; SmallVector& bucket = state.sameClassConsumerIndex[lookupKey]; if (!llvm::is_contained(bucket, producerKey)) @@ -2913,7 +3065,7 @@ LogicalResult collectProducerDestinations(MaterializerState& state) { continue; } - appendDestinationClass(state, producerKey, targetClass); + appendDestinationClass(state, producerKey, targetClass, /*ordinary=*/!projectedPlanOwnsFanout); } } @@ -3762,7 +3914,6 @@ LogicalResult collectProjectedTransfers(MaterializerState& state) { fragment.ordinal = targetClass.isBatch ? targetLane * plan.layout.fragmentsPerLogicalSlot + demand.ordinal : demand.ordinal; - fragment.channelId = state.nextChannelId++; fragment.sourceCoreId = *sourceCpu; fragment.targetCoreId = *targetCpu; plan.fragments.push_back(std::move(fragment)); @@ -3918,7 +4069,7 @@ LogicalResult collectProjectedTransfers(MaterializerState& state) { return success(); } -LogicalResult buildProjectedInputCommunicationPlan(MaterializerState& state) { +LogicalResult buildGlobalCommunicationPlan(MaterializerState& state) { SmallVector classes; classes.reserve(state.classes.size()); for (const MaterializedClass& klass : state.classes) { @@ -3927,80 +4078,55 @@ LogicalResult buildProjectedInputCommunicationPlan(MaterializerState& state) { .rank = klass.id, .isBatch = klass.isBatch, .op = klass.op, + .cpus = klass.cpus, }); } - 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(); -} + CommunicationPlan::ProducerSourceClassMap producerSourceClasses; + DenseMap producerPayloadTypes; + for (const auto& [producer, destinations] : state.producerDestClasses) { + (void)destinations; + ComputeInstance scheduledProducer = getScheduledChunkForLogicalInstance(state, producer.instance); + auto cpuIt = state.schedule.computeToCpuMap.find(scheduledProducer); + if (cpuIt == state.schedule.computeToCpuMap.end()) + return producer.instance.op->emitError("global communication planning found an unscheduled producer"); + auto classIt = state.cpuToClass.find(cpuIt->second); + if (classIt == state.cpuToClass.end()) + return producer.instance.op->emitError("global communication planning could not map producer CPU to class"); + producerSourceClasses[producer] = classIt->second; -static std::optional -collectScalarTargetProjectedDescriptor(MaterializerState& state, - MaterializedClass& targetClass, - ArrayRef keys, - bool requirePackedRunOffsetCountMatch) { - assert(!targetClass.isBatch && "scalar target projected descriptor helper expects a scalar class"); - - std::optional combined; - for (ProducerKey key : keys) { - auto producerIt = state.projectedTransfers.find(key); - if (producerIt == state.projectedTransfers.end()) - return std::nullopt; - - auto descriptorIt = producerIt->second.find(targetClass.id); - if (descriptorIt == producerIt->second.end()) - return std::nullopt; - - const ProjectedTransferDescriptor& descriptor = descriptorIt->second; - if (descriptor.fragmentOffsets.empty()) - return std::nullopt; - if (descriptor.layout.payloadFragmentCount == 0 || descriptor.layout.fragmentsPerLogicalSlot == 0) - return std::nullopt; - if (descriptor.fragmentOffsets.size() != descriptor.layout.payloadFragmentCount) - return std::nullopt; - if (descriptor.layout.payloadFragmentCount % descriptor.layout.fragmentsPerLogicalSlot != 0) - return std::nullopt; - - if (!combined) { - combined = descriptor; + if (auto batch = dyn_cast_or_null(producer.instance.op)) { + SmallVector& fragmentTypes = getBatchOutputFragmentTypesCached(state, batch); + if (producer.resultIndex >= fragmentTypes.size() || !fragmentTypes[producer.resultIndex]) + return producer.instance.op->emitError("global communication planning could not recover batch fragment type"); + producerPayloadTypes[producer] = fragmentTypes[producer.resultIndex]; continue; } - if (!(combined->inputKey == descriptor.inputKey) || combined->extractOp != descriptor.extractOp - || combined->layout.fragmentType != descriptor.layout.fragmentType - || combined->layout.fragmentShape != descriptor.layout.fragmentShape - || combined->layout.fragmentsPerLogicalSlot != descriptor.layout.fragmentsPerLogicalSlot - || combined->layout.loopLowerBounds != descriptor.layout.loopLowerBounds - || combined->layout.loopSteps != descriptor.layout.loopSteps - || combined->layout.loopTripCounts != descriptor.layout.loopTripCounts) - return std::nullopt; - - combined->layout.payloadFragmentCount += descriptor.layout.payloadFragmentCount; - llvm::append_range(combined->fragmentOffsets, descriptor.fragmentOffsets); + if (producer.resultIndex >= producer.instance.op->getNumResults()) + return producer.instance.op->emitError("global communication planning found an invalid producer result index"); + producerPayloadTypes[producer] = producer.instance.op->getResult(producer.resultIndex).getType(); } - if (!combined) - return std::nullopt; - - if (combined->fragmentOffsets.size() != combined->layout.payloadFragmentCount) - return std::nullopt; - - if (requirePackedRunOffsetCountMatch) { - if (combined->layout.payloadFragmentCount != keys.size() * combined->layout.fragmentsPerLogicalSlot) - return std::nullopt; - } - if (failed(finalizeProjectedTransferDescriptor(targetClass.op, *combined))) - return std::nullopt; - return combined; + FailureOr plan = CommunicationPlan::build(state.func, + classes, + state.cpuToClass, + producerSourceClasses, + producerPayloadTypes, + state.producerDestClasses, + state.ordinaryProducerDestClasses, + state.projectedInputTransferPlans, + state.hostOutputOwners, + state.liveExternalUseCache, + state.nextChannelId); + if (failed(plan)) + return failure(); + state.communicationPlan = std::move(*plan); + state.communicationExchangeStates.assign(state.communicationPlan.getExchanges().size(), + CommunicationExchangeEmissionState {}); + return success(); } - ArrayRef getDestinationClasses(MaterializerState& state, ProducerKey key) { auto it = state.producerDestClasses.find(key); if (it == state.producerDestClasses.end()) @@ -4013,14 +4139,6 @@ static bool hasOrdinaryDestinationClass(MaterializerState& state, ProducerKey ke 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) { @@ -4048,76 +4166,25 @@ void appendScalarSendAtCurrentInsertionPoint(MaterializerState& state, SpatChannelSendOp::create(state.rewriter, loc, channelIdValue, sourceCoreIdValue, targetCoreIdValue, payload); } -bool isMaterializedScalarCore(const MaterializerState& state, int64_t coreId) { - return llvm::any_of(state.classes, [&](const MaterializedClass& klass) { - return !klass.isBatch && klass.cpus.size() == 1 && static_cast(klass.cpus.front()) == coreId; - }); -} - -bool shouldEmitScalarPeerSendStatically(const MaterializerState& state, - const MaterializedClass& sourceClass, - int64_t sourceCoreId, - int64_t targetCoreId) { - if (sourceClass.isBatch || sourceCoreId >= targetCoreId) - return false; - return isMaterializedScalarCore(state, targetCoreId); -} - LogicalResult appendScalarSendLoopAtCurrentInsertionPoint(MaterializerState& state, MaterializedClass& sourceClass, Value payload, const MessageVector& messages, Location loc) { - MessageVector loopMessages; - MessageVector delayedPeerMessages; - loopMessages.channelIds.reserve(messages.size()); - loopMessages.sourceCoreIds.reserve(messages.size()); - loopMessages.targetCoreIds.reserve(messages.size()); - delayedPeerMessages.channelIds.reserve(messages.size()); - delayedPeerMessages.sourceCoreIds.reserve(messages.size()); - delayedPeerMessages.targetCoreIds.reserve(messages.size()); - - for (size_t index = 0; index < messages.size(); ++index) { - int64_t sourceCoreId = messages.sourceCoreIds[index]; - int64_t targetCoreId = messages.targetCoreIds[index]; - if (shouldEmitScalarPeerSendStatically(state, sourceClass, sourceCoreId, targetCoreId)) { - delayedPeerMessages.append(messages.channelIds[index], sourceCoreId, targetCoreId); - continue; - } - loopMessages.append(messages.channelIds[index], sourceCoreId, targetCoreId); - } - - auto emitStaticPeerSends = [&]() { - for (size_t index = 0; index < delayedPeerMessages.size(); ++index) - appendScalarSendAtCurrentInsertionPoint(state, - sourceClass, - payload, - delayedPeerMessages.channelIds[index], - static_cast(delayedPeerMessages.sourceCoreIds[index]), - static_cast(delayedPeerMessages.targetCoreIds[index]), - loc); - }; - - if (loopMessages.empty()) { - emitStaticPeerSends(); - return success(); - } - - if (loopMessages.size() == 1) { + if (messages.size() == 1) { appendScalarSendAtCurrentInsertionPoint(state, sourceClass, payload, - loopMessages.channelIds.front(), - static_cast(loopMessages.sourceCoreIds.front()), - static_cast(loopMessages.targetCoreIds.front()), + messages.channelIds.front(), + static_cast(messages.sourceCoreIds.front()), + static_cast(messages.targetCoreIds.front()), loc); - emitStaticPeerSends(); return success(); } Value lowerBound = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, 0); Value upperBound = - getOrCreateIndexConstant(state.constantFolder, sourceClass.op, static_cast(loopMessages.size())); + getOrCreateIndexConstant(state.constantFolder, sourceClass.op, static_cast(messages.size())); Value step = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, 1); auto sendLoop = buildNormalizedScfFor( @@ -4128,16 +4195,14 @@ LogicalResult appendScalarSendLoopAtCurrentInsertionPoint(MaterializerState& sta step, ValueRange {}, [&](OpBuilder&, Location, Value index, ValueRange, SmallVectorImpl&) { - Value channelId = createIndexedChannelId(state, sourceClass.op, loopMessages, index, loc); - Value sourceCoreId = createIndexedSourceCoreId(state, sourceClass.op, loopMessages, index, loc); - Value targetCoreId = createIndexedTargetCoreId(state, sourceClass.op, loopMessages, index, loc); + Value channelId = createIndexedChannelId(state, sourceClass.op, messages, index, loc); + Value sourceCoreId = createIndexedSourceCoreId(state, sourceClass.op, messages, index, loc); + Value targetCoreId = createIndexedTargetCoreId(state, sourceClass.op, messages, index, loc); SpatChannelSendOp::create(state.rewriter, loc, channelId, sourceCoreId, targetCoreId, payload); return success(); }); if (failed(sendLoop)) return failure(); - - emitStaticPeerSends(); return success(); } @@ -4167,182 +4232,6 @@ LogicalResult appendScalarSendLoop(MaterializerState& state, return appendScalarSendLoopAtCurrentInsertionPoint(state, sourceClass, payload, messages, loc); } -FailureOr buildProjectedPackedPayload(MaterializerState& state, - MaterializedClass& targetClass, - Value fullPayload, - const ProjectedTransferDescriptor& descriptor, - Value messageIndex, - Location loc) { - if (failed(verifyProjectedTransferDescriptor(targetClass.op, descriptor))) - return failure(); - if (descriptor.layout.payloadFragmentCount == 1) - return targetClass.op->emitError("projected packed payload builder expects a packed payload"); - - FailureOr localizedPayload = materializeTensorValueForMaterializedClassUse( - state, - targetClass, - fullPayload, - targetClass.op, - "projected packed payload tried to reuse a tensor from another materialized class"); - if (failed(localizedPayload)) - return failure(); - FailureOr localizedMessageIndex = rematerializeIndexValueInClass(state, targetClass, messageIndex, loc); - if (failed(localizedMessageIndex)) - return failure(); - - Value init = tensor::EmptyOp::create( - state.rewriter, loc, descriptor.payloadType.getShape(), descriptor.payloadType.getElementType()) - .getResult(); - - Value lowerBound = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 0); - Value upperBound = - getOrCreateIndexConstant(state.constantFolder, targetClass.op, descriptor.layout.payloadFragmentCount); - Value step = getOrCreateIndexConstant(state.constantFolder, targetClass.op, 1); - - auto loop = buildNormalizedScfFor( - state.rewriter, - loc, - lowerBound, - upperBound, - step, - ValueRange {init}, - [&](OpBuilder&, Location, Value fragmentIndex, ValueRange iterArgs, SmallVectorImpl& yielded) { - Value acc = iterArgs.front(); - MLIRContext* context = state.func.getContext(); - AffineExpr d0 = getAffineDimExpr(0, context); - AffineExpr d1 = getAffineDimExpr(1, context); - AffineMap flatIndexMap = - AffineMap::get(/*dimCount=*/2, /*symbolCount=*/0, d0 * descriptor.layout.payloadFragmentCount + d1); - Value flatIndex = createOrFoldAffineApply( - state.rewriter, loc, flatIndexMap, ValueRange {*localizedMessageIndex, fragmentIndex}, targetClass.op); - - FailureOr> fragmentOffsets = - buildProjectedFragmentOffsetsInClass(state, targetClass, descriptor, flatIndex, loc); - if (failed(fragmentOffsets)) - return failure(); - FailureOr fragment = createStaticExtractSliceInClass( - state, targetClass, loc, *localizedPayload, *fragmentOffsets, descriptor.layout.fragmentShape); - if (failed(fragment)) - return failure(); - - FailureOr packedOffset = scaleIndexByDim0SizeInClass( - state, targetClass, fragmentIndex, descriptor.layout.fragmentType.getDimSize(0), loc); - if (failed(packedOffset)) - return failure(); - FailureOr next = createDim0InsertSliceInClass(state, targetClass, loc, *fragment, acc, *packedOffset); - if (failed(next)) - return failure(); - yielded.push_back(*next); - return success(); - }); - if (failed(loop)) - return failure(); - return loop->results.front(); -} - -FailureOr buildProjectedPayloadForMessage(MaterializerState& state, - MaterializedClass& targetClass, - Value fullPayload, - const ProjectedTransferDescriptor& descriptor, - Value messageIndex, - Location loc) { - if (failed(verifyProjectedTransferDescriptor(targetClass.op, descriptor))) - return failure(); - - FailureOr localizedPayload = materializeTensorValueForMaterializedClassUse( - state, - targetClass, - fullPayload, - targetClass.op, - "projected payload tried to reuse a tensor from another materialized class"); - if (failed(localizedPayload)) - return failure(); - - if (descriptor.layout.payloadFragmentCount == 1) { - FailureOr> fragmentOffsets = - buildProjectedFragmentOffsetsInClass(state, targetClass, descriptor, messageIndex, loc); - if (failed(fragmentOffsets)) - return failure(); - return createStaticExtractSliceInClass( - state, targetClass, loc, *localizedPayload, *fragmentOffsets, descriptor.layout.fragmentShape); - } - - return buildProjectedPackedPayload(state, targetClass, *localizedPayload, descriptor, messageIndex, loc); -} - -FailureOr getProjectedCommunicationType(Operation* anchor, const ProjectedTransferDescriptor& descriptor) { - if (descriptor.layout.payloadFragmentCount == 1) - return Type(descriptor.layout.fragmentType); - - FailureOr packedType = - getPackedBatchTensorType(descriptor.layout.fragmentType, descriptor.layout.payloadFragmentCount); - if (failed(packedType)) - return anchor->emitError("cannot compute projected communication payload type"), failure(); - return Type(*packedType); -} - -LogicalResult appendProjectedScalarSendLoopAtCurrentInsertionPoint(MaterializerState& state, - MaterializedClass& sourceClass, - Value payload, - const ProjectedTransferDescriptor& descriptor, - const MessageVector& messages, - Location loc) { - if (messages.size() == 1) { - Value channelId = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, messages.channelIds.front()); - Value sourceCoreId = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, messages.sourceCoreIds.front()); - Value targetCoreId = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, messages.targetCoreIds.front()); - Value messageIndex = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, 0); - FailureOr sendPayload = - buildProjectedPayloadForMessage(state, sourceClass, payload, descriptor, messageIndex, loc); - if (failed(sendPayload)) - return failure(); - SpatChannelSendOp::create(state.rewriter, loc, channelId, sourceCoreId, targetCoreId, *sendPayload); - return success(); - } - - Value lowerBound = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, 0); - Value upperBound = - getOrCreateIndexConstant(state.constantFolder, sourceClass.op, static_cast(messages.size())); - Value step = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, 1); - - auto projectedSendLoop = buildNormalizedScfFor( - state.rewriter, - loc, - lowerBound, - upperBound, - step, - ValueRange {}, - [&](OpBuilder&, Location, Value index, ValueRange, SmallVectorImpl&) { - Value channelId = createIndexedChannelId(state, sourceClass.op, messages, index, loc); - Value sourceCoreId = createIndexedSourceCoreId(state, sourceClass.op, messages, index, loc); - Value targetCoreId = createIndexedTargetCoreId(state, sourceClass.op, messages, index, loc); - FailureOr sendPayload = - buildProjectedPayloadForMessage(state, sourceClass, payload, descriptor, index, loc); - if (failed(sendPayload)) - return failure(); - SpatChannelSendOp::create(state.rewriter, loc, channelId, sourceCoreId, targetCoreId, *sendPayload); - return success(); - }); - if (failed(projectedSendLoop)) - return failure(); - return success(); -} - -LogicalResult appendProjectedScalarSendLoop(MaterializerState& state, - MaterializedClass& sourceClass, - Value payload, - const ProjectedTransferDescriptor& descriptor, - const MessageVector& messages, - Location loc) { - assert(!sourceClass.isBatch && "projected scalar send expects scalar source class"); - assert(succeeded(messages.verify(sourceClass.op)) && "message metadata is inconsistent"); - if (failed(verifyProjectedSendDescriptor(sourceClass.op, descriptor, messages))) - return failure(); - - state.rewriter.setInsertionPoint(sourceClass.body->getTerminator()); - return appendProjectedScalarSendLoopAtCurrentInsertionPoint(state, sourceClass, payload, descriptor, messages, loc); -} - LogicalResult appendSend(MaterializerState& state, MaterializedClass& sourceClass, Value payload, @@ -4425,178 +4314,152 @@ Value appendReceive( loc); } -LogicalResult registerLazyPackedScalarReceives(MaterializerState& state, - MaterializedClass& sourceClass, - MaterializedClass& targetClass, - ArrayRef keys, - Type fragmentType, - ArrayRef channelIds, - ArrayRef sourceCoreIds, - ArrayRef targetCoreIds) { - if (!sourceClass.isBatch) - return sourceClass.op->emitError("lazy packed scalar receives expect a batch source class"); - - if (targetClass.isBatch) - return targetClass.op->emitError("lazy packed scalar receives expect a scalar target class"); - - if (keys.empty()) - return sourceClass.op->emitError("lazy packed scalar receive expects at least one producer key"); - - if (keys.size() != sourceClass.cpus.size()) - return sourceClass.op->emitError("lazy packed scalar receive expects one producer key per source lane"); - +FailureOr buildMessageVectorForExchangeIds(MaterializerState& state, + ArrayRef ids, + CommunicationExchangeKind expectedKind, + ClassId expectedSourceClass, + ClassId expectedTargetClass, + Type expectedPayloadType, + Operation* anchor) { MessageVector messages; - messages.append(channelIds, sourceCoreIds, targetCoreIds); - if (failed(messages.verify(targetClass.op))) - return failure(); - - if (keys.size() != messages.size()) - return targetClass.op->emitError("lazy packed scalar receive metadata is inconsistent"); - - auto rankedFragmentType = dyn_cast(fragmentType); - if (!rankedFragmentType || !rankedFragmentType.hasStaticShape() || rankedFragmentType.getRank() == 0) - return targetClass.op->emitError("lazy packed scalar receive expects a static ranked fragment type"); - - if (failed(verifyPackableFragmentType( - targetClass.op, fragmentType, keys.size(), "cannot create lazy packed scalar receive type"))) - return failure(); - - Operation* sourceOp = keys.front().instance.op; - size_t resultIndex = keys.front().resultIndex; - - for (ProducerKey key : keys) { - if (key.instance.op != sourceOp || key.resultIndex != resultIndex) - return sourceClass.op->emitError("lazy packed scalar receive expects one producer result"); - - if (key.instance.laneCount != 1) - return sourceClass.op->emitError("lazy packed scalar receive expects one lane per producer key"); + messages.channelIds.reserve(ids.size()); + messages.sourceCoreIds.reserve(ids.size()); + messages.targetCoreIds.reserve(ids.size()); + for (CommunicationExchangeId id : ids) { + const CommunicationExchange& exchange = state.communicationPlan.getExchange(id); + if (exchange.kind != expectedKind || exchange.sourceClass != expectedSourceClass + || exchange.targetClass != expectedTargetClass) { + return anchor->emitError("planned communication exchange does not match requested endpoint"), failure(); + } + if (expectedPayloadType && exchange.payloadType && exchange.payloadType != expectedPayloadType) { + // Packed scalar fanout can bind a wider packed payload than the per-lane + // producer type known during planning; the actual receive type is supplied + // by the materialization site. + } + messages.append(exchange.channelId, + static_cast(exchange.sourceCore), + static_cast(exchange.targetCore)); } - - PackedScalarRunValue packedRun; - packedRun.targetClass = targetClass.id; - packedRun.sourceOp = sourceOp; - packedRun.resultIndex = resultIndex; - packedRun.kind = PackedScalarRunKind::DeferredReceive; - packedRun.fragmentType = rankedFragmentType; - - packedRun.messages = std::move(messages); - - packedRun.slots.reserve(keys.size()); - for (ProducerKey key : keys) { - PackedScalarRunSlot slot; - slot.keys.push_back(key); - packedRun.slots.push_back(std::move(slot)); - } - - if (failed(validatePackedScalarRunMetadata(targetClass.op, packedRun))) + if (failed(messages.verify(anchor))) return failure(); + return messages; +} + +LogicalResult bindCommunicationPayload(MaterializerState& state, + CommunicationExchangeId id, + MaterializedClass& sourceClass, + Value payload, + Operation* anchor, + Location loc) { + if (id.value >= state.communicationExchangeStates.size()) + return anchor->emitError("communication payload binding references an invalid exchange id"); + const CommunicationExchange& exchange = state.communicationPlan.getExchange(id); + if (exchange.sourceClass != sourceClass.id) + return anchor->emitError("communication payload binding source class mismatch"); + + CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[id.value]; + if (emission.producerBound && emission.producerPayload != payload) + return anchor->emitError("communication exchange was bound to two different payload values"); + + emission.producerBound = true; + emission.producerPayload = payload; + emission.loc = loc; + if (Operation* payloadDef = payload.getDefiningOp()) + emission.producerAnchor = getTopLevelMaterializedClassBodyOp(payloadDef, sourceClass); - state.availableValues.recordPackedRun(std::move(packedRun)); return success(); } -struct ScalarSourceReceivePlan { - ClassId targetClass = 0; - MessageVector messages; - Type receiveType; - Operation* projectedExtractOp = nullptr; - ProjectedFragmentLayout projectedLayout; -}; +static LogicalResult emitPlannedSend(MaterializerState& state, + CommunicationExchangeId id, + MaterializedClass& sourceClass, + Location loc) { + const CommunicationExchange& exchange = state.communicationPlan.getExchange(id); + CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[id.value]; + if (exchange.sourceClass != sourceClass.id) + return sourceClass.op->emitError("planned send source class mismatch"); + if (emission.sendEmitted) + return sourceClass.op->emitError("planned send emitted twice"); + if (!emission.producerBound) + return success(); -struct ProjectedScalarSendGroup { - MessageVector messages; - ProjectedTransferDescriptor descriptor; -}; + FailureOr messages = buildMessageVectorForExchangeIds(state, + ArrayRef(id), + exchange.kind, + exchange.sourceClass, + exchange.targetClass, + Type(), + sourceClass.op); + if (failed(messages)) + return failure(); -ProjectedScalarSendGroup orderedProjectedScalarSendGroup(const ProjectedScalarSendGroup& group); -LogicalResult finalizeSplitProjectedScalarSendGroup(Operation* anchor, ProjectedScalarSendGroup& group); - -struct PeerExchangeOrderKey { - int64_t maxCore = 0; - int64_t sourceCore = 0; - int64_t targetCore = 0; - int64_t channelId = 0; -}; - -PeerExchangeOrderKey makePeerExchangeOrderKey(int64_t sourceCore, int64_t targetCore, int64_t channelId) { - return PeerExchangeOrderKey { - .maxCore = std::max(sourceCore, targetCore), - .sourceCore = sourceCore, - .targetCore = targetCore, - .channelId = channelId, - }; -} - -PeerExchangeOrderKey makePeerExchangeOrderKey(const MessageVector& messages, size_t index) { - return makePeerExchangeOrderKey(messages.sourceCoreIds[index], messages.targetCoreIds[index], messages.channelIds[index]); -} - -bool peerExchangeOrderLess(PeerExchangeOrderKey lhs, PeerExchangeOrderKey rhs) { - if (lhs.maxCore != rhs.maxCore) - return lhs.maxCore > rhs.maxCore; - if (lhs.sourceCore != rhs.sourceCore) - return lhs.sourceCore > rhs.sourceCore; - if (lhs.targetCore != rhs.targetCore) - return lhs.targetCore > rhs.targetCore; - 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)); - }); -} - -MessageVector orderedMessageVector(const MessageVector& messages) { - SmallVector indices; - indices.reserve(messages.size()); - for (size_t index = 0; index < messages.size(); ++index) - indices.push_back(index); - sortMessageIndicesByPeerExchangeOrder(messages, indices); - - MessageVector ordered; - for (size_t index : indices) - ordered.append(messages.channelIds[index], messages.sourceCoreIds[index], messages.targetCoreIds[index]); - return ordered; -} - -std::optional getBestPeerExchangeOrderKey(const MessageVector& messages) { - if (messages.empty()) - return std::nullopt; - - PeerExchangeOrderKey best = makePeerExchangeOrderKey(messages, 0); - for (size_t index = 1; index < messages.size(); ++index) { - PeerExchangeOrderKey candidate = makePeerExchangeOrderKey(messages, index); - if (peerExchangeOrderLess(candidate, best)) - best = candidate; + if (sourceClass.isBatch) { + auto batch = dyn_cast(sourceClass.op); + if (!batch) + return sourceClass.op->emitError("planned batch send requires a scheduled compute_batch"); + std::optional laneArg = batch.getLaneArgument(); + if (!laneArg) + return sourceClass.op->emitError("planned batch send requires a lane argument"); + state.rewriter.setInsertionPoint(sourceClass.body->getTerminator()); + Value sourceLane = getOrCreateIndexConstant(state.constantFolder, sourceClass.op, exchange.sourceLane); + Value isSender = arith::CmpIOp::create(state.rewriter, loc, arith::CmpIPredicate::eq, *laneArg, sourceLane); + auto ifOp = scf::IfOp::create(state.rewriter, loc, TypeRange {}, isSender, false); + OpBuilder::InsertionGuard guard(state.rewriter); + Block& thenBlock = ifOp.getThenRegion().front(); + if (!thenBlock.empty()) + if (auto yieldOp = dyn_cast(thenBlock.back())) + yieldOp.erase(); + state.rewriter.setInsertionPointToEnd(&thenBlock); + SpatChannelSendOp::create(state.rewriter, + loc, + getOrCreateIndexConstant(state.constantFolder, sourceClass.op, exchange.channelId), + getOrCreateIndexConstant(state.constantFolder, sourceClass.op, exchange.sourceCore), + getOrCreateIndexConstant(state.constantFolder, sourceClass.op, exchange.targetCore), + emission.producerPayload); + scf::YieldOp::create(state.rewriter, loc); } - return best; + else if (failed(appendSend(state, sourceClass, emission.producerPayload, *messages, loc))) { + return failure(); + } + emission.sendEmitted = true; + if (Operation* payloadDef = emission.producerPayload.getDefiningOp()) + emission.sendAnchor = getTopLevelMaterializedClassBodyOp(payloadDef, sourceClass); + return success(); } -void sortBatchRunSendPlansByPeerExchangeOrder(SmallVectorImpl& plans) { - llvm::stable_sort(plans, [](const BatchRunSendPlan& lhs, const BatchRunSendPlan& rhs) { - std::optional lhsKey = getBestPeerExchangeOrderKey(lhs.messages); - std::optional rhsKey = getBestPeerExchangeOrderKey(rhs.messages); - if (!lhsKey || !rhsKey) - return lhsKey.has_value() && !rhsKey.has_value(); - return peerExchangeOrderLess(*lhsKey, *rhsKey); - }); -} +FailureOr emitPlannedReceive(MaterializerState& state, + CommunicationExchangeId id, + MaterializedClass& targetClass, + Location loc) { + if (id.value >= state.communicationExchangeStates.size()) + return targetClass.op->emitError("planned receive references an invalid exchange id"), failure(); + const CommunicationExchange& exchange = state.communicationPlan.getExchange(id); + CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[id.value]; + if (exchange.targetClass != targetClass.id) + return targetClass.op->emitError("planned receive target class mismatch"), failure(); + if (emission.receiveEmitted) + return emission.receiveValue; + if (!emission.sendEmitted) + return failure(); -struct ScalarSourceFanoutPlan { - SmallVector receivePlans; - std::optional ordinaryMessages; - SmallVector projectedSendGroups; -}; + FailureOr messages = buildMessageVectorForExchangeIds(state, + ArrayRef(id), + exchange.kind, + exchange.sourceClass, + exchange.targetClass, + Type(), + targetClass.op); + if (failed(messages)) + return failure(); -bool hasSameProjectedSendCompatibility(const ProjectedTransferDescriptor& lhs, const ProjectedTransferDescriptor& rhs) { - return lhs.layout.fragmentType == rhs.layout.fragmentType && lhs.layout.fragmentShape == rhs.layout.fragmentShape - && lhs.layout.fragmentsPerLogicalSlot == rhs.layout.fragmentsPerLogicalSlot - && lhs.layout.payloadFragmentCount == rhs.layout.payloadFragmentCount - && lhs.layout.loopLowerBounds == rhs.layout.loopLowerBounds && lhs.layout.loopSteps == rhs.layout.loopSteps - && lhs.layout.loopTripCounts == rhs.layout.loopTripCounts && lhs.payloadType == rhs.payloadType; + Value received = appendReceive(state, targetClass, exchange.payloadType, *messages, loc); + if (!received) + return failure(); + emission.receiveEmitted = true; + emission.receiveValue = received; + if (Operation* receiveDef = received.getDefiningOp()) + emission.receiveAnchor = getTopLevelMaterializedClassBodyOp(receiveDef, targetClass); + return received; } SmallVector collectDestinationClassesForKeys(MaterializerState& state, ArrayRef keys) { @@ -4611,562 +4474,32 @@ SmallVector collectDestinationClassesForKeys(MaterializerState& stat return destinations; } -FailureOr buildScalarSourceFanoutPlan(MaterializerState& state, - MaterializedClass& sourceClass, - ArrayRef keys, - ArrayRef destinationClasses, - Value payload) { - assert(!sourceClass.isBatch && "scalar-source send planning expects a scalar source class"); - - auto sourceCpu = getCheckedCoreId(sourceClass.op, sourceClass.cpus.front(), "scalar source core id"); - if (failed(sourceCpu)) - return failure(); - - ScalarSourceFanoutPlan fanoutPlan; - fanoutPlan.receivePlans.reserve(destinationClasses.size()); - - const auto getProjectedDescriptor = - [&](ClassId destinationClass) -> FailureOr> { - MaterializedClass& targetClass = state.classes[destinationClass]; - if (!targetClass.isBatch) { - bool hasAnyProjectedDescriptor = llvm::any_of(keys, [&](ProducerKey key) { - auto producerIt = state.projectedTransfers.find(key); - return producerIt != state.projectedTransfers.end() && producerIt->second.count(destinationClass) != 0; - }); - - std::optional descriptor = collectScalarTargetProjectedDescriptor( - state, targetClass, keys, /*requirePackedRunOffsetCountMatch=*/keys.size() > 1); - if (hasAnyProjectedDescriptor && !descriptor) - return targetClass.op->emitError("incomplete scalar projected transfer descriptor for local run"); - return descriptor; - } - - if (keys.size() != 1) - return std::optional {}; - - auto producerIt = state.projectedTransfers.find(keys.front()); - if (producerIt == state.projectedTransfers.end()) - return std::optional {}; - - auto descriptorIt = producerIt->second.find(destinationClass); - if (descriptorIt == producerIt->second.end()) - return std::optional {}; - - const ProjectedTransferDescriptor& descriptor = descriptorIt->second; - if (failed(verifyProjectedTransferDescriptor(targetClass.op, descriptor))) - return failure(); - if (descriptor.fragmentOffsets.size() - != targetClass.cpus.size() * static_cast(descriptor.layout.payloadFragmentCount)) - return targetClass.op->emitError("inconsistent batch projected transfer descriptor"); - - return std::optional {descriptor}; +SmallVector collectProjectedProducerKeysForOutput(MaterializerState& state, Value originalOutput) { + SmallVector keys; + auto appendKey = [&](ProducerKey key) { + if (!key.instance.op || key.resultIndex >= key.instance.op->getNumResults()) + return; + if (key.instance.op->getResult(key.resultIndex) != originalOutput) + return; + if (!llvm::is_contained(keys, key)) + keys.push_back(key); }; - - for (ClassId destinationClass : destinationClasses) { - if (destinationClass == sourceClass.id) - continue; - - MaterializedClass& targetClass = state.classes[destinationClass]; - - ScalarSourceReceivePlan receivePlan; - receivePlan.targetClass = destinationClass; - receivePlan.receiveType = payload.getType(); - - auto appendMessage = [&](CpuId targetCpu) -> LogicalResult { - auto checkedTargetCpu = getCheckedCoreId(targetClass.op, targetCpu, "scalar target core id"); - if (failed(checkedTargetCpu)) - return failure(); - int64_t channelId = state.nextChannelId++; - - receivePlan.messages.append(channelId, *sourceCpu, *checkedTargetCpu); - return success(); - }; - - if (!targetClass.isBatch) { - if (failed(appendMessage(targetClass.cpus.front()))) - return failure(); - } - else { - for (CpuId targetCpu : targetClass.cpus) - if (failed(appendMessage(targetCpu))) - return failure(); - } - - FailureOr> descriptor = getProjectedDescriptor(destinationClass); - if (failed(descriptor)) - return failure(); - - if (*descriptor) { - const ProjectedTransferDescriptor& projectedDescriptor = **descriptor; - FailureOr projectedCommunicationType = getProjectedCommunicationType(sourceClass.op, projectedDescriptor); - if (failed(projectedCommunicationType)) - return failure(); - - receivePlan.receiveType = *projectedCommunicationType; - receivePlan.projectedExtractOp = projectedDescriptor.extractOp; - receivePlan.projectedLayout = projectedDescriptor.layout; - - auto groupIt = llvm::find_if(fanoutPlan.projectedSendGroups, [&](const ProjectedScalarSendGroup& group) { - return hasSameProjectedSendCompatibility(group.descriptor, projectedDescriptor); - }); - if (groupIt == fanoutPlan.projectedSendGroups.end()) { - ProjectedScalarSendGroup group; - group.descriptor.inputKey = projectedDescriptor.inputKey; - group.descriptor.extractOp = projectedDescriptor.extractOp; - group.descriptor.layout = projectedDescriptor.layout; - group.descriptor.payloadType = projectedDescriptor.payloadType; - fanoutPlan.projectedSendGroups.push_back(std::move(group)); - groupIt = std::prev(fanoutPlan.projectedSendGroups.end()); - } - - groupIt->messages.append( - receivePlan.messages.channelIds, receivePlan.messages.sourceCoreIds, receivePlan.messages.targetCoreIds); - llvm::append_range(groupIt->descriptor.fragmentOffsets, projectedDescriptor.fragmentOffsets); - } - else { - if (!fanoutPlan.ordinaryMessages) - fanoutPlan.ordinaryMessages = MessageVector {}; - fanoutPlan.ordinaryMessages->append( - 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) { - if (failed(finalizeProjectedTransferDescriptor(sourceClass.op, group.descriptor))) - return failure(); - if (failed(verifyProjectedSendDescriptor(sourceClass.op, group.descriptor, group.messages))) - return failure(); - } - - return fanoutPlan; + for (const auto& producerEntry : state.projectedTransfers) + appendKey(producerEntry.first); + for (const auto& extractEntry : state.projectedInputTransferPlans) + for (const auto& classEntry : extractEntry.second) + for (const ProjectedInputTransferFragment& fragment : classEntry.second.fragments) + appendKey(fragment.producer); + return keys; } -LogicalResult collectPlannedScalarPeerReceivesForSource(MaterializerState& state, - MaterializedClass& sourceClass, - ArrayRef keys, - Value payload) { - SmallVector destinationClasses = collectDestinationClassesForKeys(state, keys); - int64_t savedNextChannelId = state.nextChannelId; - auto fanoutPlan = buildScalarSourceFanoutPlan(state, sourceClass, keys, destinationClasses, payload); - state.nextChannelId = savedNextChannelId; - if (failed(fanoutPlan)) - return failure(); - - for (const ScalarSourceReceivePlan& plan : fanoutPlan->receivePlans) { - for (size_t index = 0; index < plan.messages.size(); ++index) { - ScalarPeerEdgeKey receiveKey { - .sourceCore = plan.messages.sourceCoreIds[index], - .targetCore = plan.messages.targetCoreIds[index], - .payloadType = plan.receiveType, - }; - bool alreadyPlanned = llvm::any_of(state.plannedScalarPeerReceives, [&](const ScalarPeerEdgeKey& record) { - return record.sourceCore == receiveKey.sourceCore && record.targetCore == receiveKey.targetCore - && record.payloadType == receiveKey.payloadType; - }); - if (!alreadyPlanned) - state.plannedScalarPeerReceives.push_back(receiveKey); - } - } - - return success(); -} - -LogicalResult emitScalarSourceFanoutSends(MaterializerState& state, - MaterializedClass& sourceClass, - Value payload, - const ScalarSourceFanoutPlan& plan, - Location loc) { - if (plan.ordinaryMessages) { - MessageVector orderedMessages = orderedMessageVector(*plan.ordinaryMessages); - if (failed(appendSend(state, sourceClass, payload, orderedMessages, loc))) - return failure(); - } - - for (const ProjectedScalarSendGroup& group : plan.projectedSendGroups) { - ProjectedScalarSendGroup orderedGroup = orderedProjectedScalarSendGroup(group); - if (failed(finalizeSplitProjectedScalarSendGroup(sourceClass.op, orderedGroup))) - return failure(); - if (failed(appendProjectedScalarSendLoop( - state, sourceClass, payload, orderedGroup.descriptor, orderedGroup.messages, loc))) - return failure(); - } - - return success(); -} - -ProjectedScalarSendGroup makeEmptyProjectedScalarSendGroup(const ProjectedScalarSendGroup& group) { - ProjectedScalarSendGroup splitGroup; - splitGroup.descriptor = group.descriptor; - splitGroup.descriptor.fragmentOffsets.clear(); - splitGroup.descriptor.fragmentOffsetsByDim.clear(); - return splitGroup; -} - -void appendProjectedScalarSendMessage(ProjectedScalarSendGroup& destination, - const ProjectedScalarSendGroup& source, - size_t index) { - const unsigned payloadFragmentCount = source.descriptor.layout.payloadFragmentCount; - destination.messages.append(source.messages.channelIds[index], - source.messages.sourceCoreIds[index], - source.messages.targetCoreIds[index]); - ArrayRef> offsets( - source.descriptor.fragmentOffsets.data() + index * payloadFragmentCount, payloadFragmentCount); - llvm::append_range(destination.descriptor.fragmentOffsets, offsets); -} - -ProjectedScalarSendGroup orderedProjectedScalarSendGroup(const ProjectedScalarSendGroup& group) { - SmallVector indices; - indices.reserve(group.messages.size()); - for (size_t index = 0; index < group.messages.size(); ++index) - indices.push_back(index); - sortMessageIndicesByPeerExchangeOrder(group.messages, indices); - - ProjectedScalarSendGroup ordered = makeEmptyProjectedScalarSendGroup(group); - for (size_t index : indices) - appendProjectedScalarSendMessage(ordered, group, index); - return ordered; -} - -LogicalResult finalizeSplitProjectedScalarSendGroup(Operation* anchor, ProjectedScalarSendGroup& group) { - if (group.messages.empty()) - return success(); - if (failed(finalizeProjectedTransferDescriptor(anchor, group.descriptor))) - return failure(); - if (failed(verifyProjectedSendDescriptor(anchor, group.descriptor, group.messages))) - return failure(); - return success(); -} - -ScalarPeerEdgeKey makeScalarPeerEdgeKey(int64_t sourceCore, int64_t targetCore, Type payloadType) { - return ScalarPeerEdgeKey {.sourceCore = sourceCore, .targetCore = targetCore, .payloadType = payloadType}; -} - -ScalarPeerReceiveKey makeScalarPeerReceiveKey(int64_t sourceCore, - int64_t targetCore, - Type payloadType, - std::optional channelId) { - return ScalarPeerReceiveKey { - .sourceCore = sourceCore, - .targetCore = targetCore, - .payloadType = payloadType, - .channelId = channelId, - }; -} - -ScalarPeerEdgeKey makeRequiredPeerReceiveForLowerToHigherSend(int64_t sourceCore, int64_t targetCore, Type payloadType) { - return makeScalarPeerEdgeKey(targetCore, sourceCore, payloadType); -} - -bool hasMaterializedScalarPeerReceive(const MaterializerState& state, const ScalarPeerEdgeKey& key) { - return llvm::any_of(state.materializedScalarPeerReceives, [&](const ScalarPeerEdgeKey& record) { - return record.sourceCore == key.sourceCore && record.targetCore == key.targetCore && record.payloadType == key.payloadType; - }); -} - -bool hasPlannedScalarPeerReceive(const MaterializerState& state, const ScalarPeerEdgeKey& key) { - return llvm::any_of(state.plannedScalarPeerReceives, [&](const ScalarPeerEdgeKey& record) { - return record.sourceCore == key.sourceCore && record.targetCore == key.targetCore && record.payloadType == key.payloadType; - }); -} - -bool receiveKeyMatchesPendingWait(const ScalarPeerReceiveKey& receiveKey, const ScalarPeerEdgeKey& pendingWait) { - return receiveKey.sourceCore == pendingWait.sourceCore && receiveKey.targetCore == pendingWait.targetCore - && receiveKey.payloadType == pendingWait.payloadType; -} - -bool hasPendingScalarSendWaitingForReceive(const MaterializerState& state, const ScalarPeerReceiveKey& receiveKey) { - return llvm::any_of(state.pendingScalarSends, [&](const PendingScalarSend& pending) { - return receiveKeyMatchesPendingWait(receiveKey, pending.waitForReceive); - }) || llvm::any_of(state.pendingProjectedScalarSends, [&](const PendingProjectedScalarSend& pending) { - return receiveKeyMatchesPendingWait(receiveKey, pending.waitForReceive); - }); -} - -ReceiveMessagePartition partitionReceivesByPendingSendUnlocks(const MaterializerState& state, - const MaterializedClass& targetClass, - Type payloadType, - const MessageVector& messages, - bool preserveMessageOrder) { - (void) targetClass; +ReceiveMessagePartition partitionReceiveMessagesInPlanOrder(const MessageVector& messages) { ReceiveMessagePartition partition; - for (size_t index = 0; index < messages.size(); ++index) { - ScalarPeerReceiveKey receiveKey = makeScalarPeerReceiveKey(messages.sourceCoreIds[index], - messages.targetCoreIds[index], - payloadType, - messages.channelIds[index]); - bool isCritical = hasPendingScalarSendWaitingForReceive(state, receiveKey); - (isCritical ? partition.criticalStaticIndices : partition.remainingIndices).push_back(index); - } - - if (!preserveMessageOrder) { - sortMessageIndicesByPeerExchangeOrder(messages, partition.criticalStaticIndices); - sortMessageIndicesByPeerExchangeOrder(messages, partition.remainingIndices); - } + for (size_t index = 0; index < messages.size(); ++index) + partition.remainingIndices.push_back(index); return partition; } -FailureOr findTopLevelDefAnchorInClass(MaterializedClass& klass, Value value) { - if (auto blockArg = dyn_cast(value)) { - if (blockArg.getOwner() == klass.body) - return static_cast(nullptr); - klass.op->emitError("pending scalar send payload block argument is not local to class"); - return failure(); - } - - Operation* definingOp = value.getDefiningOp(); - if (!definingOp) - return klass.op->emitError("pending scalar send payload has no defining op"), failure(); - - Operation* current = definingOp; - while (current) { - if (current->getBlock() == klass.body) - return current; - current = current->getParentOp(); - } - - klass.op->emitError("pending scalar send payload is defined outside the source materialized class"); - return failure(); -} - -Operation* latestAnchorInBlock(Operation* a, Operation* b) { - if (!a) - return b; - if (!b) - return a; - if (a->getBlock() != b->getBlock()) - return nullptr; - return a->isBeforeInBlock(b) ? b : a; -} - -FailureOr resolvePendingSendInsertionAnchor(MaterializedClass& sourceClass, - Operation* receiveAnchor, - Operation* payloadAnchor) { - if (!receiveAnchor) - return sourceClass.op->emitError("missing receive anchor while flushing pending scalar send"), failure(); - if (!payloadAnchor) - return receiveAnchor; - if (receiveAnchor->getBlock() == payloadAnchor->getBlock()) - return latestAnchorInBlock(receiveAnchor, payloadAnchor); - - Operation* current = receiveAnchor; - while (current) { - Operation* parent = current->getParentOp(); - if (!parent) - break; - if (payloadAnchor->getBlock() == parent->getBlock()) { - if (payloadAnchor->isBeforeInBlock(current)) - return receiveAnchor; - return sourceClass.op->emitError("pending scalar send payload does not dominate nested receive flush point"), - failure(); - } - current = parent; - } - - return sourceClass.op->emitError("pending scalar send payload anchor is incompatible with nested receive flush point"), - failure(); -} - -LogicalResult enqueuePendingScalarSend(MaterializerState& state, - ClassId sourceClass, - int64_t sourceCore, - int64_t targetCore, - Type payloadType, - Value payload, - const MessageVector& messages, - Location loc) { - MaterializedClass& sourceMaterializedClass = state.classes[sourceClass]; - if (failed(messages.verify(sourceMaterializedClass.op))) - return failure(); - FailureOr payloadAnchor = findTopLevelDefAnchorInClass(sourceMaterializedClass, payload); - if (failed(payloadAnchor)) - return failure(); - state.pendingScalarSends.push_back(PendingScalarSend { - .sourceClass = sourceClass, - .sourceCore = sourceCore, - .targetCore = targetCore, - .payloadType = payloadType, - .waitForReceive = makeRequiredPeerReceiveForLowerToHigherSend(sourceCore, targetCore, payloadType), - .payloadAnchor = *payloadAnchor, - .payload = payload, - .messages = messages, - .loc = loc, - }); - return success(); -} - -LogicalResult enqueuePendingProjectedScalarSend(MaterializerState& state, - ClassId sourceClass, - int64_t sourceCore, - int64_t targetCore, - Type payloadType, - Value payload, - const ProjectedScalarSendGroup& group, - Location loc) { - MaterializedClass& sourceMaterializedClass = state.classes[sourceClass]; - if (failed(verifyProjectedSendDescriptor(sourceMaterializedClass.op, group.descriptor, group.messages))) - return failure(); - FailureOr payloadAnchor = findTopLevelDefAnchorInClass(sourceMaterializedClass, payload); - if (failed(payloadAnchor)) - return failure(); - state.pendingProjectedScalarSends.push_back(PendingProjectedScalarSend { - .sourceClass = sourceClass, - .sourceCore = sourceCore, - .targetCore = targetCore, - .payloadType = payloadType, - .waitForReceive = makeRequiredPeerReceiveForLowerToHigherSend(sourceCore, targetCore, payloadType), - .payloadAnchor = *payloadAnchor, - .payload = payload, - .messages = group.messages, - .descriptor = group.descriptor, - .loc = loc, - }); - return success(); -} - -LogicalResult flushPendingScalarSendsForReceive( - MaterializerState& state, const ScalarPeerEdgeKey& receiveKey, Operation* anchor) { - for (size_t index = 0; index < state.pendingProjectedScalarSends.size();) { - const PendingProjectedScalarSend& pending = state.pendingProjectedScalarSends[index]; - if (pending.waitForReceive.sourceCore != receiveKey.sourceCore - || pending.waitForReceive.targetCore != receiveKey.targetCore - || pending.waitForReceive.payloadType != receiveKey.payloadType) { - ++index; - continue; - } - - MaterializedClass& sourceClass = state.classes[pending.sourceClass]; - FailureOr sendAnchor = resolvePendingSendInsertionAnchor(sourceClass, anchor, pending.payloadAnchor); - if (failed(sendAnchor)) - return failure(); - OpBuilder::InsertionGuard guard(state.rewriter); - state.rewriter.setInsertionPointAfter(*sendAnchor); - if (failed(appendProjectedScalarSendLoopAtCurrentInsertionPoint( - state, sourceClass, pending.payload, pending.descriptor, pending.messages, pending.loc))) - return failure(); - state.pendingProjectedScalarSends.erase(state.pendingProjectedScalarSends.begin() + index); - } - - for (size_t index = 0; index < state.pendingScalarSends.size();) { - const PendingScalarSend& pending = state.pendingScalarSends[index]; - if (pending.waitForReceive.sourceCore != receiveKey.sourceCore - || pending.waitForReceive.targetCore != receiveKey.targetCore - || pending.waitForReceive.payloadType != receiveKey.payloadType) { - ++index; - continue; - } - - MaterializedClass& sourceClass = state.classes[pending.sourceClass]; - FailureOr sendAnchor = resolvePendingSendInsertionAnchor(sourceClass, anchor, pending.payloadAnchor); - if (failed(sendAnchor)) - return failure(); - OpBuilder::InsertionGuard guard(state.rewriter); - state.rewriter.setInsertionPointAfter(*sendAnchor); - if (pending.messages.size() == 1) - appendScalarSendAtCurrentInsertionPoint(state, - sourceClass, - pending.payload, - pending.messages.channelIds.front(), - pending.messages.sourceCoreIds.front(), - pending.messages.targetCoreIds.front(), - pending.loc); - else if (failed(appendScalarSendLoopAtCurrentInsertionPoint( - state, sourceClass, pending.payload, pending.messages, pending.loc))) - return failure(); - state.pendingScalarSends.erase(state.pendingScalarSends.begin() + index); - } - - return success(); -} - -FailureOr createScalarPeerReceiveAndFlush(MaterializerState& state, - MaterializedClass& targetClass, - Type payloadType, - Value channelId, - Value sourceCoreId, - Value targetCoreId, - std::optional staticKey, - Location loc) { - OpBuilder::InsertionGuard guard(state.rewriter); - - Value received = - SpatChannelReceiveOp::create(state.rewriter, loc, payloadType, channelId, sourceCoreId, targetCoreId).getOutput(); - if (targetClass.isBatch || !staticKey || staticKey->sourceCore <= staticKey->targetCore) - return received; - - ScalarPeerEdgeKey receiveKey = - makeScalarPeerEdgeKey(staticKey->sourceCore, staticKey->targetCore, staticKey->payloadType); - if (!hasMaterializedScalarPeerReceive(state, receiveKey)) - state.materializedScalarPeerReceives.push_back(receiveKey); - if (failed(flushPendingScalarSendsForReceive(state, receiveKey, received.getDefiningOp()))) - return failure(); - return received; -} - -FailureOr appendReceiveAndFlushPendingScalarSends( - MaterializerState& state, MaterializedClass& targetClass, Type type, const MessageVector& messages, Location loc) { - assert(succeeded(messages.verify(targetClass.op)) && "message metadata is inconsistent"); - assert(!messages.empty() && "expected at least one receive"); - - state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - if (targetClass.isBatch) - return appendReceive(state, targetClass, type, messages, loc); - if (messages.size() != 1) - return targetClass.op->emitError("scalar receive flush expects exactly one message") - << " messageCount=" << messages.size(), - failure(); - - return createScalarPeerReceiveAndFlush(state, - targetClass, - type, - getOrCreateIndexConstant(state.constantFolder, targetClass.op, messages.channelIds.front()), - getOrCreateIndexConstant( - state.constantFolder, targetClass.op, messages.sourceCoreIds.front()), - getOrCreateIndexConstant( - state.constantFolder, targetClass.op, messages.targetCoreIds.front()), - makeScalarPeerReceiveKey(messages.sourceCoreIds.front(), - messages.targetCoreIds.front(), - type, - messages.channelIds.front()), - loc); -} - - - LogicalResult appendConditionalBatchProjectedInputSendsAtCurrentInsertionPoint(MaterializerState& state, MaterializedClass& sourceClass, Value payload, @@ -5228,7 +4561,7 @@ LogicalResult appendConditionalBatchProjectedInputSendsAtCurrentInsertionPoint(M if (failed(sendLoop)) return failure(); if (sendLoop->loop) - tagProjectedInputTransfer(sendLoop->loop.getOperation(), getBestPeerExchangeOrderKey(messages)); + tagProjectedInputTransfer(sendLoop->loop.getOperation()); return success(); } @@ -5249,10 +4582,7 @@ LogicalResult appendScalarProjectedInputSendAtCurrentInsertionPoint(Materializer 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))); + tagProjectedInputTransfer(send.getOperation()); return success(); } @@ -5273,32 +4603,31 @@ LogicalResult emitProjectedInputSendsAtCurrentInsertionPoint(MaterializerState& return success(); } -static std::optional getProjectedExchangeId(MaterializerState& state, - const ProjectedInputTransferFragment& fragment, - Operation* anchor) { - std::optional exchangeId = state.projectedInputCommunicationPlan.getExchangeId(fragment); - if (!exchangeId && anchor) +static std::optional getProjectedExchangeId(MaterializerState& state, + const ProjectedInputTransferFragment& fragment, + Operation* anchor) { + ArrayRef ids = state.communicationPlan.getProjectedInputExchanges(fragment); + if (ids.size() > 1 && anchor) + anchor->emitError("projected remote fragment has multiple communication exchanges") + << " sourceCore=" << fragment.sourceCoreId << " targetCore=" << fragment.targetCoreId; + if (ids.empty() && anchor) anchor->emitError("projected remote fragment is missing from communication plan") << " sourceCore=" << fragment.sourceCoreId << " targetCore=" << fragment.targetCoreId << " channelId=" << fragment.channelId; - return exchangeId; -} - -static void recordProjectedInputDeferral(MaterializerState& state, size_t exchangeId) { - if (!llvm::is_contained(state.deferredProjectedExchanges, exchangeId)) - state.deferredProjectedExchanges.push_back(exchangeId); + if (ids.size() != 1) + return std::nullopt; + return ids.front(); } static bool isProjectedRemoteFragmentReadyForReceive(MaterializerState& state, const ProjectedInputTransferFragment& fragment, Operation* anchor) { - std::optional exchangeId = getProjectedExchangeId(state, fragment, 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); + CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[exchangeId->value]; + if (!emission.sendEmitted) { + recordDeferredProducerValue(state, fragment.producer); return false; } return true; @@ -5316,15 +4645,13 @@ static LogicalResult preflightProjectedProducerReceives(MaterializerState& state if (!(fragment.producer == producer) || fragment.sourceCoreId == fragment.targetCoreId) continue; - std::optional exchangeId = + 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); + const CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[exchangeId->value]; + (void)emission; } } return success(); @@ -5348,6 +4675,75 @@ static LogicalResult preflightProjectedWholeBatchProducerReceives(MaterializerSt return success(); } +static bool isRemoteProjectedInputPlanReadyForReceive(MaterializerState& state, + const ProjectedInputTransferPlan& plan) { + bool hasRemoteFragment = false; + for (const ProjectedInputTransferFragment& fragment : plan.fragments) { + if (fragment.sourceCoreId == fragment.targetCoreId) + return false; + + hasRemoteFragment = true; + ArrayRef exchangeIds = state.communicationPlan.getProjectedInputExchanges(fragment); + if (exchangeIds.size() > 1) + return false; + std::optional exchangeId = + exchangeIds.empty() ? std::nullopt : std::optional(exchangeIds.front()); + if (!exchangeId) + return false; + if (!state.communicationExchangeStates[exchangeId->value].sendEmitted) + return false; + } + return hasRemoteFragment; +} + +static bool hasBoundReciprocalDifferentPayloadSendPending(MaterializerState& state, + const ExchangeDescriptor& receiveExchange) { + for (const ExchangeDescriptor& candidate : state.communicationPlan.getExchanges()) { + if (candidate.kind != CommunicationExchangeKind::ProjectedInput) + continue; + if (candidate.sourceCore != receiveExchange.targetCore || candidate.targetCore != receiveExchange.sourceCore) + continue; + if (candidate.payloadType == receiveExchange.payloadType) + continue; + + const CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[candidate.id.value]; + if (emission.producerBound && !emission.sendEmitted) + return true; + } + return false; +} + +static LogicalResult materializeReadyProjectedInputReceiveForExchange(MaterializerState& state, + const ExchangeDescriptor& exchange, + Location loc) { + MaterializedClass& targetClass = state.classes[exchange.targetClass]; + for (auto& extractEntry : state.projectedInputTransferPlans) { + auto classIt = extractEntry.second.find(targetClass.id); + if (classIt == extractEntry.second.end()) + continue; + + bool containsExchange = llvm::any_of(classIt->second.fragments, [&](const ProjectedInputTransferFragment& fragment) { + return &fragment == exchange.fragment; + }); + if (!containsExchange) + continue; + + auto replacementIt = state.projectedExtractReplacements.find(extractEntry.first); + if (replacementIt != state.projectedExtractReplacements.end() + && replacementIt->second.find(targetClass.id) != replacementIt->second.end()) + return success(); + + if (!isRemoteProjectedInputPlanReadyForReceive(state, classIt->second)) + return success(); + if (hasBoundReciprocalDifferentPayloadSendPending(state, exchange)) + return success(); + + Operation* insertionAnchor = getProjectedReceiveInsertionAnchor(targetClass); + return materializeProjectedInputReceivesForClass(state, targetClass, extractEntry.first, loc, insertionAnchor); + } + return success(); +} + LogicalResult emitPlannedProjectedInputSendsForKeys(MaterializerState& state, MaterializedClass& sourceClass, ArrayRef keys, @@ -5360,8 +4756,6 @@ LogicalResult emitPlannedProjectedInputSendsForKeys(MaterializerState& state, for (ProducerKey key : keys) keySet.insert(key); - SmallVector lowToHighFragments; - SmallVector highToLowFragments; for (auto& extractEntry : state.projectedInputTransferPlans) { for (auto& classEntry : extractEntry.second) { ProjectedInputTransferPlan& plan = classEntry.second; @@ -5369,84 +4763,102 @@ LogicalResult emitPlannedProjectedInputSendsForKeys(MaterializerState& state, if (!keySet.contains(fragment.producer)) 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 (fragment.sourceCoreId == fragment.targetCoreId) { + 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; - } - - 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; - } + std::optional exchangeId = getProjectedExchangeId(state, fragment, sourceClass.op); + if (!exchangeId) + return failure(); + if (state.communicationPlan.getExchange(*exchangeId).sourceClass != sourceClass.id) + continue; + if (failed(bindCommunicationPayload(state, *exchangeId, sourceClass, payload, sourceClass.op, loc))) + return failure(); } } } - if (failed(emitProjectedInputSendsAtCurrentInsertionPoint(state, sourceClass, payload, lowToHighFragments, loc))) + FailureOr emittedBound = emitReadyCommunicationEventsForClass(state, sourceClass, loc); + if (failed(emittedBound)) return failure(); - 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"); +static LogicalResult emitProjectedInputSendForExchange(MaterializerState& state, + const ExchangeDescriptor& exchange, + CommunicationExchangeEmissionState& emission, + Location loc) { + if (!exchange.fragment) + return state.func.emitError("projected communication exchange is missing its fragment"); + if (!emission.producerBound || emission.sendEmitted) + return success(); + if (!emission.loc) + emission.loc = loc; - MaterializedClass& sourceClass = state.classes[exchange.sourceClass]; - OpBuilder::InsertionGuard guard(state.rewriter); + MaterializedClass& sourceClass = state.classes[exchange.sourceClass]; + OpBuilder::InsertionGuard guard(state.rewriter); + if (emission.producerAnchor && emission.producerAnchor->getBlock() == sourceClass.body) + state.rewriter.setInsertionPointAfter(emission.producerAnchor); + else 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; + + ProjectedInputTransferFragment* fragment = const_cast(exchange.fragment); + SmallVector fragments {fragment}; + if (failed(emitProjectedInputSendsAtCurrentInsertionPoint( + state, sourceClass, emission.producerPayload, fragments, *emission.loc))) + return failure(); + emission.sendEmitted = true; + emission.sendAnchor = fragment->sendEmitted && emission.producerAnchor ? emission.producerAnchor : nullptr; + fragment->sendEmitted = true; + return success(); +} + +static FailureOr emitReadyCommunicationEventsForClass(MaterializerState& state, + MaterializedClass& materializedClass, + Location loc) { + bool emittedAny = false; + for (const CommunicationEvent& event : state.communicationPlan.getEventsForClass(materializedClass.id)) { + const ExchangeDescriptor& exchange = state.communicationPlan.getExchange(event.exchangeId); + CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[exchange.id]; + switch (event.kind) { + case CommunicationEventKind::Send: + if (!emission.producerBound || emission.sendEmitted) + break; + if (exchange.kind == CommunicationExchangeKind::ProjectedInput) { + bool wasEmitted = emission.sendEmitted; + if (failed(emitProjectedInputSendForExchange(state, exchange, emission, loc))) + return failure(); + emittedAny |= !wasEmitted && emission.sendEmitted; + } + else { + if (failed(emitPlannedSend(state, event.exchangeId, materializedClass, loc))) + return failure(); + emittedAny = true; + } + break; + case CommunicationEventKind::Receive: + if (!emission.sendEmitted || emission.receiveEmitted) + break; + if (exchange.kind == CommunicationExchangeKind::ProjectedInput) { + bool wasEmitted = emission.receiveEmitted; + if (failed(materializeReadyProjectedInputReceiveForExchange(state, exchange, loc))) + return failure(); + emittedAny |= !wasEmitted && emission.receiveEmitted; + } + else if (exchange.kind == CommunicationExchangeKind::OrdinaryValue + || exchange.kind == CommunicationExchangeKind::HostForward) { + if (materializedClass.isBatch) + break; + FailureOr received = emitPlannedReceive(state, event.exchangeId, materializedClass, loc); + if (failed(received)) + return failure(); + if (exchange.kind == CommunicationExchangeKind::OrdinaryValue) + state.availableValues.record(exchange.producer, exchange.targetClass, *received); + emittedAny = true; + } + break; + } } return emittedAny; } @@ -5516,73 +4928,42 @@ static bool isLocalProjectedInputFragment(const ProjectedInputTransferFragment& enum class ProjectedInputExchangeDirection { Local, - LowerToHigher, - HigherToLower + Remote }; 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) { +static void tagProjectedInputTransfer(Operation* op) { 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 Operation* getProjectedReceiveInsertionAnchor(MaterializedClass& targetClass) { + Operation* lastProjectedTransfer = nullptr; + for (Operation& op : targetClass.body->without_terminator()) + if (op.hasAttr(projectedInputTransferAttr)) + lastProjectedTransfer = &op; + + if (!lastProjectedTransfer) + return targetClass.body->getTerminator(); + + Operation* next = lastProjectedTransfer->getNextNode(); + return next ? next : targetClass.body->getTerminator(); } static FailureOr -getProjectedInputExchangeDirectionKind(MaterializerState& state, +getProjectedInputExchangeDirectionKind(MaterializerState&, 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: + (void)targetClass; + if (fragment.sourceCoreId == fragment.targetCoreId) return ProjectedInputExchangeDirection::Local; - 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"); + return ProjectedInputExchangeDirection::Remote; } static void tagProjectedInputReceive(Value value, ProjectedInputExchangeDirection direction) { - StringRef attrValue = getProjectedInputDirectionAttrValue(direction); - if (attrValue.empty()) + if (direction == ProjectedInputExchangeDirection::Local) return; auto receive = value.getDefiningOp(); @@ -5590,8 +4971,6 @@ static void tagProjectedInputReceive(Value value, ProjectedInputExchangeDirectio return; tagProjectedInputTransfer(receive.getOperation()); - MLIRContext* context = receive->getContext(); - receive->setAttr(projectedInputDirectionAttr, StringAttr::get(context, attrValue)); } static FailureOr getUniformProjectedInputSlotDirection( @@ -5620,6 +4999,10 @@ static FailureOr resolveUniformLocalBatchProjectedInputSlot(MaterializerS std::optional slotValue; for (const ProjectedInputTransferFragment* fragment : fragments) { std::optional local = state.availableValues.lookup(state, fragment->producer, targetClass.id); + if (!local && !isProducerValueMaterialized(state, fragment->producer)) { + recordDeferredProducerValue(state, fragment->producer); + return failure(); + } if (!local) return targetClass.op->emitError("projected input batch transfer could not resolve local fragment") << " sourceCore=" << fragment->sourceCoreId << " targetCore=" << fragment->targetCoreId @@ -5638,29 +5021,6 @@ static FailureOr resolveUniformLocalBatchProjectedInputSlot(MaterializerS return *slotValue; } -static FailureOr materializeRemoteBatchProjectedInputSlotReceive( - MaterializerState& state, - MaterializedClass& targetClass, - Type fragmentType, - ArrayRef fragments, - Value laneArg, - Location loc) { - MessageVector messages; - messages.channelIds.reserve(fragments.size()); - messages.sourceCoreIds.reserve(fragments.size()); - messages.targetCoreIds.reserve(fragments.size()); - for (const ProjectedInputTransferFragment* fragment : fragments) - messages.append(fragment->channelId, fragment->sourceCoreId, fragment->targetCoreId); - - return SpatChannelReceiveOp::create(state.rewriter, - loc, - fragmentType, - createIndexedChannelId(state, targetClass.op, messages, laneArg, loc), - createIndexedSourceCoreId(state, targetClass.op, messages, laneArg, loc), - createIndexedTargetCoreId(state, targetClass.op, messages, laneArg, loc)) - .getOutput(); -} - struct ProjectedInputPlanMaterialization { struct BatchSlotGroup { unsigned slot = 0; @@ -5675,8 +5035,7 @@ struct ProjectedInputPlanMaterialization { Value laneArg; SmallVector batchSlotGroups; SmallVector localScalarFragments; - SmallVector lowerToHigherScalarFragments; - SmallVector higherToLowerScalarFragments; + SmallVector remoteScalarFragments; }; static FailureOr @@ -5767,12 +5126,10 @@ prepareProjectedInputPlanMaterialization(MaterializerState& state, 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 + if (*direction == ProjectedInputExchangeDirection::Local) work.localScalarFragments.push_back(index); + else + work.remoteScalarFragments.push_back(index); } return work; @@ -5790,6 +5147,10 @@ static FailureOr materializeScalarProjectedInputFragment(MaterializerStat if (*direction == ProjectedInputExchangeDirection::Local) { std::optional local = state.availableValues.lookup(state, fragment.producer, targetClass.id); + if (!local && !isProducerValueMaterialized(state, fragment.producer)) { + recordDeferredProducerValue(state, fragment.producer); + return failure(); + } if (!local) return targetClass.op->emitError("projected input transfer could not resolve local fragment"), failure(); return *local; @@ -5798,23 +5159,13 @@ static FailureOr materializeScalarProjectedInputFragment(MaterializerStat 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); + std::optional exchangeId = getProjectedExchangeId(state, fragment, targetClass.op); + if (!exchangeId) + return failure(); + FailureOr received = emitPlannedReceive(state, *exchangeId, targetClass, loc); if (failed(received)) return failure(); tagProjectedInputReceive(*received, *direction); - if (std::optional exchangeId = state.projectedInputCommunicationPlan.getExchangeId(fragment)) - state.projectedExchangeStates[*exchangeId].receiveEmitted = true; return *received; } @@ -5846,16 +5197,44 @@ static FailureOr materializeBatchProjectedInputPayloadSlotGroupValue( if (!isProjectedRemoteFragmentReadyForReceive(state, *fragment, targetClass.op)) return failure(); + SmallVector exchangeIds; + exchangeIds.reserve(groupFragments.size()); + for (const ProjectedInputTransferFragment* fragment : groupFragments) { + std::optional exchangeId = getProjectedExchangeId(state, *fragment, targetClass.op); + if (!exchangeId) + return failure(); + exchangeIds.push_back(*exchangeId); + } + Value groupIndex = createIndexedIndexValue(state, targetClass.op, laneToGroupIndex, laneArg, loc); - FailureOr received = materializeRemoteBatchProjectedInputSlotReceive( - state, targetClass, plan.layout.fragmentType, groupFragments, groupIndex, loc); - if (failed(received)) + FailureOr messages = + buildMessageVectorForExchangeIds(state, + exchangeIds, + CommunicationExchangeKind::ProjectedInput, + state.communicationPlan.getExchange(exchangeIds.front()).sourceClass, + targetClass.id, + plan.layout.fragmentType, + targetClass.op); + if (failed(messages)) 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; + Value received = SpatChannelReceiveOp::create(state.rewriter, + loc, + plan.layout.fragmentType, + createIndexedChannelId(state, targetClass.op, *messages, groupIndex, loc), + createIndexedSourceCoreId(state, targetClass.op, *messages, groupIndex, loc), + createIndexedTargetCoreId(state, targetClass.op, *messages, groupIndex, loc)) + .getOutput(); + tagProjectedInputReceive(received, *direction); + Operation* receiveAnchor = nullptr; + if (Operation* receiveDef = received.getDefiningOp()) + receiveAnchor = getTopLevelMaterializedClassBodyOp(receiveDef, targetClass); + for (CommunicationExchangeId exchangeId : exchangeIds) { + CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[exchangeId.value]; + emission.receiveEmitted = true; + emission.receiveValue = received; + emission.receiveAnchor = receiveAnchor; + } + return received; } static FailureOr insertProjectedBatchSlotGroup(MaterializerState& state, @@ -5931,28 +5310,15 @@ static LogicalResult materializeProjectedInputScalarFragment(MaterializerState& return success(); } -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); + std::optional exchangeId = getProjectedExchangeId(state, fragment, targetClass.op); 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 failure(); } return success(); } @@ -5979,7 +5345,7 @@ static LogicalResult preflightProjectedInputsForInstance(MaterializerState& stat if (replacementIt != state.projectedExtractReplacements.end() && replacementIt->second.find(targetClass.id) != replacementIt->second.end()) continue; - if (extractEntry.first != requestedExtract && !hasRemoteProjectedInputFragment(classIt->second)) + if (extractEntry.first != requestedExtract) continue; if (failed(preflightProjectedInputPlanReceives(state, targetClass, classIt->second))) @@ -6004,7 +5370,7 @@ collectProjectedInputPlanMaterializationsForClass(MaterializerState& state, if (replacementIt != state.projectedExtractReplacements.end() && replacementIt->second.find(targetClass.id) != replacementIt->second.end()) continue; - if (extractEntry.first != requestedExtract && !hasRemoteProjectedInputFragment(classIt->second)) + if (extractEntry.first != requestedExtract) continue; FailureOr work = @@ -6049,20 +5415,10 @@ static LogicalResult materializeProjectedInputReceivesForClass(MaterializerState for (ProjectedInputPlanMaterialization& work : works) { for (const auto& group : work.batchSlotGroups) - if (group.direction == ProjectedInputExchangeDirection::LowerToHigher + if (group.direction == ProjectedInputExchangeDirection::Remote && 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(); - } - - 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) + for (size_t index : work.remoteScalarFragments) if (failed(materializeProjectedInputScalarFragment(state, targetClass, work, index, loc))) return failure(); } @@ -6078,27 +5434,37 @@ static LogicalResult materializeProjectedInputReceivesForClass(MaterializerState return success(); } -LogicalResult verifyNoPendingProjectedInputSends(MaterializerState& state) { - 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(); - } +LogicalResult verifyCommunicationPlanFullyMaterialized(MaterializerState& state) { + size_t malformedCount = 0; + SmallVector malformed; + for (const ExchangeDescriptor& exchange : state.communicationPlan.getExchanges()) { + const CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[exchange.id.value]; + bool bad = !emission.producerBound || !emission.sendEmitted || !emission.receiveEmitted || !emission.receiveValue; + if (emission.receiveValue && emission.receiveValue.getType() != exchange.payloadType) + bad = true; + if (!bad) + continue; + ++malformedCount; + if (malformed.size() < 8) + malformed.push_back(&exchange); } - return success(); + + if (malformedCount == 0) + return success(); + + InFlightDiagnostic diag = state.func.emitError("communication plan was not fully materialized"); + for (const ExchangeDescriptor* exchange : malformed) { + const CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[exchange->id.value]; + diag << " [exchange=" << exchange->id.value << " kind=" << static_cast(exchange->kind) + << " sourceClass=" << exchange->sourceClass << " sourceCore=" << exchange->sourceCore + << " targetClass=" << exchange->targetClass << " targetCore=" << exchange->targetCore + << " channelId=" << exchange->channelId << " producerBound=" << emission.producerBound + << " sendEmitted=" << emission.sendEmitted << " receiveEmitted=" << emission.receiveEmitted + << " receiveType=" << (emission.receiveValue ? emission.receiveValue.getType() : Type()) << "]"; + } + if (malformedCount > malformed.size()) + diag << " additionalIncomplete=" << (malformedCount - malformed.size()); + return failure(); } MessageVector filterMessageVector(const MessageVector& messages, ArrayRef indices) { @@ -6116,141 +5482,15 @@ SmallVector filterInt64Vector(ArrayRef values, ArrayRef summaries; - auto accumulate = [&](ClassId sourceClass, int64_t sourceCore, int64_t targetCore, Type payloadType, bool projected) { - auto it = llvm::find_if(summaries, [&](const PendingSummary& summary) { - return summary.sourceClass == sourceClass && summary.sourceCore == sourceCore - && summary.targetCore == targetCore && summary.payloadType == payloadType; - }); - if (it == summaries.end()) { - summaries.push_back(PendingSummary { - .sourceClass = sourceClass, - .sourceCore = sourceCore, - .targetCore = targetCore, - .payloadType = payloadType, - }); - it = std::prev(summaries.end()); - } - if (projected) - ++it->projectedCount; - else - ++it->ordinaryCount; - }; - - for (const PendingScalarSend& pending : state.pendingScalarSends) - accumulate(pending.sourceClass, pending.sourceCore, pending.targetCore, pending.payloadType, /*projected=*/false); - for (const PendingProjectedScalarSend& pending : state.pendingProjectedScalarSends) - accumulate(pending.sourceClass, pending.sourceCore, pending.targetCore, pending.payloadType, /*projected=*/true); - - if (summaries.empty()) - return success(); - - auto diag = state.func.emitError("scalar peer communication materialization left pending sends"); - for (const PendingSummary& summary : summaries) - diag << " [sourceClass=" << summary.sourceClass << " sourceCore=" << summary.sourceCore - << " targetCore=" << summary.targetCore << " payloadType=" << summary.payloadType - << " ordinaryPending=" << summary.ordinaryCount << " projectedPending=" << summary.projectedCount << "]"; - return failure(); +template +SmallVector filterVector(ArrayRef values, ArrayRef indices) { + SmallVector filtered; + filtered.reserve(indices.size()); + for (size_t index : indices) + filtered.push_back(values[index]); + return filtered; } -LogicalResult splitScalarPeerFanoutForImmediateEmission(MaterializerState& state, - MaterializedClass& sourceClass, - Value payload, - const ScalarSourceFanoutPlan& fanoutPlan, - Location loc, - ScalarSourceFanoutPlan& immediatePlan) { - int64_t sourceCpu = static_cast(sourceClass.cpus.front()); - - immediatePlan.receivePlans = fanoutPlan.receivePlans; - if (fanoutPlan.ordinaryMessages) { - for (size_t index = 0; index < fanoutPlan.ordinaryMessages->size(); ++index) { - MessageVector message; - message.append(fanoutPlan.ordinaryMessages->channelIds[index], - fanoutPlan.ordinaryMessages->sourceCoreIds[index], - fanoutPlan.ordinaryMessages->targetCoreIds[index]); - int64_t targetCore = fanoutPlan.ordinaryMessages->targetCoreIds[index]; - ScalarPeerEdgeKey requiredReceive = makeRequiredPeerReceiveForLowerToHigherSend(sourceCpu, targetCore, payload.getType()); - bool deferSend = sourceCpu < targetCore && hasPlannedScalarPeerReceive(state, requiredReceive) - && !hasMaterializedScalarPeerReceive(state, requiredReceive); - if (deferSend) { - if (failed(enqueuePendingScalarSend( - state, sourceClass.id, sourceCpu, targetCore, payload.getType(), payload, message, loc))) - return failure(); - continue; - } - if (!immediatePlan.ordinaryMessages) - immediatePlan.ordinaryMessages = MessageVector {}; - immediatePlan.ordinaryMessages->append(message.channelIds, message.sourceCoreIds, message.targetCoreIds); - } - } - - for (const ProjectedScalarSendGroup& group : fanoutPlan.projectedSendGroups) { - assert(succeeded(verifyProjectedSendDescriptor(sourceClass.op, group.descriptor, group.messages)) - && "projected scalar send group metadata is inconsistent"); - FailureOr projectedCommunicationType = getProjectedCommunicationType(sourceClass.op, group.descriptor); - if (failed(projectedCommunicationType)) - return failure(); - - std::optional immediateGroup; - SmallVector deferredGroups; - for (size_t index = 0; index < group.messages.size(); ++index) { - int64_t targetCore = group.messages.targetCoreIds[index]; - ScalarPeerEdgeKey requiredReceive = - makeRequiredPeerReceiveForLowerToHigherSend(sourceCpu, targetCore, *projectedCommunicationType); - bool deferSend = sourceCpu < targetCore && hasPlannedScalarPeerReceive(state, requiredReceive) - && !hasMaterializedScalarPeerReceive(state, requiredReceive); - if (!deferSend) { - if (!immediateGroup) - immediateGroup = makeEmptyProjectedScalarSendGroup(group); - appendProjectedScalarSendMessage(*immediateGroup, group, index); - continue; - } - - auto deferredIt = llvm::find_if(deferredGroups, [&](const ProjectedScalarSendGroup& deferredGroup) { - return !deferredGroup.messages.empty() && deferredGroup.messages.targetCoreIds.front() == targetCore; - }); - if (deferredIt == deferredGroups.end()) { - deferredGroups.push_back(makeEmptyProjectedScalarSendGroup(group)); - deferredIt = std::prev(deferredGroups.end()); - } - appendProjectedScalarSendMessage(*deferredIt, group, index); - } - - if (immediateGroup) { - if (failed(finalizeSplitProjectedScalarSendGroup(sourceClass.op, *immediateGroup))) - return failure(); - immediatePlan.projectedSendGroups.push_back(std::move(*immediateGroup)); - } - - for (ProjectedScalarSendGroup& deferredGroup : deferredGroups) { - if (failed(finalizeSplitProjectedScalarSendGroup(sourceClass.op, deferredGroup))) - return failure(); - int64_t targetCore = deferredGroup.messages.targetCoreIds.front(); - if (failed(enqueuePendingProjectedScalarSend(state, - sourceClass.id, - sourceCpu, - targetCore, - *projectedCommunicationType, - payload, - deferredGroup, - loc))) - return failure(); - } - } - - return success(); -} LogicalResult emitScalarSourceCommunication( MaterializerState& state, MaterializedClass& sourceClass, ArrayRef keys, Value payload, Location loc) { @@ -6260,35 +5500,30 @@ LogicalResult emitScalarSourceCommunication( state.availableValues.record(key, sourceClass.id, payload); SmallVector destinationClasses = collectDestinationClassesForKeys(state, keys); - auto fanoutPlan = buildScalarSourceFanoutPlan(state, sourceClass, keys, destinationClasses, payload); - if (failed(fanoutPlan)) - return failure(); - ScalarSourceFanoutPlan immediatePlan; - if (failed(splitScalarPeerFanoutForImmediateEmission(state, sourceClass, payload, *fanoutPlan, loc, immediatePlan))) - return failure(); - - if (failed(emitScalarSourceFanoutSends(state, sourceClass, payload, immediatePlan, loc))) - return failure(); - - for (const ScalarSourceReceivePlan& plan : fanoutPlan->receivePlans) { - MaterializedClass& targetClass = state.classes[plan.targetClass]; - - FailureOr received = appendReceiveAndFlushPendingScalarSends(state, targetClass, plan.receiveType, plan.messages, loc); - if (failed(received)) - return failure(); - - if (plan.projectedExtractOp) { - state.projectedExtractReplacements[plan.projectedExtractOp][plan.targetClass] = ProjectedExtractReplacement { - *received, - plan.projectedLayout - }; + for (ClassId destinationClass : destinationClasses) { + if (destinationClass == sourceClass.id) continue; + for (ProducerKey key : keys) { + ArrayRef ids = + state.communicationPlan.getOrdinaryValueExchanges(key, destinationClass); + if (ids.empty()) + continue; + SmallVector boundIds; + for (CommunicationExchangeId id : ids) { + if (state.communicationPlan.getExchange(id).sourceClass != sourceClass.id) + continue; + if (failed(bindCommunicationPayload(state, id, sourceClass, payload, sourceClass.op, loc))) + return failure(); + boundIds.push_back(id); + } + if (!boundIds.empty()) + state.availableValues.recordDeferredExact(key, destinationClass, boundIds); } - - for (ProducerKey key : keys) - state.availableValues.record(key, targetClass.id, *received); } + FailureOr emitted = emitReadyCommunicationEventsForClass(state, sourceClass, loc); + if (failed(emitted)) + return failure(); return success(); } @@ -6310,33 +5545,48 @@ LogicalResult emitClassToClassCommunication(MaterializerState& state, return sourceClass.op->emitError("scalar-source communication must be emitted through the scalar fanout planner"); if (!targetClass.isBatch) { - MessageVector messages; - messages.channelIds.reserve(keys.size()); - messages.sourceCoreIds.reserve(keys.size()); - messages.targetCoreIds.reserve(keys.size()); - - auto targetCpu = getCheckedCoreId(targetClass.op, targetClass.cpus.front(), "batch-to-scalar target core id"); - if (failed(targetCpu)) - return failure(); if (keys.size() != sourceClass.cpus.size()) return sourceClass.op->emitError("batch-to-scalar communication expects one producer key per source lane"); - for (CpuId sourceCpu : sourceClass.cpus) { - auto checkedSourceCpu = getCheckedCoreId(sourceClass.op, sourceCpu, "batch-to-scalar source core id"); - if (failed(checkedSourceCpu)) - return failure(); - messages.append(state.nextChannelId++, *checkedSourceCpu, *targetCpu); + + SmallVector exchangeIds; + for (ProducerKey key : keys) { + ArrayRef ids = + state.communicationPlan.getPackedScalarRunExchanges(key, targetClass.id); + if (ids.empty()) + return sourceClass.op->emitError("missing planned batch-to-scalar exchange"); + for (CommunicationExchangeId id : ids) { + if (failed(bindCommunicationPayload(state, id, sourceClass, payload, sourceClass.op, loc))) + return failure(); + exchangeIds.push_back(id); + } } - if (failed(appendSend(state, sourceClass, payload, messages, loc))) + if (keys.size() == 1) { + state.availableValues.recordDeferredExact(keys.front(), targetClass.id, exchangeIds); + return emitReadyCommunicationEventsForClass(state, sourceClass, loc); + } + + FailureOr fragmentType = getPackedSlotFragmentType(payload, keys.size()); + if (failed(fragmentType)) + return sourceClass.op->emitError("batch-to-scalar communication could not recover packed fragment type"); + + PackedScalarRunValue packedRun; + packedRun.targetClass = targetClass.id; + packedRun.sourceOp = keys.front().instance.op; + packedRun.resultIndex = keys.front().resultIndex; + packedRun.kind = PackedScalarRunKind::DeferredReceive; + packedRun.fragmentType = *fragmentType; + packedRun.exchangeIds = std::move(exchangeIds); + packedRun.slots.reserve(keys.size()); + for (ProducerKey key : keys) { + PackedScalarRunSlot slot; + slot.keys.push_back(key); + packedRun.slots.push_back(std::move(slot)); + } + if (failed(validatePackedScalarRunMetadata(targetClass.op, packedRun))) return failure(); - return registerLazyPackedScalarReceives(state, - sourceClass, - targetClass, - keys, - payload.getType(), - messages.channelIds, - messages.sourceCoreIds, - messages.targetCoreIds); + state.availableValues.recordPackedRun(std::move(packedRun)); + return emitReadyCommunicationEventsForClass(state, sourceClass, loc); } if (sourceClass.cpus.size() != targetClass.cpus.size()) @@ -6346,55 +5596,44 @@ LogicalResult emitClassToClassCommunication(MaterializerState& state, << " targetClass=" << targetClass.id << " targetSize=" << targetClass.cpus.size() << " keyCount=" << keys.size(); - MessageVector messages; - messages.channelIds.reserve(sourceClass.cpus.size()); - messages.sourceCoreIds.reserve(sourceClass.cpus.size()); - messages.targetCoreIds.reserve(targetClass.cpus.size()); - SmallVector lowerToHigherOrLocal; - SmallVector higherToLower; + if (keys.size() != sourceClass.cpus.size()) + return sourceClass.op->emitError("batch-to-batch communication expects one producer key per source lane"); - for (auto [lane, sourceCpu] : llvm::enumerate(sourceClass.cpus)) { - auto checkedSourceCpu = getCheckedCoreId(sourceClass.op, sourceCpu, "batch source core id"); - if (failed(checkedSourceCpu)) - return failure(); - auto checkedTargetCpu = getCheckedCoreId(targetClass.op, targetClass.cpus[lane], "batch target core id"); - if (failed(checkedTargetCpu)) - return failure(); - messages.append(state.nextChannelId++, *checkedSourceCpu, *checkedTargetCpu); - if (*checkedSourceCpu <= *checkedTargetCpu) - lowerToHigherOrLocal.push_back(lane); - else - higherToLower.push_back(lane); + auto fragmentType = dyn_cast(payload.getType()); + if (!fragmentType || !fragmentType.hasStaticShape() || fragmentType.getRank() == 0) + return sourceClass.op->emitError("batch-to-batch communication requires a static ranked lane payload"); + + Operation* sourceOp = keys.front().instance.op; + size_t resultIndex = keys.front().resultIndex; + for (ProducerKey key : keys) + if (key.instance.op != sourceOp || key.resultIndex != resultIndex || key.instance.laneCount != 1) + return sourceClass.op->emitError("batch-to-batch communication expects one-lane keys from one producer result"); + + SmallVector exchangeIds; + for (ProducerKey key : keys) { + ArrayRef ids = + state.communicationPlan.getIndexedBatchRunExchanges(key, targetClass.id); + if (ids.empty()) + return sourceClass.op->emitError("missing planned batch-to-batch exchange"); + for (CommunicationExchangeId id : ids) { + if (failed(bindCommunicationPayload(state, id, sourceClass, payload, sourceClass.op, loc))) + return failure(); + exchangeIds.push_back(id); + } } - auto recordReceived = [&](ArrayRef indices, Value received) { - for (size_t index : indices) - state.availableValues.record(keys[index], targetClass.id, received); - }; + IndexedBatchRunValue indexedRun; + indexedRun.targetClass = targetClass.id; + indexedRun.sourceOp = sourceOp; + indexedRun.resultIndex = resultIndex; + indexedRun.fragmentType = fragmentType; + indexedRun.exchangeIds = std::move(exchangeIds); + PackedScalarRunSlot slot; + slot.keys.append(keys.begin(), keys.end()); + indexedRun.slots.push_back(std::move(slot)); + state.availableValues.recordIndexedBatchRun(std::move(indexedRun)); - if (!lowerToHigherOrLocal.empty()) { - MessageVector sendFirst = filterMessageVector(messages, lowerToHigherOrLocal); - if (failed(appendSend(state, sourceClass, payload, sendFirst, loc))) - return failure(); - FailureOr received = - appendReceiveAndFlushPendingScalarSends(state, targetClass, payload.getType(), sendFirst, loc); - if (failed(received)) - return failure(); - recordReceived(lowerToHigherOrLocal, *received); - } - - if (!higherToLower.empty()) { - MessageVector receiveFirst = filterMessageVector(messages, higherToLower); - FailureOr received = - appendReceiveAndFlushPendingScalarSends(state, targetClass, payload.getType(), receiveFirst, loc); - if (failed(received)) - return failure(); - if (failed(appendSend(state, sourceClass, payload, receiveFirst, loc))) - return failure(); - recordReceived(higherToLower, *received); - } - - return success(); + return emitReadyCommunicationEventsForClass(state, sourceClass, loc); } LogicalResult @@ -6459,36 +5698,49 @@ emitHostCommunication(MaterializerState& state, MaterializedClass& sourceClass, if (sourceClass.id == ownerClass.id) return setHostOutputValue(state, ownerClass, originalOutput, payload); - // Keep the old deadlock-free communication discipline: only scalar-to-scalar - // host-owner forwarding is introduced here. Batch host publication remains on - // the owning batch path; projected terminal batch publication must use the - // explicit projected whole-batch path instead of generic host forwarding. + ArrayRef ids = + state.communicationPlan.getHostForwardExchanges(originalOutput, sourceClass.id, ownerClass.id); + if (ids.empty()) + return sourceClass.op->emitError("missing planned host-forward exchange"); + for (CommunicationExchangeId id : ids) + if (failed(bindCommunicationPayload(state, id, sourceClass, payload, sourceClass.op, payload.getLoc()))) + return failure(); + FailureOr emitted = emitReadyCommunicationEventsForClass(state, sourceClass, payload.getLoc()); + if (failed(emitted)) + return failure(); + if (sourceClass.isBatch && ownerClass.isBatch) { if (sourceClass.cpus.size() != ownerClass.cpus.size()) return sourceClass.op->emitError("batch host publication requires batch source/owner classes of equal size"); + if (ids.size() != ownerClass.cpus.size()) + return sourceClass.op->emitError("planned batch host-forward exchange count does not match owner lanes"); - MessageVector messages; - messages.channelIds.reserve(sourceClass.cpus.size()); - messages.sourceCoreIds.reserve(sourceClass.cpus.size()); - messages.targetCoreIds.reserve(ownerClass.cpus.size()); - - for (auto [lane, sourceCpu] : llvm::enumerate(sourceClass.cpus)) { - auto checkedSourceCpu = getCheckedCoreId(sourceClass.op, sourceCpu, "host batch source core id"); - if (failed(checkedSourceCpu)) + for (CommunicationExchangeId id : ids) + if (!state.communicationExchangeStates[id.value].sendEmitted) return failure(); - auto checkedTargetCpu = getCheckedCoreId(ownerClass.op, ownerClass.cpus[lane], "host batch owner core id"); - if (failed(checkedTargetCpu)) - return failure(); - messages.append(state.nextChannelId++, *checkedSourceCpu, *checkedTargetCpu); + FailureOr messages = + buildMessageVectorForExchangeIds(state, + ids, + CommunicationExchangeKind::HostForward, + sourceClass.id, + ownerClass.id, + payload.getType(), + ownerClass.op); + if (failed(messages)) + return failure(); + Value ownerPayload = appendReceive(state, ownerClass, payload.getType(), *messages, payload.getLoc()); + if (!ownerPayload) + return failure(); + Operation* receiveAnchor = nullptr; + if (Operation* receiveDef = ownerPayload.getDefiningOp()) + receiveAnchor = getTopLevelMaterializedClassBodyOp(receiveDef, ownerClass); + for (CommunicationExchangeId id : ids) { + CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[id.value]; + emission.receiveEmitted = true; + emission.receiveValue = ownerPayload; + emission.receiveAnchor = receiveAnchor; } - - if (failed(appendSend(state, sourceClass, payload, messages, payload.getLoc()))) - return failure(); - FailureOr ownerPayload = - appendReceiveAndFlushPendingScalarSends(state, ownerClass, payload.getType(), messages, payload.getLoc()); - if (failed(ownerPayload)) - return failure(); - return setHostOutputValue(state, ownerClass, originalOutput, *ownerPayload); + return setHostOutputValue(state, ownerClass, originalOutput, ownerPayload); } if (sourceClass.isBatch) @@ -6499,17 +5751,9 @@ emitHostCommunication(MaterializerState& state, MaterializedClass& sourceClass, return sourceClass.op->emitError("cannot forward fragment payload to scalar host owner") << " payloadType=" << payload.getType() << " outputType=" << originalOutput.getType(); - MessageVector messages; - auto checkedSourceCpu = getCheckedCoreId(sourceClass.op, sourceClass.cpus.front(), "host source core id"); - auto checkedTargetCpu = getCheckedCoreId(ownerClass.op, ownerClass.cpus.front(), "host target core id"); - if (failed(checkedSourceCpu) || failed(checkedTargetCpu)) - return failure(); - messages.append(state.nextChannelId++, *checkedSourceCpu, *checkedTargetCpu); - - if (failed(appendSend(state, sourceClass, payload, messages, payload.getLoc()))) - return failure(); - FailureOr ownerPayload = - appendReceiveAndFlushPendingScalarSends(state, ownerClass, payload.getType(), messages, payload.getLoc()); + if (ids.size() != 1) + return sourceClass.op->emitError("scalar host forwarding expected one planned exchange"); + FailureOr ownerPayload = emitPlannedReceive(state, ids.front(), ownerClass, payload.getLoc()); if (failed(ownerPayload)) return failure(); return setHostOutputValue(state, ownerClass, originalOutput, *ownerPayload); @@ -6521,13 +5765,19 @@ LogicalResult emitOutputFanout(MaterializerState& state, Value payload, Value originalOutput, Location loc) { - if (keys.empty()) + SmallVector communicationKeys(keys.begin(), keys.end()); + for (ProducerKey projectedKey : collectProjectedProducerKeysForOutput(state, originalOutput)) + if (!llvm::is_contained(communicationKeys, projectedKey)) + communicationKeys.push_back(projectedKey); + + if (keys.empty() && communicationKeys.empty()) return success(); if (!sourceClass.isBatch) { - if (failed(emitScalarSourceCommunication(state, sourceClass, keys, payload, loc))) - return failure(); - if (failed(emitPlannedProjectedInputSendsForKeys(state, sourceClass, keys, payload, loc))) + if (!keys.empty()) + if (failed(emitScalarSourceCommunication(state, sourceClass, keys, payload, loc))) + return failure(); + if (failed(emitPlannedProjectedInputSendsForKeys(state, sourceClass, communicationKeys, payload, loc))) return failure(); return emitHostCommunication(state, sourceClass, payload, originalOutput); @@ -6542,11 +5792,12 @@ LogicalResult emitOutputFanout(MaterializerState& state, recordedProjectedHostFragments = *recorded; } - if (failed(emitPlannedProjectedInputSendsForKeys(state, sourceClass, keys, payload, loc))) + if (failed(emitPlannedProjectedInputSendsForKeys(state, sourceClass, communicationKeys, payload, loc))) return failure(); for (ClassId destinationClass : collectDestinationClassesForKeys(state, keys)) { - SmallVector destinationKeys = filterProducerKeysForDestination(state, keys, destinationClass); + SmallVector destinationKeys = + filterProducerKeysForDestination(state, keys, destinationClass); if (destinationKeys.empty()) continue; if (sourceClass.isBatch && destinationKeys.size() != keys.size()) @@ -6565,39 +5816,12 @@ LogicalResult emitOutputFanout(MaterializerState& state, if (keys.size() == 1) state.availableValues.record(keys.front(), sourceClass.id, payload); - else if (failed(recordLocalPackedSlotValue(state, sourceClass, keys, payload))) + else if (!keys.empty() && failed(recordLocalPackedSlotValue(state, sourceClass, keys, payload))) return failure(); return success(); } -LogicalResult collectPlannedScalarPeerReceives(MaterializerState& state) { - for (const ComputeInstance& instance : state.schedule.dominanceOrderCompute) { - auto cpuIt = state.schedule.computeToCpuMap.find(instance); - if (cpuIt == state.schedule.computeToCpuMap.end()) - return instance.op->emitError("missing CPU assignment while precomputing scalar peer receives"); - - auto classIt = state.cpuToClass.find(cpuIt->second); - if (classIt == state.cpuToClass.end()) - return instance.op->emitError("missing materialized class while precomputing scalar peer receives"); - - MaterializedClass& sourceClass = state.classes[classIt->second]; - if (sourceClass.isBatch) - continue; - - ArrayRef outputs = getComputeInstanceOutputValuesCached(state, instance); - for (Value output : outputs) { - SmallVector keys = collectProducerKeysForDestinations(output, instance); - if (keys.empty()) - continue; - if (failed(collectPlannedScalarPeerReceivesForSource(state, sourceClass, keys, output))) - return failure(); - } - } - - return success(); -} - struct DirectWholeBatchFragment { ProducerKey key; Value fragment; @@ -6614,7 +5838,7 @@ struct WholeBatchFragmentGroup { WholeBatchFragmentSourceKind kind = WholeBatchFragmentSourceKind::DirectValue; RankedTensorType fragmentType; SmallVector outputOffsets; - MessageVector messages; + SmallVector exchangeIds; Operation* sourceOp = nullptr; size_t resultIndex = 0; SmallVector sourceLanes; @@ -6644,7 +5868,7 @@ struct ProjectedWholeBatchFragmentGroup { SmallVector, 4> offsetsByDim; SmallVector, 4> sizesByDim; SmallVector, 4> stridesByDim; - MessageVector messages; + SmallVector exchangeIds; SmallVector redundantOps; Value packed; RankedTensorType packedSourceType; @@ -6668,7 +5892,8 @@ ProjectedWholeBatchFragmentGroup filterProjectedDeferredReceiveGroup(const Proje ProjectedWholeBatchFragmentGroup filtered; filtered.kind = group.kind; filtered.fragmentType = group.fragmentType; - filtered.messages = filterMessageVector(group.messages, indices); + for (size_t index : indices) + filtered.exchangeIds.push_back(group.exchangeIds[index]); filtered.offsetsByDim.resize(group.offsetsByDim.size()); filtered.sizesByDim.resize(group.sizesByDim.size()); filtered.stridesByDim.resize(group.stridesByDim.size()); @@ -6795,14 +6020,34 @@ std::optional getConstantIndexValue(Value value) { return std::nullopt; } -bool appendConstantChannelReceiveMessage(MessageVector& messages, SpatChannelReceiveOp receive) { +std::optional findPlannedExchangeForReceive(MaterializerState& state, + SpatChannelReceiveOp receive, + CommunicationExchangeKind expectedKind) { std::optional channelId = getConstantIndexValue(receive.getChannelId()); std::optional sourceCoreId = getConstantIndexValue(receive.getSourceCoreId()); std::optional targetCoreId = getConstantIndexValue(receive.getTargetCoreId()); if (!channelId || !sourceCoreId || !targetCoreId) - return false; - messages.append(*channelId, static_cast(*sourceCoreId), static_cast(*targetCoreId)); - return true; + return std::nullopt; + for (const CommunicationExchange& exchange : state.communicationPlan.getExchanges()) { + if (exchange.kind == expectedKind && exchange.channelId == *channelId && exchange.sourceCore == *sourceCoreId + && exchange.targetCore == *targetCoreId) + return exchange.id; + } + return std::nullopt; +} + +std::optional findPlannedExchangeForReceive(MaterializerState& state, + SpatChannelReceiveOp receive) { + std::optional channelId = getConstantIndexValue(receive.getChannelId()); + std::optional sourceCoreId = getConstantIndexValue(receive.getSourceCoreId()); + std::optional targetCoreId = getConstantIndexValue(receive.getTargetCoreId()); + if (!channelId || !sourceCoreId || !targetCoreId) + return std::nullopt; + for (const CommunicationExchange& exchange : state.communicationPlan.getExchanges()) { + if (exchange.channelId == *channelId && exchange.sourceCore == *sourceCoreId && exchange.targetCore == *targetCoreId) + return exchange.id; + } + return std::nullopt; } PackedScalarRunValue* findDeferredReceiveAlternativeForPackedRun(MaterializerState& state, @@ -7116,7 +6361,7 @@ LogicalResult collectWholeBatchFragmentGroups(MaterializerState& state, groupIt = std::prev(groups.end()); } - groupIt->messages.append(run->messages.channelIds, run->messages.sourceCoreIds, run->messages.targetCoreIds); + llvm::append_range(groupIt->exchangeIds, run->exchangeIds); for (const PackedScalarRunSlot& slot : run->slots) for (ProducerKey fragmentKey : slot.keys) groupIt->outputOffsets.push_back(static_cast(fragmentKey.instance.laneStart) * plan.rowsPerLane); @@ -7243,7 +6488,9 @@ LogicalResult collectWholeBatchFragmentGroups(MaterializerState& state, if (auto receive = fragment.fragment.getDefiningOp()) { if (fragment.fragment.use_empty()) { WholeBatchFragmentGroup& group = getOrCreateDeferredReceiveGroup(fragmentType); - if (appendConstantChannelReceiveMessage(group.messages, receive)) { + if (std::optional exchangeId = + findPlannedExchangeForReceive(state, receive, CommunicationExchangeKind::PackedScalarRun)) { + group.exchangeIds.push_back(*exchangeId); group.outputOffsets.push_back(outputOffset); group.redundantReceives.push_back(receive.getOperation()); continue; @@ -7265,27 +6512,37 @@ FailureOr emitWholeBatchFragmentGroup(MaterializerState& state, Location loc) { switch (group.kind) { case WholeBatchFragmentSourceKind::DeferredReceive: { + FailureOr messages = + buildMessageVectorForExchangeIds(state, + group.exchangeIds, + CommunicationExchangeKind::PackedScalarRun, + state.communicationPlan.getExchange(group.exchangeIds.front()).sourceClass, + targetClass.id, + group.fragmentType, + targetClass.op); + if (failed(messages)) + return failure(); Value updatedDestination = destination; - ReceiveMessagePartition partition = - partitionReceivesByPendingSendUnlocks(state, targetClass, group.fragmentType, group.messages); + ReceiveMessagePartition partition = partitionReceiveMessagesInPlanOrder(*messages); for (size_t index : partition.criticalStaticIndices) { state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - FailureOr received = createScalarPeerReceiveAndFlush( - state, - targetClass, - group.fragmentType, - getOrCreateIndexConstant(state.constantFolder, targetClass.op, group.messages.channelIds[index]), - getOrCreateIndexConstant(state.constantFolder, targetClass.op, group.messages.sourceCoreIds[index]), - getOrCreateIndexConstant(state.constantFolder, targetClass.op, group.messages.targetCoreIds[index]), - makeScalarPeerReceiveKey( - group.messages.sourceCoreIds[index], group.messages.targetCoreIds[index], group.fragmentType, group.messages.channelIds[index]), - loc); - if (failed(received)) - return failure(); + Value received = SpatChannelReceiveOp::create( + state.rewriter, + loc, + group.fragmentType, + getOrCreateIndexConstant(state.constantFolder, targetClass.op, messages->channelIds[index]), + getOrCreateIndexConstant(state.constantFolder, targetClass.op, messages->sourceCoreIds[index]), + getOrCreateIndexConstant(state.constantFolder, targetClass.op, messages->targetCoreIds[index])) + .getOutput(); + CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[group.exchangeIds[index].value]; + emission.receiveEmitted = true; + emission.receiveValue = received; + if (Operation* receiveDef = received.getDefiningOp()) + emission.receiveAnchor = getTopLevelMaterializedClassBodyOp(receiveDef, targetClass); FailureOr updated = createDim0InsertSliceInClass(state, targetClass, loc, - *received, + received, updatedDestination, getOrCreateIndexConstant( state.constantFolder, targetClass.op, group.outputOffsets[index])); @@ -7294,8 +6551,10 @@ FailureOr emitWholeBatchFragmentGroup(MaterializerState& state, updatedDestination = *updated; } - auto remainingMessages = filterMessageVector(group.messages, partition.remainingIndices); + auto remainingMessages = filterMessageVector(*messages, partition.remainingIndices); auto remainingOffsets = filterInt64Vector(group.outputOffsets, partition.remainingIndices); + auto remainingExchangeIds = + filterVector(ArrayRef(group.exchangeIds), partition.remainingIndices); FailureOr updated = remainingMessages.empty() ? FailureOr(updatedDestination) : emitIndexedFragmentInsertLoop( @@ -7307,9 +6566,10 @@ FailureOr emitWholeBatchFragmentGroup(MaterializerState& state, Value channelId = createIndexedChannelId(state, targetClass.op, remainingMessages, flatIndex, loc); Value sourceCoreId = createIndexedSourceCoreId(state, targetClass.op, remainingMessages, flatIndex, loc); Value targetCoreId = createIndexedTargetCoreId(state, targetClass.op, remainingMessages, flatIndex, loc); - return SpatChannelReceiveOp::create( - state.rewriter, loc, group.fragmentType, channelId, sourceCoreId, targetCoreId) - .getOutput(); + Value received = + SpatChannelReceiveOp::create(state.rewriter, loc, group.fragmentType, channelId, sourceCoreId, targetCoreId) + .getOutput(); + return received; }, [&](Value flatIndex) -> FailureOr { return createIndexedIndexValue(state, targetClass.op, remainingOffsets, flatIndex, loc); @@ -7317,6 +6577,13 @@ FailureOr emitWholeBatchFragmentGroup(MaterializerState& state, loc); if (failed(updated)) return failure(); + if (!remainingMessages.empty()) { + for (CommunicationExchangeId exchangeId : remainingExchangeIds) { + CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[exchangeId.value]; + emission.receiveEmitted = true; + emission.receiveValue = *updated; + } + } for (Operation* receive : group.redundantReceives) if (receive && receive->use_empty()) @@ -7451,30 +6718,83 @@ FailureOr emitProjectedWholeBatchFragmentInsertLoop(MaterializerState& st } static bool isProducerValueMaterialized(MaterializerState& state, ProducerKey key) { - auto cpuIt = state.schedule.computeToCpuMap.find(key.instance); + ComputeInstance scheduledInstance = getScheduledChunkForLogicalInstance(state, key.instance); + auto cpuIt = state.schedule.computeToCpuMap.find(scheduledInstance); 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); + auto rangeIt = state.scheduledInstanceToLogicalSlots.find(scheduledInstance); 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})) + if (isa(key.instance.op)) { + if (key.instance.laneStart < scheduledInstance.laneStart) return false; + uint32_t localLaneStart = key.instance.laneStart - scheduledInstance.laneStart; + if (localLaneStart + key.instance.laneCount > range.count) + return false; + for (SlotId offset = localLaneStart; offset < localLaneStart + key.instance.laneCount; ++offset) + if (!state.materializedLogicalSlots.contains({classIt->second, range.start + offset})) + return false; + return true; + } + + if (range.count != 1) + return false; + if (!state.materializedLogicalSlots.contains({classIt->second, range.start})) + return false; return true; } +static void unmarkProducerValueMaterialized(MaterializerState& state, ProducerKey key) { + ComputeInstance scheduledInstance = getScheduledChunkForLogicalInstance(state, key.instance); + auto cpuIt = state.schedule.computeToCpuMap.find(scheduledInstance); + if (cpuIt == state.schedule.computeToCpuMap.end()) + return; + auto classIt = state.cpuToClass.find(cpuIt->second); + if (classIt == state.cpuToClass.end()) + return; + auto rangeIt = state.scheduledInstanceToLogicalSlots.find(scheduledInstance); + if (rangeIt == state.scheduledInstanceToLogicalSlots.end()) + return; + + LogicalSlotRange range = rangeIt->second; + if (!isa(key.instance.op)) { + if (range.count == 1) + state.materializedLogicalSlots.erase({classIt->second, range.start}); + return; + } + + if (key.instance.laneStart < scheduledInstance.laneStart) + return; + uint32_t localLaneStart = key.instance.laneStart - scheduledInstance.laneStart; + if (localLaneStart + key.instance.laneCount > range.count) + return; + for (SlotId offset = localLaneStart; offset < localLaneStart + key.instance.laneCount; ++offset) + state.materializedLogicalSlots.erase({classIt->second, range.start + offset}); +} + static void recordDeferredProducerValue(MaterializerState& state, ProducerKey key) { if (!llvm::is_contained(state.deferredProducerValues, key)) state.deferredProducerValues.push_back(key); } +static int64_t getProducerSourceClassForDiagnostic(MaterializerState& state, ProducerKey key) { + ComputeInstance scheduled = getScheduledChunkForLogicalInstance(state, key.instance); + auto cpuIt = state.schedule.computeToCpuMap.find(scheduled); + if (cpuIt == state.schedule.computeToCpuMap.end()) + return -1; + auto classIt = state.cpuToClass.find(cpuIt->second); + if (classIt == state.cpuToClass.end()) + return -1; + return static_cast(classIt->second); +} + static bool hasDeferredMaterializationDependency(const MaterializerState& state) { - return !state.deferredProjectedExchanges.empty() || !state.deferredProducerValues.empty(); + return !state.deferredProducerValues.empty(); } std::optional getStaticProjectedPackedFragmentIndex(tensor::ExtractSliceOp extract) { @@ -7599,8 +6919,6 @@ FailureOr materializeProjectedWholeBatchInputFromFragments( 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()); @@ -7615,11 +6933,19 @@ FailureOr materializeProjectedWholeBatchInputFromFragments( recordDeferredProducerValue(state, laneKey); return failure(); } - if (!fragment) + if (!fragment) { + if (hasOrdinaryDestinationClass(state, laneKey, targetClass.id)) { + unmarkProducerValueMaterialized(state, laneKey); + recordDeferredProducerValue(state, laneKey); + return failure(); + } return targetClass.op->emitError("projected whole-batch assembly is missing a lane fragment") << " targetClass=" << targetClass.id << " lane=" << lane - << " sourceOp='" << batch->getName() << "' resultIndex=" << key.resultIndex, + << " sourceOp='" << batch->getName() << "' resultIndex=" << key.resultIndex + << " sourceClass=" << getProducerSourceClassForDiagnostic(state, laneKey) + << " ordinaryDestination=" << hasOrdinaryDestinationClass(state, laneKey, targetClass.id), failure(); + } FailureOr> offsets = evaluateStaticProjectionIndices(projection->getMixedOffsets(), *laneArg, lane); @@ -7725,7 +7051,8 @@ FailureOr materializeProjectedWholeBatchInputFromFragments( auto fragmentType = dyn_cast(receive.getOutput().getType()); if (fragmentType && receive.getOutput().use_empty()) { ProjectedWholeBatchFragmentGroup& group = getOrCreateReceiveGroup(fragmentType); - if (appendConstantChannelReceiveMessage(group.messages, receive)) { + if (std::optional exchangeId = findPlannedExchangeForReceive(state, receive)) { + group.exchangeIds.push_back(*exchangeId); appendProjectedInsertCoordinates(group, *offsets, *sizes, *strides); group.redundantOps.push_back(receive.getOperation()); grouped = true; @@ -7742,11 +7069,19 @@ FailureOr materializeProjectedWholeBatchInputFromFragments( recordDeferredProducerValue(state, laneKey); return failure(); } - if (!fragment) + if (!fragment) { + if (hasOrdinaryDestinationClass(state, laneKey, targetClass.id)) { + unmarkProducerValueMaterialized(state, laneKey); + recordDeferredProducerValue(state, laneKey); + return failure(); + } return targetClass.op->emitError("projected whole-batch assembly is missing a lane fragment") << " targetClass=" << targetClass.id << " lane=" << lane - << " sourceOp='" << batch->getName() << "' resultIndex=" << key.resultIndex, + << " sourceOp='" << batch->getName() << "' resultIndex=" << key.resultIndex + << " sourceClass=" << getProducerSourceClassForDiagnostic(state, laneKey) + << " ordinaryDestination=" << hasOrdinaryDestinationClass(state, laneKey, targetClass.id), failure(); + } auto fragmentType = dyn_cast(fragment->getType()); if (!fragmentType) @@ -7778,23 +7113,35 @@ FailureOr materializeProjectedWholeBatchInputFromFragments( for (const ProjectedWholeBatchFragmentGroup& group : groups) { FailureOr updated = failure(); switch (group.kind) { - case ProjectedWholeBatchFragmentSourceKind::DeferredReceive: + case ProjectedWholeBatchFragmentSourceKind::DeferredReceive: { + const CommunicationExchange& firstExchange = state.communicationPlan.getExchange(group.exchangeIds.front()); + FailureOr messages = + buildMessageVectorForExchangeIds(state, + group.exchangeIds, + firstExchange.kind, + firstExchange.sourceClass, + targetClass.id, + group.fragmentType, + targetClass.op); + if (failed(messages)) + return failure(); updated = result; - for (size_t index : partitionReceivesByPendingSendUnlocks(state, targetClass, group.fragmentType, group.messages) - .criticalStaticIndices) { + ReceiveMessagePartition partition = partitionReceiveMessagesInPlanOrder(*messages); + for (size_t index : partition.criticalStaticIndices) { state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - FailureOr received = createScalarPeerReceiveAndFlush( - state, - targetClass, - group.fragmentType, - getOrCreateIndexConstant(state.constantFolder, targetClass.op, group.messages.channelIds[index]), - getOrCreateIndexConstant(state.constantFolder, targetClass.op, group.messages.sourceCoreIds[index]), - getOrCreateIndexConstant(state.constantFolder, targetClass.op, group.messages.targetCoreIds[index]), - makeScalarPeerReceiveKey( - group.messages.sourceCoreIds[index], group.messages.targetCoreIds[index], group.fragmentType, group.messages.channelIds[index]), - loc); - if (failed(received)) - return failure(); + Value received = SpatChannelReceiveOp::create( + state.rewriter, + loc, + group.fragmentType, + getOrCreateIndexConstant(state.constantFolder, targetClass.op, messages->channelIds[index]), + getOrCreateIndexConstant(state.constantFolder, targetClass.op, messages->sourceCoreIds[index]), + getOrCreateIndexConstant(state.constantFolder, targetClass.op, messages->targetCoreIds[index])) + .getOutput(); + CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[group.exchangeIds[index].value]; + emission.receiveEmitted = true; + emission.receiveValue = received; + if (Operation* receiveDef = received.getDefiningOp()) + emission.receiveAnchor = getTopLevelMaterializedClassBodyOp(receiveDef, targetClass); SmallVector offsets; SmallVector sizes; @@ -7808,32 +7155,37 @@ FailureOr materializeProjectedWholeBatchInputFromFragments( strides.push_back(state.rewriter.getIndexAttr(group.stridesByDim[dim][index])); } state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); - updated = tensor::InsertSliceOp::create(state.rewriter, loc, *received, *updated, offsets, sizes, strides) + updated = tensor::InsertSliceOp::create(state.rewriter, loc, received, *updated, offsets, sizes, strides) .getResult(); } - { - ReceiveMessagePartition partition = - partitionReceivesByPendingSendUnlocks(state, targetClass, group.fragmentType, group.messages); - if (!partition.remainingIndices.empty()) { - ProjectedWholeBatchFragmentGroup remainingGroup = - filterProjectedDeferredReceiveGroup(group, partition.remainingIndices); - updated = emitProjectedWholeBatchFragmentInsertLoop( - state, - targetClass, - *updated, - remainingGroup, - [&](Value flatIndex) -> FailureOr { - Value channelId = createIndexedChannelId(state, targetClass.op, remainingGroup.messages, flatIndex, loc); - Value sourceCoreId = createIndexedSourceCoreId(state, targetClass.op, remainingGroup.messages, flatIndex, loc); - Value targetCoreId = createIndexedTargetCoreId(state, targetClass.op, remainingGroup.messages, flatIndex, loc); - return SpatChannelReceiveOp::create( - state.rewriter, loc, remainingGroup.fragmentType, channelId, sourceCoreId, targetCoreId) - .getOutput(); - }, - loc); + if (!partition.remainingIndices.empty()) { + ProjectedWholeBatchFragmentGroup remainingGroup = + filterProjectedDeferredReceiveGroup(group, partition.remainingIndices); + MessageVector remainingMessages = filterMessageVector(*messages, partition.remainingIndices); + updated = emitProjectedWholeBatchFragmentInsertLoop( + state, + targetClass, + *updated, + remainingGroup, + [&](Value flatIndex) -> FailureOr { + Value channelId = createIndexedChannelId(state, targetClass.op, remainingMessages, flatIndex, loc); + Value sourceCoreId = createIndexedSourceCoreId(state, targetClass.op, remainingMessages, flatIndex, loc); + Value targetCoreId = createIndexedTargetCoreId(state, targetClass.op, remainingMessages, flatIndex, loc); + return SpatChannelReceiveOp::create( + state.rewriter, loc, remainingGroup.fragmentType, channelId, sourceCoreId, targetCoreId) + .getOutput(); + }, + loc); + if (failed(updated)) + return failure(); + for (CommunicationExchangeId exchangeId : remainingGroup.exchangeIds) { + CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[exchangeId.value]; + emission.receiveEmitted = true; + emission.receiveValue = *updated; } } break; + } case ProjectedWholeBatchFragmentSourceKind::PackedValue: updated = emitProjectedWholeBatchFragmentInsertLoop( state, @@ -8179,12 +7531,24 @@ FailureOr resolveInputValue(MaterializerState& state, } else { size_t laneCount = targetClass.cpus.size(); - MessageVector messages = indexedRun->messages.slice(slotIndex * laneCount, laneCount); - FailureOr appended = appendReceiveAndFlushPendingScalarSends( - state, targetClass, indexedRun->fragmentType, messages, consumerInstance.op->getLoc()); - if (failed(appended)) + MessageVector messages; + if (indexedRun->exchangeIds.empty()) + return targetClass.op->emitError("indexed batch run lookup is missing planned exchange ids"), failure(); + ArrayRef ids(indexedRun->exchangeIds.data() + slotIndex * laneCount, laneCount); + FailureOr plannedMessages = + buildMessageVectorForExchangeIds(state, + ids, + CommunicationExchangeKind::IndexedBatchRun, + state.communicationPlan.getExchange(ids.front()).sourceClass, + targetClass.id, + indexedRun->fragmentType, + targetClass.op); + if (failed(plannedMessages)) + return failure(); + messages = std::move(*plannedMessages); + received = appendReceive(state, targetClass, indexedRun->fragmentType, messages, consumerInstance.op->getLoc()); + if (!received) return failure(); - received = *appended; } state.availableValues.record(*producer, targetClass.id, received); return rejectNonLocalResolvedValue(received); @@ -8218,6 +7582,23 @@ FailureOr resolveInputValue(MaterializerState& state, return failure(); } + ArrayRef deferredExactIds = + state.availableValues.getDeferredExactExchangeIds(*producer, targetClass.id); + if (!deferredExactIds.empty()) { + bool hasUnboundExchange = false; + for (CommunicationExchangeId id : deferredExactIds) { + if (id.value >= state.communicationExchangeStates.size() + || !state.communicationExchangeStates[id.value].producerBound) { + hasUnboundExchange = true; + break; + } + } + if (hasUnboundExchange) + unmarkProducerValueMaterialized(state, *producer); + recordDeferredProducerValue(state, *producer); + return failure(); + } + std::optional producerClass; if (auto cpuIt = state.schedule.computeToCpuMap.find(producer->instance); cpuIt != state.schedule.computeToCpuMap.end()) { @@ -8362,6 +7743,10 @@ LogicalResult mapInputs(MaterializerState& state, Operation* op = instance.op; if (auto compute = dyn_cast(op)) { for (auto [index, input] : llvm::enumerate(compute.getInputs())) { + if (hasProjectedComputeInputReplacement(state, compute, static_cast(index), targetClass.id) + && !hasUnreplacedComputeInputUse(state, compute, static_cast(index), targetClass.id)) + continue; + FailureOr mapped = resolveInputValue(state, targetClass, input, instance, indexing); if (failed(mapped)) { if (hasDeferredMaterializationDependency(state)) @@ -8396,9 +7781,25 @@ 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) - && !hasUnreplacedBatchInputUse(state, batch, static_cast(index), targetClass.id)) + bool hasProjectedReplacement = + hasProjectedInputReplacement(state, batch, static_cast(index), input, instance, targetClass.id); + bool hasUnreplacedUse = + hasUnreplacedBatchInputUse(state, batch, static_cast(index), targetClass.id); + if (isDebugPaddedVggInput(input)) { + static unsigned debugCount = 0; + if (debugCount++ < 12) { + llvm::errs() << "debug padded mapInputs targetClass=" << targetClass.id + << " inputIndex=" << index + << " hasProjectedReplacement=" << hasProjectedReplacement + << " hasUnreplacedUse=" << hasUnreplacedUse + << " canUseProjectedLane=" + << canUseProjectedLaneInput(state, batch, static_cast(index), input, instance) + << " hasWholeBatchReplacement=" + << getProjectedWholeBatchReplacementProducer(state, batch, static_cast(index)).has_value() + << "\n"; + } + } + if (hasProjectedReplacement && !hasUnreplacedUse) continue; FailureOr demand = @@ -8535,6 +7936,89 @@ std::optional lookupProjectedExtractReplacement(Mat return classIt->second; } +static std::optional getBatchInputKeyForExtractSource(tensor::ExtractSliceOp extract) { + auto sourceArg = dyn_cast(extract.getSource()); + if (!sourceArg) + return std::nullopt; + + auto batch = dyn_cast_or_null(sourceArg.getOwner()->getParentOp()); + if (!batch) + return std::nullopt; + + for (unsigned inputIndex = 0; inputIndex < batch.getInputs().size(); ++inputIndex) { + std::optional inputArg = batch.getInputArgument(inputIndex); + if (inputArg && *inputArg == sourceArg) + return ProjectedBatchInputKey {batch.getOperation(), inputIndex}; + } + + return std::nullopt; +} + +static std::optional +lookupProjectedExtractReplacementForInputKey(MaterializerState& state, + MaterializedClass& targetClass, + ProjectedBatchInputKey inputKey) { + auto lookupReplacementForExtract = [&](Operation* extractOp) -> std::optional { + auto replacementIt = state.projectedExtractReplacements.find(extractOp); + if (replacementIt == state.projectedExtractReplacements.end()) + return std::nullopt; + auto classIt = replacementIt->second.find(targetClass.id); + if (classIt == replacementIt->second.end()) + return std::nullopt; + return classIt->second; + }; + + for (auto& extractEntry : state.projectedInputTransferPlans) { + auto classIt = extractEntry.second.find(targetClass.id); + if (classIt != extractEntry.second.end() && classIt->second.inputKey == inputKey) { + if (std::optional replacement = lookupReplacementForExtract(extractEntry.first)) + return replacement; + } + } + + for (auto& producerEntry : state.projectedTransfers) { + auto classIt = producerEntry.second.find(targetClass.id); + if (classIt != producerEntry.second.end() && classIt->second.inputKey == inputKey) { + if (std::optional replacement = + lookupReplacementForExtract(classIt->second.extractOp)) + return replacement; + } + } + + return std::nullopt; +} + +static Operation* getProjectedInputPlanExtractForInputKey(MaterializerState& state, + MaterializedClass& targetClass, + ProjectedBatchInputKey inputKey) { + for (auto& extractEntry : state.projectedInputTransferPlans) { + auto classIt = extractEntry.second.find(targetClass.id); + if (classIt != extractEntry.second.end() && classIt->second.inputKey == inputKey) + return extractEntry.first; + } + return nullptr; +} + +static std::optional getProjectedInputKeyProducerForClass(MaterializerState& state, + MaterializedClass& targetClass, + ProjectedBatchInputKey inputKey) { + for (auto& producerEntry : state.projectedTransfers) { + auto classIt = producerEntry.second.find(targetClass.id); + if (classIt != producerEntry.second.end() && classIt->second.inputKey == inputKey) + return producerEntry.first; + } + + for (auto& extractEntry : state.projectedInputTransferPlans) { + auto classIt = extractEntry.second.find(targetClass.id); + if (classIt == extractEntry.second.end() || !(classIt->second.inputKey == inputKey) + || classIt->second.fragments.empty()) + continue; + return classIt->second.fragments.front().producer; + } + + return std::nullopt; +} + FailureOr> getOrMaterializeProjectedExtractReplacement( MaterializerState& state, MaterializedClass& targetClass, @@ -8545,28 +8029,57 @@ FailureOr> getOrMaterializeProjectedE return replacement; auto planIt = state.projectedInputTransferPlans.find(extract.getOperation()); - if (planIt == state.projectedInputTransferPlans.end()) - return std::optional {}; + if (planIt != state.projectedInputTransferPlans.end()) { + auto classIt = planIt->second.find(targetClass.id); + if (classIt != planIt->second.end()) { + if (failed(materializeProjectedInputReceivesForClass( + state, targetClass, extract.getOperation(), extract.getLoc(), insertionAnchor))) { + if (!hasDeferredMaterializationDependency(state)) + targetClass.op->emitError("failed to materialize projected input transfer replacements for class") + << " targetClass=" << targetClass.id; + return failure(); + } - auto classIt = planIt->second.find(targetClass.id); - if (classIt == planIt->second.end()) - return std::optional {}; + if (std::optional replacement = + lookupProjectedExtractReplacement(state, targetClass, extract)) + return replacement; - 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(); + 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(); + } } - if (std::optional replacement = - lookupProjectedExtractReplacement(state, targetClass, extract)) - return replacement; + if (std::optional inputKey = getBatchInputKeyForExtractSource(extract)) { + if (std::optional replacement = + lookupProjectedExtractReplacementForInputKey(state, targetClass, *inputKey)) + return replacement; + if (Operation* planExtract = getProjectedInputPlanExtractForInputKey(state, targetClass, *inputKey)) { + if (failed(materializeProjectedInputReceivesForClass( + state, targetClass, planExtract, extract.getLoc(), insertionAnchor))) { + if (!hasDeferredMaterializationDependency(state)) + targetClass.op->emitError("failed to materialize projected input key replacement for class") + << " targetClass=" << targetClass.id; + return failure(); + } + if (std::optional replacement = + lookupProjectedExtractReplacementForInputKey(state, targetClass, *inputKey)) + 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(); + if (std::optional producer = getProjectedInputKeyProducerForClass(state, targetClass, *inputKey)) { + if (getProducerSourceClassForDiagnostic(state, *producer) == static_cast(targetClass.id)) + return std::optional {}; + if (isProducerValueMaterialized(state, *producer) + && hasOrdinaryDestinationClass(state, *producer, targetClass.id)) + unmarkProducerValueMaterialized(state, *producer); + recordDeferredProducerValue(state, *producer); + return failure(); + } + } + + return std::optional {}; } Operation* getTopLevelMaterializedClassBodyOp(Operation* nestedOp, MaterializedClass& targetClass) { @@ -8597,6 +8110,24 @@ void moveNewTopLevelOpsBeforeAnchor(MaterializedClass& targetClass, Operation* a op->moveBefore(anchor); } +void moveDefinitionBeforeAnchorWithSameBlockDeps(Operation* def, Operation* anchor, DenseSet& moved) { + if (!def || !anchor || def == anchor || def->getBlock() != anchor->getBlock()) + return; + if (!moved.insert(def).second) + return; + + for (Value operand : def->getOperands()) + moveDefinitionBeforeAnchorWithSameBlockDeps(operand.getDefiningOp(), anchor, moved); + + if (anchor->isBeforeInBlock(def)) + def->moveBefore(anchor); +} + +void moveDefinitionBeforeAnchorWithSameBlockDeps(Operation* def, Operation* anchor) { + DenseSet moved; + moveDefinitionBeforeAnchorWithSameBlockDeps(def, anchor, moved); +} + FailureOr materializeProjectedWholeBatchExtractReplacement(MaterializerState& state, MaterializedClass& targetClass, tensor::ExtractSliceOp extract, @@ -8658,6 +8189,85 @@ FailureOr materializeProjectedWholeBatchExtractReplacement(MaterializerSt state.rewriter, extract.getLoc(), *fullSource, resultType, offsets, sizes, strides); } +FailureOr> materializeLocalProjectedInputExtractReplacement(MaterializerState& state, + MaterializedClass& targetClass, + tensor::ExtractSliceOp extract, + IRMapping* mapper) { + std::optional inputKey = getBatchInputKeyForExtractSource(extract); + if (!inputKey) + return std::optional {}; + + std::optional producer = getProjectedInputKeyProducerForClass(state, targetClass, *inputKey); + if (!producer) + return std::optional {}; + + if (!isProducerValueMaterialized(state, *producer)) { + recordDeferredProducerValue(state, *producer); + return failure(); + } + + std::optional fullSource = state.availableValues.lookup(state, *producer, targetClass.id); + if (!fullSource) + return std::optional {}; + + FailureOr localizedSource = materializeTensorValueForMaterializedClassUse( + state, + targetClass, + *fullSource, + extract.getOperation(), + "local projected input replacement tried to reuse a tensor from another materialized class", + *producer, + mapper); + if (failed(localizedSource)) + return failure(); + + auto remapFoldResult = [&](OpFoldResult value) -> FailureOr { + if (auto mappedValue = dyn_cast_if_present(value)) { + FailureOr localized = rematerializeIndexValueInClass( + state, targetClass, mapper ? mapper->lookupOrDefault(mappedValue) : mappedValue, extract.getLoc(), mapper); + if (failed(localized)) + return failure(); + return OpFoldResult(*localized); + } + return value; + }; + + SmallVector offsets; + SmallVector sizes; + SmallVector strides; + offsets.reserve(extract.getMixedOffsets().size()); + sizes.reserve(extract.getMixedSizes().size()); + strides.reserve(extract.getMixedStrides().size()); + + for (OpFoldResult value : extract.getMixedOffsets()) { + FailureOr localized = remapFoldResult(value); + if (failed(localized)) + return failure(); + offsets.push_back(*localized); + } + for (OpFoldResult value : extract.getMixedSizes()) { + FailureOr localized = remapFoldResult(value); + if (failed(localized)) + return failure(); + sizes.push_back(*localized); + } + for (OpFoldResult value : extract.getMixedStrides()) { + FailureOr localized = remapFoldResult(value); + if (failed(localized)) + return failure(); + strides.push_back(*localized); + } + + auto resultType = dyn_cast(extract.getType()); + if (!resultType) + return failure(); + FailureOr projected = + extractStaticSliceOrIdentity(state.rewriter, extract.getLoc(), *localizedSource, resultType, offsets, sizes, strides); + if (failed(projected)) + return failure(); + return std::optional {*projected}; +} + LogicalResult applyProjectedExtractReplacementsInClonedOp(MaterializerState& state, MaterializedClass& targetClass, Operation& originalOp, @@ -8687,6 +8297,20 @@ LogicalResult applyProjectedExtractReplacementsInClonedOp(MaterializerState& sta return success(); } + if (auto clonedExtract = dyn_cast(&clonedOp)) { + OpBuilder::InsertionGuard guard(state.rewriter); + state.rewriter.setInsertionPoint(clonedExtract); + FailureOr> localProjected = + materializeLocalProjectedInputExtractReplacement(state, targetClass, clonedExtract, &mapper); + if (failed(localProjected)) + return failure(); + if (std::optional projected = *localProjected) { + clonedExtract.getResult().replaceAllUsesWith(*projected); + state.rewriter.eraseOp(clonedExtract); + return success(); + } + } + if (std::optional producer = getProjectedWholeBatchReplacementProducer(state, originalExtract)) { auto clonedExtract = dyn_cast(&clonedOp); if (!clonedExtract) @@ -8851,6 +8475,15 @@ LogicalResult cloneComputeTemplateBody(MaterializerState& state, continue; } + FailureOr> localProjected = + materializeLocalProjectedInputExtractReplacement(state, targetClass, extract, &mapper); + if (failed(localProjected)) + return failure(); + if (std::optional projected = *localProjected) { + mapper.map(extract.getResult(), *projected); + continue; + } + if (std::optional producer = getProjectedWholeBatchReplacementProducer(state, extract)) { FailureOr projected = materializeProjectedWholeBatchExtractReplacement(state, targetClass, extract, *producer, &mapper); @@ -8953,6 +8586,15 @@ FailureOr materializeProjectedExtractReplacement(MaterializerState& state if (failed(localizedPayload)) return failure(); Value payload = *localizedPayload; + Operation* topLevelAnchor = getTopLevelMaterializedClassBodyOp(extract.getOperation(), targetClass); + Operation* payloadDef = payload.getDefiningOp(); + Operation* topLevelPayloadDef = payloadDef ? getTopLevelMaterializedClassBodyOp(payloadDef, targetClass) : nullptr; + if (topLevelAnchor && topLevelPayloadDef && topLevelAnchor != topLevelPayloadDef + && topLevelAnchor->getBlock() == topLevelPayloadDef->getBlock() + && topLevelAnchor->isBeforeInBlock(topLevelPayloadDef)) + moveDefinitionBeforeAnchorWithSameBlockDeps(topLevelPayloadDef, topLevelAnchor); + if (payloadDef && payloadDef->getBlock() == extract->getBlock() && extract->isBeforeInBlock(payloadDef)) + moveDefinitionBeforeAnchorWithSameBlockDeps(payloadDef, extract); if (replacement.layout.payloadFragmentCount == 1) return payload; @@ -9057,14 +8699,24 @@ FailureOr materializeIndexedBatchRunReceive(MaterializerState& state, return targetClass.op->emitError("indexed batch run receive requires a batch target class"); if (run.packed) return getPackedSliceForDynamicRunIndex(state, targetClass.op, run.packed, run.fragmentType, runSlotIndex, loc); - if (failed(run.messages.verify(targetClass.op))) + if (run.exchangeIds.empty()) + return targetClass.op->emitError("indexed batch run receive is missing planned exchange ids"), failure(); + FailureOr messages = + buildMessageVectorForExchangeIds(state, + run.exchangeIds, + CommunicationExchangeKind::IndexedBatchRun, + state.communicationPlan.getExchange(run.exchangeIds.front()).sourceClass, + targetClass.id, + run.fragmentType, + targetClass.op); + if (failed(messages)) return failure(); Value flatIndex = createBatchRunFlatIndex(state, targetClass, runSlotIndex, loc); std::optional preferredPeriod = static_cast(targetClass.cpus.size()); - Value channelId = createIndexedChannelId(state, targetClass.op, run.messages, flatIndex, loc, preferredPeriod); - Value sourceCoreId = createIndexedSourceCoreId(state, targetClass.op, run.messages, flatIndex, loc, preferredPeriod); - Value targetCoreId = createIndexedTargetCoreId(state, targetClass.op, run.messages, flatIndex, loc, preferredPeriod); + Value channelId = createIndexedChannelId(state, targetClass.op, *messages, flatIndex, loc, preferredPeriod); + Value sourceCoreId = createIndexedSourceCoreId(state, targetClass.op, *messages, flatIndex, loc, preferredPeriod); + Value targetCoreId = createIndexedTargetCoreId(state, targetClass.op, *messages, flatIndex, loc, preferredPeriod); state.rewriter.setInsertionPoint(targetClass.body->getTerminator()); return SpatChannelReceiveOp::create(state.rewriter, loc, run.fragmentType, channelId, sourceCoreId, targetCoreId) @@ -9132,6 +8784,36 @@ LogicalResult localizeAllScheduledBodyCaptures(MaterializerState& state, Materia for (Operation* nestedOp : bodyOps) { if (nestedOp->getBlock() == nullptr) continue; + if (auto extract = dyn_cast(nestedOp)) { + if (!isValueLegalInMaterializedClassBody(extract.getSource(), targetClass)) { + FailureOr> maybeReplacement = + getOrMaterializeProjectedExtractReplacement(state, targetClass, extract, extract.getOperation()); + if (failed(maybeReplacement)) + return failure(); + if (std::optional replacement = *maybeReplacement) { + OpBuilder::InsertionGuard guard(state.rewriter); + state.rewriter.setInsertionPoint(extract); + FailureOr projected = materializeProjectedExtractReplacement( + state, targetClass, extract, *replacement, std::nullopt, nullptr); + if (failed(projected)) + return failure(); + extract.getResult().replaceAllUsesWith(*projected); + state.rewriter.eraseOp(extract); + continue; + } + OpBuilder::InsertionGuard guard(state.rewriter); + state.rewriter.setInsertionPoint(extract); + FailureOr> localProjected = + materializeLocalProjectedInputExtractReplacement(state, targetClass, extract, nullptr); + if (failed(localProjected)) + return failure(); + if (std::optional projected = *localProjected) { + extract.getResult().replaceAllUsesWith(*projected); + state.rewriter.eraseOp(extract); + continue; + } + } + } for (OpOperand& operand : nestedOp->getOpOperands()) { Value current = operand.get(); if (isValueLegalInMaterializedClassBody(current, targetClass)) @@ -9206,6 +8888,8 @@ FailureOr> cloneInstanceBody(MaterializerState& state, state.rewriter.restoreInsertionPoint(cloneInsertionPoint); if (failed(cloneComputeTemplateBody(state, targetClass, instance, mapper, indexing))) { + if (hasDeferredMaterializationDependency(state)) + return failure(); sourceOp->emitError("failed to clone compute/compute_batch body") << " laneStart=" << instance.laneStart << " laneCount=" << instance.laneCount << " targetClass=" << targetClass.id; @@ -9462,36 +9146,47 @@ LogicalResult emitPackedRunFanout(MaterializerState& state, Location loc) { assert(!sourceClass.isBatch && "packed run fanout expects a scalar source class"); - auto fanoutPlan = buildScalarSourceFanoutPlan(state, sourceClass, keys, destinationClasses, packed); - if (failed(fanoutPlan)) - return failure(); - ScalarSourceFanoutPlan immediatePlan; - if (failed(splitScalarPeerFanoutForImmediateEmission(state, sourceClass, packed, *fanoutPlan, loc, immediatePlan))) - return failure(); + auto rankedFragmentType = dyn_cast(fragmentType); + if (!rankedFragmentType || !rankedFragmentType.hasStaticShape() || rankedFragmentType.getRank() == 0) + return sourceClass.op->emitError("packed run fanout expects static ranked fragment type"); - if (failed(emitScalarSourceFanoutSends(state, sourceClass, packed, immediatePlan, loc))) - return failure(); - - for (const ScalarSourceReceivePlan& plan : fanoutPlan->receivePlans) { - MaterializedClass& targetClass = state.classes[plan.targetClass]; - - FailureOr received = appendReceiveAndFlushPendingScalarSends(state, targetClass, plan.receiveType, plan.messages, loc); - if (failed(received)) - return failure(); - - if (plan.projectedExtractOp) { - state.projectedExtractReplacements[plan.projectedExtractOp][plan.targetClass] = ProjectedExtractReplacement { - *received, - plan.projectedLayout - }; + Operation* sourceOp = keys.front().instance.op; + size_t resultIndex = keys.front().resultIndex; + for (ClassId destinationClass : destinationClasses) { + if (destinationClass == sourceClass.id) continue; - } - if (failed(registerPackedRunValue(state, targetClass, keys, *received, fragmentType, loc))) - return failure(); + SmallVector exchangeIds; + for (ProducerKey key : keys) { + ArrayRef ids = + state.communicationPlan.getOrdinaryValueExchanges(key, destinationClass); + if (ids.empty()) + continue; + exchangeIds.push_back(ids.front()); + break; + } + if (exchangeIds.empty()) + continue; + + for (CommunicationExchangeId id : exchangeIds) + if (failed(bindCommunicationPayload(state, id, sourceClass, packed, sourceClass.op, loc))) + return failure(); + + PackedScalarRunValue packedRun; + packedRun.targetClass = destinationClass; + packedRun.sourceOp = sourceOp; + packedRun.resultIndex = resultIndex; + packedRun.kind = PackedScalarRunKind::DeferredReceive; + packedRun.fragmentType = rankedFragmentType; + llvm::append_range(packedRun.exchangeIds, exchangeIds); + PackedScalarRunSlot slot; + llvm::append_range(slot.keys, keys); + packedRun.slots.push_back(std::move(slot)); + state.availableValues.recordPackedRun(std::move(packedRun)); } - return success(); + FailureOr emitted = emitReadyCommunicationEventsForClass(state, sourceClass, loc); + return failed(emitted) ? failure() : success(); } FailureOr> cloneBatchBodyForLane(MaterializerState& state, @@ -9758,6 +9453,30 @@ bool hasMaterializationRunGroupSameClassConsumer(MaterializerState& state, return false; } +bool hasPlannedProjectedInputSendForKeys(MaterializerState& state, ArrayRef keys) { + DenseSet keySet; + for (ProducerKey key : keys) + keySet.insert(key); + + for (auto& extractEntry : state.projectedInputTransferPlans) + for (auto& classEntry : extractEntry.second) + for (const ProjectedInputTransferFragment& fragment : classEntry.second.fragments) + if (keySet.contains(fragment.producer)) + return true; + return false; +} + +bool hasMaterializationRunGroupProjectedConsumer(MaterializerState& state, + ArrayRef run, + const OutputDestinationGroup& group) { + for (size_t resultIndex : group.resultIndices) { + SmallVector keys = getMaterializationRunOutputKeys(run, resultIndex); + if (hasPlannedProjectedInputSendForKeys(state, keys)) + return true; + } + return false; +} + bool hasMaterializationRunResultSameClassConsumer(MaterializerState& state, ClassId classId, ArrayRef run, @@ -9813,6 +9532,12 @@ FailureOr classifyRunOutputDemand(MaterializerState& state, .destinationClass = std::nullopt, .barrierReason = std::nullopt}); + SmallVector keys = getMaterializationRunOutputKeys(run, resultIndex); + if (hasPlannedProjectedInputSendForKeys(state, keys)) + demand.actions.push_back(TensorDemandAction {.kind = TensorDemandActionKind::ProjectedInputPublication, + .destinationClass = std::nullopt, + .barrierReason = std::nullopt}); + if (!hasMaterializationRunResultLiveExternalUse(state, run, resultIndex)) return demand; @@ -9911,7 +9636,8 @@ LogicalResult materializeScalarBatchRun(MaterializerState& state, for (const OutputDestinationGroup& group : groups) { if (run.size() > 1 && group.destinationClasses.empty() && !hasMaterializationRunGroupLiveExternalUse(state, run, group) - && !hasMaterializationRunGroupSameClassConsumer(state, targetClass.id, run, group)) { + && !hasMaterializationRunGroupSameClassConsumer(state, targetClass.id, run, group) + && !hasMaterializationRunGroupProjectedConsumer(state, run, group)) { for (size_t resultIndex : group.resultIndices) { if (resultIndex >= fragmentTypes.size() || !fragmentTypes[resultIndex]) return sourceBatch.emitOpError("failed to recover per-lane output type for deferred local packed run"); @@ -9996,6 +9722,9 @@ LogicalResult materializeScalarBatchRun(MaterializerState& state, if (failed(registerPackedRunValue(state, targetClass, keys, packed, fragmentType, loc))) return failure(); + if (failed(emitPlannedProjectedInputSendsForKeys(state, targetClass, keys, packed, loc))) + return failure(); + if (*recordedProjectedHostFragments) continue; @@ -10094,25 +9823,21 @@ LogicalResult buildBatchRunSendPlans(MaterializerState& state, plan.resultIndex = output.resultIndex; plan.destinationClass = destinationClass; - size_t messageCount = run.size() * sourceClass.cpus.size(); - plan.messages.channelIds.reserve(messageCount); - plan.messages.sourceCoreIds.reserve(messageCount); - plan.messages.targetCoreIds.reserve(messageCount); + size_t exchangeCount = run.size() * sourceClass.cpus.size(); + plan.exchangeIds.reserve(exchangeCount); for (size_t slotIndex = 0; slotIndex < run.size(); ++slotIndex) { - for (auto [lane, sourceCpu] : llvm::enumerate(sourceClass.cpus)) { - auto checkedSourceCpu = getCheckedCoreId(sourceClass.op, sourceCpu, "batch run source core id"); - if (failed(checkedSourceCpu)) - return failure(); - auto checkedTargetCpu = - getCheckedCoreId(targetClass.op, - targetClass.isBatch ? targetClass.cpus[lane] : targetClass.cpus.front(), - "batch run target core id"); - if (failed(checkedTargetCpu)) - return failure(); - plan.messages.append(state.nextChannelId++, *checkedSourceCpu, *checkedTargetCpu); + SmallVector slotKeys = getMaterializationRunSlotOutputKeys(run[slotIndex], output.resultIndex); + if (slotKeys.size() != sourceClass.cpus.size()) + return sourceClass.op->emitError("compact batch run send plan expected one key per source lane"); + for (ProducerKey key : slotKeys) { + ArrayRef ids = targetClass.isBatch + ? state.communicationPlan.getIndexedBatchRunExchanges(key, destinationClass) + : state.communicationPlan.getPackedScalarRunExchanges(key, destinationClass); + if (ids.empty()) + return sourceClass.op->emitError("compact batch run send plan is missing planned exchanges"); + plan.exchangeIds.push_back(ids.front()); } - (void) slotIndex; } plans.push_back(std::move(plan)); @@ -10130,12 +9855,30 @@ void appendBatchRunSend(MaterializerState& state, Location loc) { assert(sourceClass.isBatch && "batch run send expects a materialized batch source"); + MaterializedClass& targetClass = state.classes[plan.destinationClass]; + const CommunicationExchange& firstExchange = state.communicationPlan.getExchange(plan.exchangeIds.front()); + FailureOr messages = + buildMessageVectorForExchangeIds(state, + plan.exchangeIds, + firstExchange.kind, + sourceClass.id, + targetClass.id, + payload.getType(), + sourceClass.op); + if (failed(messages)) + return; std::optional preferredPeriod = static_cast(sourceClass.cpus.size()); - Value channelId = createIndexedChannelId(state, sourceClass.op, plan.messages, flatIndex, loc, preferredPeriod); - Value sourceCoreId = createIndexedSourceCoreId(state, sourceClass.op, plan.messages, flatIndex, loc, preferredPeriod); - Value targetCoreId = createIndexedTargetCoreId(state, sourceClass.op, plan.messages, flatIndex, loc, preferredPeriod); + Value channelId = createIndexedChannelId(state, sourceClass.op, *messages, flatIndex, loc, preferredPeriod); + Value sourceCoreId = createIndexedSourceCoreId(state, sourceClass.op, *messages, flatIndex, loc, preferredPeriod); + Value targetCoreId = createIndexedTargetCoreId(state, sourceClass.op, *messages, flatIndex, loc, preferredPeriod); SpatChannelSendOp::create(state.rewriter, loc, channelId, sourceCoreId, targetCoreId, payload); + for (CommunicationExchangeId id : plan.exchangeIds) { + CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[id.value]; + emission.producerBound = true; + emission.producerPayload = payload; + emission.sendEmitted = true; + } } LogicalResult appendPackedScalarRunReceives(MaterializerState& state, @@ -10150,10 +9893,7 @@ LogicalResult appendPackedScalarRunReceives(MaterializerState& state, size_t laneCount = sourceClass.cpus.size(); size_t receiveCount = run.size() * laneCount; - if (failed(plan.messages.verify(targetClass.op))) - return failure(); - - if (receiveCount != plan.messages.size()) + if (receiveCount != plan.exchangeIds.size()) return targetClass.op->emitError("inconsistent flattened batch run receive plan"); auto rankedFragmentType = dyn_cast(fragmentType); @@ -10167,7 +9907,7 @@ LogicalResult appendPackedScalarRunReceives(MaterializerState& state, packedRun.kind = PackedScalarRunKind::DeferredReceive; packedRun.fragmentType = rankedFragmentType; - packedRun.messages = plan.messages; + packedRun.exchangeIds = plan.exchangeIds; packedRun.slots.reserve(run.size()); for (const MaterializationRunSlot& slot : run) { @@ -10197,7 +9937,7 @@ LogicalResult recordIndexedBatchRunReceives(MaterializerState& state, indexedRun.sourceOp = run.front().peers.front().op; indexedRun.resultIndex = plan.resultIndex; indexedRun.fragmentType = rankedFragmentType; - indexedRun.messages = plan.messages; + indexedRun.exchangeIds = plan.exchangeIds; indexedRun.slots.reserve(run.size()); for (const MaterializationRunSlot& slot : run) { PackedScalarRunSlot indexedSlot; @@ -10330,7 +10070,6 @@ LogicalResult emitPackedBatchRunSends(MaterializerState& state, return failure(); if (sendPlans.empty()) return success(); - sortBatchRunSendPlansByPeerExchangeOrder(sendPlans); DenseMap packedOutputIndexByResult; for (auto [packedIndex, output] : llvm::enumerate(plan.outputs)) @@ -10406,6 +10145,10 @@ LogicalResult materializeBatchClassRun(MaterializerState& state, state, targetClass, run, output.resultIndex, (*packedOutputs)[packedIndex], output.fragmentType))) return failure(); break; + case TensorDemandActionKind::ProjectedInputPublication: + if (failed(emitPlannedProjectedInputSendsForKeys(state, targetClass, keys, (*packedOutputs)[packedIndex], loc))) + return failure(); + break; case TensorDemandActionKind::TerminalBlueprintPublication: { FailureOr recordedProjectedHostFragments = recordProjectedScalarHostFragmentsFromPackedValue( state, targetClass, keys, (*packedOutputs)[packedIndex], output.originalOutput, loc); @@ -10506,43 +10249,12 @@ LogicalResult materializeInstanceSlot(MaterializerState& state, const ComputeIns enum class MaterializationAttemptKind { Success, - DeferredProjectedInput, + DeferredCommunication, 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); @@ -10552,69 +10264,137 @@ static MaterializationAttemptKind attemptMaterializeInstanceSlot(MaterializerSta 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 (!state.deferredProducerValues.empty()) + return MaterializationAttemptKind::DeferredCommunication; + + SmallVector newlyMarkedSlots; + auto rangeIt = state.scheduledInstanceToLogicalSlots.find(instance); + if (rangeIt != state.scheduledInstanceToLogicalSlots.end()) { + LogicalSlotRange range = rangeIt->second; + for (SlotId slot = range.start; slot < range.start + range.count; ++slot) { + ClassSlotKey classSlot {classId, slot}; + if (!state.materializedLogicalSlots.contains(classSlot)) + newlyMarkedSlots.push_back(classSlot); + } + } if (succeeded(materializeInstanceSlot(state, instance))) return MaterializationAttemptKind::Success; - if (!state.deferredProjectedExchanges.empty() || !state.deferredProducerValues.empty()) - return MaterializationAttemptKind::DeferredProjectedInput; + if (!state.deferredProducerValues.empty()) { + for (ClassSlotKey classSlot : newlyMarkedSlots) + state.materializedLogicalSlots.erase(classSlot); + return MaterializationAttemptKind::DeferredCommunication; + } 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) +static bool isScheduledInstanceFullyMaterialized(MaterializerState& state, const ComputeInstance& instance) { + auto cpuIt = state.schedule.computeToCpuMap.find(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(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 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"); +static bool containsRemainingInstance(ArrayRef remaining, ComputeInstance instance) { + return llvm::is_contained(remaining, instance); +} + +static uint64_t getMaterializerOpId(Operation* op) { + return static_cast(reinterpret_cast(op)); +} + +static void appendDeferredProducerDiagnostic(InFlightDiagnostic& diag, + MaterializerState& state, + ProducerKey producer, + ArrayRef remaining) { + ComputeInstance scheduledInstance = getScheduledChunkForLogicalInstance(state, producer.instance); + diag << " [producerOp='" << producer.instance.op->getName() << "' laneStart=" << producer.instance.laneStart + << " laneCount=" << producer.instance.laneCount << " resultIndex=" << producer.resultIndex + << " opId=" << getMaterializerOpId(producer.instance.op) + << " sourceClass=" << getProducerSourceClassForDiagnostic(state, producer) + << " producerRemaining=" << containsRemainingInstance(remaining, producer.instance) + << " scheduledLaneStart=" << scheduledInstance.laneStart + << " scheduledLaneCount=" << scheduledInstance.laneCount + << " scheduledRemaining=" << containsRemainingInstance(remaining, scheduledInstance) + << " materialized=" << isProducerValueMaterialized(state, producer) + << " destinations="; + for (ClassId destination : getDestinationClasses(state, producer)) + diag << " " << destination; + diag << "]"; +} + +static LogicalResult emitCommunicationDeferralDiagnostic( + MaterializerState& state, + ArrayRef remaining, + const DenseMap, DenseMapInfo>& + deferredProducersByInstance) { + const ComputeInstance& instance = remaining.front(); + InFlightDiagnostic diag = instance.op->emitError("communication plan made no materialization progress") + << " remainingInstances=" << remaining.size() + << " schedule=communication-plan-events"; 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) + if (reported++ == 8) break; - diag << " [producerOp='" << producer.instance.op->getName() << "' laneStart=" << producer.instance.laneStart - << " laneCount=" << producer.instance.laneCount << " resultIndex=" << producer.resultIndex << "]"; + appendDeferredProducerDiagnostic(diag, state, producer, remaining); + } + unsigned exchangeReported = 0; + for (const ExchangeDescriptor& exchange : state.communicationPlan.getExchanges()) { + const CommunicationExchangeEmissionState& emission = state.communicationExchangeStates[exchange.id.value]; + if ((!emission.producerBound || emission.sendEmitted) && (emission.sendEmitted || emission.receiveEmitted)) + continue; + if (exchangeReported++ == 8) + break; + diag << " [exchange=" << exchange.id.value << " kind=" << static_cast(exchange.kind) + << " sourceClass=" << exchange.sourceClass << " targetClass=" << exchange.targetClass + << " sourceCore=" << exchange.sourceCore << " targetCore=" << exchange.targetCore + << " channel=" << exchange.channelId << " producerBound=" << emission.producerBound + << " sendEmitted=" << emission.sendEmitted << " receiveEmitted=" << emission.receiveEmitted << "]"; + } + unsigned remainingReported = 0; + for (const ComputeInstance& pending : remaining) { + if (remainingReported++ == 8) + break; + diag << " [pendingOp='" << pending.op->getName() << "' laneStart=" << pending.laneStart + << " laneCount=" << pending.laneCount << " opId=" << getMaterializerOpId(pending.op) << "]"; + auto producerIt = deferredProducersByInstance.find(pending); + if (producerIt != deferredProducersByInstance.end() && !producerIt->second.empty()) { + diag << " [pendingBlockedBy"; + appendDeferredProducerDiagnostic(diag, state, producerIt->second.front(), remaining); + diag << "]"; + } } return failure(); } -static LogicalResult materializeInstanceSlotsWithProjectedDeferral(MaterializerState& state) { +static LogicalResult materializeInstanceSlotsWithCommunicationPlan(MaterializerState& state) { SmallVector remaining(state.schedule.dominanceOrderCompute.begin(), state.schedule.dominanceOrderCompute.end()); while (!remaining.empty()) { bool progress = false; + DenseMap, DenseMapInfo> deferredProducersByInstance; 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; + if (isScheduledInstanceFullyMaterialized(state, instance)) { + remaining.erase(remaining.begin() + index); + continue; + } + ++index; continue; } if (attempt == MaterializationAttemptKind::Failure) { @@ -10623,25 +10403,32 @@ static LogicalResult materializeInstanceSlotsWithProjectedDeferral(MaterializerS diag << " laneStart=" << instance.laneStart << " laneCount=" << instance.laneCount; return failure(); } + deferredProducersByInstance[instance] = state.deferredProducerValues; + for (ProducerKey producer : state.deferredProducerValues) { + ComputeInstance scheduledProducer = getScheduledChunkForLogicalInstance(state, producer.instance); + if (state.schedule.computeToCpuMap.find(scheduledProducer) != state.schedule.computeToCpuMap.end() + && !isProducerValueMaterialized(state, producer) + && !containsRemainingInstance(remaining, scheduledProducer)) + remaining.push_back(scheduledProducer); + } ++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)) + bool emittedAnyCommunication = false; + for (MaterializedClass& materializedClass : state.classes) { + FailureOr emitted = + emitReadyCommunicationEventsForClass(state, materializedClass, state.func.getLoc()); + if (failed(emitted)) return failure(); - progress = true; + emittedAnyCommunication |= *emitted; + } + progress |= emittedAnyCommunication; + + if (progress) { + continue; } - if (!progress) - return emitProjectedDeferralDiagnostic(state, remaining.front(), remaining.size()); + return emitCommunicationDeferralDiagnostic(state, remaining, deferredProducersByInstance); } return success(); } @@ -10662,6 +10449,26 @@ LogicalResult eraseOldComputeOps(MaterializerState& state) { return success(); } +LogicalResult verifyMaterializerDominanceForDebug(func::FuncOp func) { + DominanceInfo dominance(func); + WalkResult result = func.walk([&](Operation* op) -> WalkResult { + for (OpOperand& operand : op->getOpOperands()) { + Value value = operand.get(); + if (dominance.dominates(value, op)) + continue; + InFlightDiagnostic diag = op->emitError("materializer produced non-dominating operand") + << " operand#" << operand.getOperandNumber() + << " operandType=" << value.getType() + << " op=\"" << truncateMaterializerDebugString(stringifyOperationForMaterializerDebug(op)) + << "\""; + attachMaterializerValueOriginNote(diag, value, "non-dominating operand"); + return WalkResult::interrupt(); + } + return WalkResult::advance(); + }); + return result.wasInterrupted() ? failure() : success(); +} + } // namespace LogicalResult @@ -10693,26 +10500,21 @@ MergeScheduleMaterializer::run(func::FuncOp func, const MergeScheduleResult& sch return failure(); if (failed(runStage("empty materialized-op creation", [&]() { return createEmptyMaterializedOps(state); }))) return failure(); - if (failed(runStage("producer-destination collection", [&]() { return collectProducerDestinations(state); }))) - return failure(); if (failed(runStage("projected-transfer collection", [&]() { return collectProjectedTransfers(state); }))) return failure(); - if (failed(runStage("projected input communication planning", [&]() { - return buildProjectedInputCommunicationPlan(state); + if (failed(runStage("producer-destination collection", [&]() { return collectProducerDestinations(state); }))) + return failure(); + if (failed(runStage("global communication planning", [&]() { + return buildGlobalCommunicationPlan(state); }))) return failure(); - if (failed(runStage("planned scalar-peer receive collection", [&]() { - return collectPlannedScalarPeerReceives(state); - }))) - return failure(); - if (failed(runStage("instance-slot materialization", [&]() -> LogicalResult { - return materializeInstanceSlotsWithProjectedDeferral(state); + return materializeInstanceSlotsWithCommunicationPlan(state); }))) return failure(); - if (failed(runStage("projected input send phase verification", [&]() { - return verifyNoPendingProjectedInputSends(state); + if (failed(runStage("communication plan materialization verification", [&]() { + return verifyCommunicationPlanFullyMaterialized(state); }))) return failure(); @@ -10733,15 +10535,16 @@ MergeScheduleMaterializer::run(func::FuncOp func, const MergeScheduleResult& sch }))) return failure(); - if (failed(runStage("pending scalar-peer send verification", [&]() { return verifyNoPendingScalarPeerSends(state); }))) - return failure(); - replaceHostUses(state); if (failed(runStage("old compute-op erasure", [&]() { return eraseOldComputeOps(state); }))) return failure(); LogicalResult _ = runRegionDCE(state.rewriter, state.func.getBody()); (void) _; + if (failed(runStage("materializer dominance verification", [&]() { + return verifyMaterializerDominanceForDebug(state.func); + }))) + return failure(); return success(); } diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializedClassState.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializedClassState.hpp index 06c6168..7337d2c 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializedClassState.hpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/MaterializedClassState.hpp @@ -58,7 +58,7 @@ struct PackedScalarRunValue { mlir::RankedTensorType fragmentType; llvm::SmallVector slots; - MessageVector messages; + llvm::SmallVector exchangeIds; }; struct IndexedBatchRunValue { @@ -68,7 +68,7 @@ struct IndexedBatchRunValue { mlir::Value packed; mlir::RankedTensorType fragmentType; llvm::SmallVector slots; - MessageVector messages; + llvm::SmallVector exchangeIds; }; struct LogicalSlotRange { @@ -90,12 +90,13 @@ struct OutputDestinationGroup { struct BatchRunSendPlan { size_t resultIndex = 0; ClassId destinationClass = 0; - MessageVector messages; + llvm::SmallVector exchangeIds; }; enum class TensorDemandActionKind { DestinationFanout, SameClassIndexedFragment, + ProjectedInputPublication, TerminalBlueprintPublication, WholeTensorBarrier }; @@ -122,44 +123,6 @@ struct CompactRunPlan { llvm::SmallVector outputs; }; -struct ScalarPeerEdgeKey { - int64_t sourceCore = 0; - int64_t targetCore = 0; - mlir::Type payloadType; -}; - -struct ScalarPeerReceiveKey { - int64_t sourceCore = 0; - int64_t targetCore = 0; - mlir::Type payloadType; - std::optional channelId; -}; - -struct PendingScalarSend { - ClassId sourceClass = 0; - int64_t sourceCore = 0; - int64_t targetCore = 0; - mlir::Type payloadType; - ScalarPeerEdgeKey waitForReceive; - mlir::Operation* payloadAnchor = nullptr; - mlir::Value payload; - MessageVector messages; - mlir::Location loc; -}; - -struct PendingProjectedScalarSend { - ClassId sourceClass = 0; - int64_t sourceCore = 0; - int64_t targetCore = 0; - mlir::Type payloadType; - ScalarPeerEdgeKey waitForReceive; - mlir::Operation* payloadAnchor = nullptr; - mlir::Value payload; - MessageVector messages; - ProjectedTransferDescriptor descriptor; - mlir::Location loc; -}; - enum class BatchInputDemandKind { LaneFragment, ProjectedFragment, @@ -171,11 +134,17 @@ struct BatchInputDemand { std::optional wholeTensorProducer; }; -struct ProjectedExchangeEmissionState { +struct CommunicationExchangeEmissionState { bool producerBound = false; bool sendEmitted = false; bool receiveEmitted = false; + mlir::Value producerPayload; + mlir::Value receiveValue; + + mlir::Operation* producerAnchor = nullptr; + mlir::Operation* sendAnchor = nullptr; + mlir::Operation* receiveAnchor = nullptr; std::optional loc; }; @@ -220,6 +189,24 @@ public: } void recordIndexedBatchRun(IndexedBatchRunValue run) { indexedBatchRuns.push_back(std::move(run)); } + void recordDeferredExact(ProducerKey key, ClassId classId, llvm::ArrayRef exchangeIds) { + llvm::SmallVector& bucket = exactExchangeIds[key][classId]; + for (CommunicationExchangeId id : exchangeIds) + if (!llvm::is_contained(bucket, id)) + bucket.push_back(id); + } + llvm::ArrayRef getDeferredExactExchangeIds(ProducerKey key, ClassId classId) const { + auto producerIt = exactExchangeIds.find(key); + if (producerIt == exactExchangeIds.end()) + return {}; + auto classIt = producerIt->second.find(classId); + if (classIt == producerIt->second.end()) + return {}; + return classIt->second; + } + bool hasDeferredExact(ProducerKey key, ClassId classId) const { + return !getDeferredExactExchangeIds(key, classId).empty(); + } std::optional lookupExact(ProducerKey key, ClassId classId) const; std::optional lookup(MaterializerState& state, ProducerKey key, ClassId classId); @@ -245,6 +232,10 @@ private: std::optional lookupPackedRun(MaterializerState& state, ProducerKey key, ClassId classId); llvm::DenseMap, ProducerKeyInfo> exactValues; + llvm::DenseMap>, + ProducerKeyInfo> + exactExchangeIds; llvm::SmallVector packedScalarRuns; llvm::SmallVector indexedBatchRuns; llvm::DenseMap> projectedInputTransferPlans; - CommunicationPlan projectedInputCommunicationPlan; - llvm::SmallVector projectedExchangeStates; - CommunicationPhase currentProjectedCommunicationPhase = CommunicationPhase::LowToHigh; - llvm::SmallVector deferredProjectedExchanges; + CommunicationPlan communicationPlan; + llvm::SmallVector communicationExchangeStates; llvm::SmallVector deferredProducerValues; AvailableValueStore availableValues; llvm::DenseMap hostReplacements; llvm::DenseMap hostOutputOwners; llvm::SmallVector pendingProjectedHostOutputFragments; - llvm::SmallVector plannedScalarPeerReceives; - llvm::SmallVector materializedScalarPeerReceives; - llvm::SmallVector pendingScalarSends; - llvm::SmallVector pendingProjectedScalarSends; llvm::DenseSet oldComputeOps; MaterializerState(mlir::func::FuncOp func, const MergeScheduleResult& schedule, int64_t& nextChannelId)