diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp index 731c02f..1089ee9 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredBoundaryRealization.cpp @@ -2,6 +2,7 @@ #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "DeferredBoundaryRealization.hpp" +#include "DeferredProjectionAnalysis.hpp" #include "DeferredResultRealization.hpp" #include "src/Accelerators/PIM/Common/IR/LoopUtils.hpp" #include "src/Accelerators/PIM/Common/IR/StaticIntGrid.hpp" @@ -428,9 +429,9 @@ static FailureOr insertProjectionFragment(Value fragment, Value specializ DeferredExchangePlan &exchange, bool grouped, DeferredEmissionContext &context) { Value shaped = fragment; if (leaf.form == DeferredLeafForm::GraphBatchProjection) { - SmallVector shape(leaf.leadingRankReduced ? leaf.reconstructedType.getShape() : leaf.reconstructedType.getShape().drop_front()); + RankedTensorType projectedType = getDeferredProjectedFragmentType(leaf); shaped = extractMixedSliceOrIdentity(context.rewriter, exchange.deferred.getLoc(), shaped, - RankedTensorType::get(shape, leaf.reconstructedType.getElementType()), + projectedType, lookupGeometry(geometry, geometryRow, runtimeLane, exchange.deferred, context, exchange.deferred.getLoc())); if (!shaped) return failure(); } @@ -494,8 +495,26 @@ static LogicalResult emitLeafCollectionUpdate(const EmitReceiveAssemblyRun &run, return emitCollectionUpdate(run.lanes, lane, laneCount, key, current, exchange.deferred, context, emit); } -static FailureOr transformAssemblySource(Value fragment, const DeferredInsertAssemblyEntryTemplate &entry, - DeferredExchangePlan &exchange, DeferredEmissionContext &context) { +static FailureOr transformAssemblySource( + Value fragment, const DeferredInsertAssemblyEntryTemplate &entry, + Value runtimeLane, const DeferredResultPlan &resultPlan, + DeferredExchangePlan &exchange, DeferredEmissionContext &context) { + if (entry.coordinate.leafIndex >= exchange.program.leaves.size() + || entry.coordinate.leafIndex >= resultPlan.innerGeometry.size()) + return failure(); + const DeferredProjectionLeafTemplate &leaf = + exchange.program.leaves[entry.coordinate.leafIndex]; + if (leaf.form == DeferredLeafForm::GraphBatchProjection) { + fragment = extractMixedSliceOrIdentity( + context.rewriter, exchange.deferred.getLoc(), fragment, + getDeferredProjectedFragmentType(leaf), + lookupGeometry(resultPlan.innerGeometry[entry.coordinate.leafIndex], + context.constants.getIndex(0), runtimeLane, + exchange.deferred, context, + exchange.deferred.getLoc())); + if (!fragment) + return failure(); + } switch (entry.sourceTransform) { case DeferredAssemblySourceTransform::Identity: return fragment.getType() == entry.sourceType ? FailureOr(fragment) : FailureOr(failure()); @@ -510,12 +529,15 @@ static FailureOr transformAssemblySource(Value fragment, const DeferredIn static FailureOr materializeLoopedLocalAssemblySource( RequirementFamily &requirement, const DeferredInsertAssemblyEntryTemplate &entry, - Value localOffset, DeferredExchangePlan &exchange, + Value localOffset, Value runtimeLane, + const DeferredResultPlan &resultPlan, DeferredExchangePlan &exchange, DeferredEmissionContext &context) { Value payload = requirement.producer->payload; auto payloadType = dyn_cast(payload.getType()); RankedTensorType sourceType = entry.sourceType; - if (entry.sourceTransform + if (exchange.program.leaves[entry.coordinate.leafIndex].form + != DeferredLeafForm::GraphBatchProjection + && entry.sourceTransform == DeferredAssemblySourceTransform::RemoveLeadingUnitDimension && payloadType && sourceType && payloadType.getRank() > sourceType.getRank() @@ -542,7 +564,8 @@ static FailureOr materializeLoopedLocalAssemblySource( exchange.deferred.getLoc()); if (failed(fragment)) return failure(); - return transformAssemblySource(*fragment, entry, exchange, context); + return transformAssemblySource( + *fragment, entry, runtimeLane, resultPlan, exchange, context); } static LogicalResult emitLoopedLocalCollectionUpdate( @@ -590,7 +613,8 @@ static LogicalResult emitLoopedLocalCollectionUpdate( RequirementFamily &requirement = *run.families.front()->requirement; FailureOr source = entry ? materializeLoopedLocalAssemblySource( - requirement, *entry, localOffset, exchange, context) + requirement, *entry, localOffset, runtimeLane, resultPlan, + exchange, context) : materializeSendPayload( requirement, localOffset, nullptr, context, loc); if (failed(source)) @@ -650,7 +674,8 @@ static LogicalResult emitInsertAssemblyUpdate(const EmitReceiveAssemblyRun &run, auto emit = [&](Value initial) { return emitReceiveAssembly(run, lane, laneCount, initial, context, [&](Value fragment, Value position, Value, Value runtimeLane, Value assembled) -> FailureOr { - auto shaped = transformAssemblySource(fragment, sourceEntry, exchange, context); + auto shaped = transformAssemblySource( + fragment, sourceEntry, runtimeLane, resultPlan, exchange, context); if (failed(shaped) || shaped->getType() != sourceEntry.sourceType) return failure(); return insertMixedSlice(context.rewriter, exchange.deferred.getLoc(), *shaped, assembled, lookupGeometry(resultPlan.assemblyGeometry, position, runtimeLane, exchange.deferred, context, exchange.deferred.getLoc())); @@ -775,7 +800,9 @@ static LogicalResult emitLocalCollectionUpdate(const EmitLocalCollectionRun &upd auto emit = [&](Value assembled) -> FailureOr { auto fragment = materialize(); if (failed(fragment)) return failure(); - auto shaped = transformAssemblySource(*fragment, entry, exchange, context); + auto shaped = transformAssemblySource( + *fragment, entry, lane ? lane : context.constants.getIndex(0), + resultPlan, exchange, context); if (failed(shaped) || shaped->getType() != entry.sourceType) return failure(); return insertMixedSlice(context.rewriter, exchange.deferred.getLoc(), *shaped, assembled, lookupGeometry(resultPlan.assemblyGeometry, context.constants.getIndex(update.collectionPosition), lane ? lane : context.constants.getIndex(0), diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp index d88ecbd..d29e072 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.cpp @@ -354,6 +354,16 @@ FailureOr evaluateDeferredIndex( return evaluate(value, environment, visiting); } +RankedTensorType getDeferredProjectedFragmentType( + const DeferredProjectionLeafTemplate &leaf) { + if (leaf.form == DeferredLeafForm::GraphBatchProjection + && !leaf.leadingRankReduced) + return RankedTensorType::get( + leaf.reconstructedType.getShape().drop_front(), + leaf.reconstructedType.getElementType()); + return leaf.reconstructedType; +} + FailureOr evaluateDeferredIndex( OpFoldResult value, const StaticIndexEnvironment &environment) { if (auto attr = dyn_cast(value)) diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp index 60b7491..aba4830 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp @@ -18,6 +18,9 @@ mlir::FailureOr evaluateDeferredIndex( mlir::FailureOr analyzeDeferredProgramTemplate( SpatDeferredCommunicationOp deferred); +mlir::RankedTensorType getDeferredProjectedFragmentType( + const DeferredProjectionLeafTemplate &leaf); + class DeferredLaneValueEvaluator { public: DeferredLaneValueEvaluator(const DeferredProgramTemplate &program, diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredResultRealization.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredResultRealization.cpp index 5e841e8..ac036d6 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredResultRealization.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredResultRealization.cpp @@ -72,11 +72,7 @@ static LogicalResult buildLeafCollections(DeferredExchangePlan &exchange, << ": publication fragment types differ or are unranked"; RankedTensorType normalized = fragmentType; if (leaf.form == DeferredLeafForm::GraphBatchProjection) - normalized = leaf.leadingRankReduced - ? leaf.reconstructedType - : RankedTensorType::get( - leaf.reconstructedType.getShape().drop_front(), - leaf.reconstructedType.getElementType()); + normalized = getDeferredProjectedFragmentType(leaf); bool direct = positionCount == 1 && normalized == leaf.reconstructedType; bool leading = leaf.reconstructedType.getRank() == normalized.getRank() + 1 && leaf.reconstructedType.getDimSize(0) == positionCount @@ -113,11 +109,20 @@ static LogicalResult buildLeafCollections(DeferredExchangePlan &exchange, } static bool supportsAssemblyTransform( - Type publicationType, const DeferredInsertAssemblyEntryTemplate &entry) { + Type publicationType, const DeferredInsertAssemblyEntryTemplate &entry, + const DeferredProjectionLeafTemplate &leaf) { auto publication = dyn_cast(publicationType); RankedTensorType source = entry.sourceType; if (!publication || !source) return false; + if (leaf.form == DeferredLeafForm::GraphBatchProjection) { + auto physical = dyn_cast(leaf.sourceRoot.getType()); + if (!physical || physical.getRank() != publication.getRank() + 1 + || physical.getElementType() != publication.getElementType() + || physical.getShape().drop_front() != publication.getShape()) + return false; + publication = getDeferredProjectedFragmentType(leaf); + } switch (entry.sourceTransform) { case DeferredAssemblySourceTransform::Identity: return publication == source; @@ -150,11 +155,16 @@ static LogicalResult buildInsertAssemblyCollection( for (RequirementFamily &requirement : exchange.requirements) { if (!(requirement.coordinate == entry.coordinate)) continue; + const DeferredProjectionLeafTemplate &leaf = + exchange.program.leaves[entry.coordinate.leafIndex]; if (!supportsAssemblyTransform( - requirement.publicationFragmentType, entry)) + requirement.publicationFragmentType, entry, leaf)) return exchange.deferred.emitOpError( "insert assembly source transform does not match publication type at entry ") - << position; + << position << ": publication " + << requirement.publicationFragmentType << ", assembly source " + << entry.sourceType << ", transform " + << static_cast(entry.sourceTransform); if (!collected.insert(&requirement).second) return exchange.deferred.emitOpError( "insert assembly requirement is owned by multiple entries at entry ") diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp index 93f80c4..5f7dc58 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp @@ -25,6 +25,7 @@ #include "src/Accelerators/PIM/Compiler/PimCompilerOptions.hpp" #include "src/Accelerators/PIM/Common/IR/AffineUtils.hpp" #include "src/Accelerators/PIM/Common/IR/ConstantUtils.hpp" +#include "src/Accelerators/PIM/Common/IR/ShapingUtils.hpp" #include "src/Support/TypeUtilities.hpp" namespace onnx_mlir { @@ -534,6 +535,27 @@ evaluateIndexLike(Value value, const DenseMap& bindings, std::op return evaluateAffineApply(affineApply, [&](Value operand) { return evaluateIndexLike(operand, bindings, lane, laneArg); }); + Operation* op = value.getDefiningOp(); + if (!op || !isPureIndexComputationOp(op)) + return failure(); + SmallVector operands; + Builder builder(op->getContext()); + for (Value operand : op->getOperands()) { + FailureOr folded = + evaluateIndexLike(operand, bindings, lane, laneArg); + if (failed(folded)) + return failure(); + operands.push_back(builder.getIntegerAttr(operand.getType(), *folded)); + } + SmallVector results; + if (failed(op->fold(operands, results)) || results.size() != 1) + return failure(); + if (auto attribute = dyn_cast(results.front())) + if (auto integer = dyn_cast(attribute)) + return integer.getInt(); + if (auto folded = dyn_cast(results.front())) + return evaluateIndexLike(folded, bindings, lane, laneArg); + return failure(); } diff --git a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp index 499698a..da9b38c 100644 --- a/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp +++ b/src/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp @@ -12,10 +12,13 @@ #include #include "PeftScheduler.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp" namespace onnx_mlir { namespace spatial { +using namespace mlir; + namespace { // Pressure means distinct weights exceed half the fleet's one-copy capacity. @@ -209,6 +212,79 @@ std::vector planCrossbarReservations(const ComputeGraph& graph, return reservations; } +using LanePublicationSignatures = llvm::SmallVector, 8>; + +FailureOr +buildLanePublicationSignatures(SpatComputeBatch batch, GraphBatchPublicationCache& publicationCache) { + LanePublicationSignatures signatures(batch.getLaneCount()); + for (auto [resultIndex, result] : llvm::enumerate(batch.getResults())) { + auto publicationMap = getGraphBatchPublicationMap(batch, resultIndex, publicationCache); + if (failed(publicationMap)) + return failure(); + + for (auto [useIndex, use] : llvm::enumerate(result.getUses())) { + auto blueprint = dyn_cast(use.getOwner()); + if (!blueprint || blueprint.getMode() != "fragment_assembly") + continue; + auto operandIndices = blueprint.getFragmentOperandIndices(); + auto sourceSlots = blueprint.getFragmentSourceSlots(); + auto sourceOffsets = blueprint.getFragmentSourceOffsets(); + auto fragmentStrides = blueprint.getFragmentStrides(); + auto outputType = dyn_cast(blueprint.getOutput().getType()); + if (!operandIndices || !sourceSlots || !sourceOffsets || !fragmentStrides + || !outputType || !outputType.hasStaticShape()) + return blueprint.emitOpError("PEFT publication compatibility requires complete static fragment metadata"), + failure(); + + llvm::ArrayRef fragmentOffsets = blueprint.getFragmentOffsets(); + llvm::ArrayRef fragmentSizes = blueprint.getFragmentSizes(); + int64_t rank = outputType.getRank(); + if (rank <= 0 || fragmentOffsets.size() != fragmentSizes.size() + || fragmentOffsets.size() != fragmentStrides->size() + || fragmentOffsets.size() != operandIndices->size() * rank + || sourceSlots->size() != operandIndices->size() + || sourceOffsets->size() != operandIndices->size()) + return blueprint.emitOpError("PEFT publication compatibility found inconsistent fragment metadata"), + failure(); + + llvm::SmallVector, 8> fragmentsByLane(batch.getLaneCount()); + for (auto [fragmentIndex, operandIndex] : llvm::enumerate(*operandIndices)) { + if (operandIndex != static_cast(use.getOperandNumber())) + continue; + int64_t slot = (*sourceSlots)[fragmentIndex]; + if (slot < 0 || slot >= static_cast((*publicationMap)->physicalSlotToGraphLane.size())) + return blueprint.emitOpError("PEFT publication fragment source slot is out of range"), failure(); + int64_t graphLane = (*publicationMap)->physicalSlotToGraphLane[slot]; + if (graphLane < 0 || graphLane >= batch.getLaneCount()) + return blueprint.emitOpError("PEFT publication fragment has no graph lane owner"), failure(); + fragmentsByLane[graphLane].push_back(fragmentIndex); + } + + for (auto [lane, fragments] : llvm::enumerate(fragmentsByLane)) { + if (fragments.empty()) + continue; + llvm::SmallVector& signature = signatures[lane]; + signature.push_back(resultIndex); + signature.push_back(useIndex); + signature.push_back(fragments.size()); + signature.push_back(rank); + size_t firstFragment = fragments.front(); + for (size_t fragmentIndex : fragments) { + signature.push_back((*sourceOffsets)[fragmentIndex] - (*sourceOffsets)[firstFragment]); + for (int64_t dim = 0; dim < rank; ++dim) { + size_t index = fragmentIndex * rank + dim; + size_t firstIndex = firstFragment * rank + dim; + signature.push_back(fragmentOffsets[index] - fragmentOffsets[firstIndex]); + signature.push_back(fragmentSizes[index]); + signature.push_back((*fragmentStrides)[index]); + } + } + } + } + } + return signatures; +} + } // namespace Time getPeftTransferTime(Time transferCost, size_t sourceProcessor, size_t targetProcessor, size_t processorCount) { @@ -450,6 +526,28 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu // 5. Check if equal schedule in two level llvm::DenseMap> equivalentClass; + GraphBatchPublicationCache publicationCache; + llvm::DenseMap publicationSignatures; + llvm::DenseSet invalidPublicationSignatures; + auto haveCompatiblePublications = [&](const ComputeInstance& lhs, const ComputeInstance& rhs) { + auto batch = dyn_cast(lhs.op); + if (!batch || batch.getNumResults() == 0) + return true; + auto signatures = publicationSignatures.find(lhs.op); + if (signatures == publicationSignatures.end() && !invalidPublicationSignatures.contains(lhs.op)) { + auto built = buildLanePublicationSignatures(batch, publicationCache); + if (failed(built)) + invalidPublicationSignatures.insert(lhs.op); + else + signatures = publicationSignatures.try_emplace(lhs.op, std::move(*built)).first; + } + if (invalidPublicationSignatures.contains(lhs.op)) + return false; + for (uint32_t lane = 0; lane < lhs.laneCount; ++lane) + if (signatures->second[lhs.laneStart + lane] != signatures->second[rhs.laneStart + lane]) + return false; + return true; + }; for (size_t currentProcessor = 0; currentProcessor < processorCount - 1; ++currentProcessor) { for (size_t controlProcessor = currentProcessor; controlProcessor < processorCount; ++controlProcessor) { if (tasksByProcessor[currentProcessor].size() != tasksByProcessor[controlProcessor].size()) @@ -462,7 +560,8 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu const ComputeInstance currentComputeInstance = graph.nodes[currentTask].instance; const ComputeInstance controlComputeInstance = graph.nodes[controlTask].instance; if (currentComputeInstance.op != controlComputeInstance.op - || currentComputeInstance.laneCount != controlComputeInstance.laneCount) { + || currentComputeInstance.laneCount != controlComputeInstance.laneCount + || !haveCompatiblePublications(currentComputeInstance, controlComputeInstance)) { equalSchedule = false; break; }