diff --git a/.agents/invariants/PIPELINE_SCHEDULING_INVARIANT.md b/.agents/invariants/PIPELINE_SCHEDULING_INVARIANT.md index 506f1a3..da04e03 100644 --- a/.agents/invariants/PIPELINE_SCHEDULING_INVARIANT.md +++ b/.agents/invariants/PIPELINE_SCHEDULING_INVARIANT.md @@ -30,8 +30,9 @@ they must not infer it from a physical core number after placement. ## Ownership -Physical PEFT supplies the initial placement. Stage partitioning is the first -phase of pipeline scheduling and owns stage adjacency. It must construct a valid +Latency PEFT and pipeline scheduling independently consume the shared compute +graph and target cost model. Stage partitioning is the first phase of pipeline +scheduling and owns stage adjacency. It must construct a valid operation-level partition before physical-core packing. Operations split for physical capacity retain one shared stage identity. Repacking may move work only within its assigned stage. Deferred-transfer planning and diff --git a/src/PIM/Common/IR/ShapeUtils.cpp b/src/PIM/Common/IR/ShapeUtils.cpp index 14da6a1..5e0fd55 100644 --- a/src/PIM/Common/IR/ShapeUtils.cpp +++ b/src/PIM/Common/IR/ShapeUtils.cpp @@ -15,6 +15,25 @@ llvm::SmallVector computeRowMajorStrides(llvm::ArrayRef shape) return strides; } +TransposeRun getTransposeRun(llvm::ArrayRef sourceShape, + llvm::ArrayRef permutation) { + auto sourceStrides = computeRowMajorStrides(sourceShape); + TransposeRun run{sourceShape.size(), 1, 1}; + for (size_t dim = permutation.size(); dim-- > 0;) { + int64_t sourceDim = permutation[dim]; + if (sourceShape[sourceDim] == 1) + continue; + int64_t stride = sourceStrides[sourceDim]; + if (run.start != sourceShape.size() && stride != run.elements * run.sourceStride) + break; + run.start = dim; + if (run.elements == 1) + run.sourceStride = stride; + run.elements *= sourceShape[sourceDim]; + } + return run; +} + llvm::SmallVector delinearizeIndex(int64_t linearIndex, llvm::ArrayRef shape, llvm::ArrayRef strides) { llvm::SmallVector indices(shape.size(), 0); diff --git a/src/PIM/Common/IR/ShapeUtils.hpp b/src/PIM/Common/IR/ShapeUtils.hpp index 50ea76c..ce1be65 100644 --- a/src/PIM/Common/IR/ShapeUtils.hpp +++ b/src/PIM/Common/IR/ShapeUtils.hpp @@ -21,6 +21,15 @@ using CoreId = size_t; llvm::SmallVector computeRowMajorStrides(llvm::ArrayRef shape); +struct TransposeRun { + size_t start; + int64_t elements; + int64_t sourceStride; +}; + +TransposeRun getTransposeRun(llvm::ArrayRef sourceShape, + llvm::ArrayRef permutation); + llvm::SmallVector delinearizeIndex(int64_t linearIndex, llvm::ArrayRef shape, llvm::ArrayRef strides); diff --git a/src/PIM/Compiler/PimCompilerUtils.cpp b/src/PIM/Compiler/PimCompilerUtils.cpp index 5452f52..d12b999 100644 --- a/src/PIM/Compiler/PimCompilerUtils.cpp +++ b/src/PIM/Compiler/PimCompilerUtils.cpp @@ -78,6 +78,8 @@ spatial::SchedulingTarget getDefaultPimSchedulingTarget() { target.processorCount = static_cast(coresCount.getValue()); target.residentWeightCapacity = crossbarCountInCore.getValue(); target.matrixRows = crossbarSize.getValue(); + // The oracle charges eight input reads for each matrix operation. + target.matrixInputReadMultiplier = 8; target.matrixColumns = crossbarSize.getValue(); target.synchronizationRegisterCount = kPimEventRegisterCount; diff --git a/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/PimBufferizationPass.cpp b/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/PimBufferizationPass.cpp index ac16b52..2a33276 100644 --- a/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/PimBufferizationPass.cpp +++ b/src/PIM/Dialect/Pim/Passes/Transforms/Bufferization/PimBufferizationPass.cpp @@ -275,9 +275,12 @@ static void forwardSingleConsumerReceiveCopies(func::FuncOp funcOp) { ? receive.getOutputBuffer().getDefiningOp() : memref::AllocOp(); auto sourceOffset = getConstantIntValue(copy.getSourceOffset()); + auto receiveOffset = receive + ? getConstantIntValue(receive.getOutputOffset()) : std::nullopt; if (!receive || !outputAlloc || !receive.getOutput().hasOneUse() || !outputAlloc.getResult().hasOneUse() || !sourceOffset || *sourceOffset != 0 + || !receiveOffset || *receiveOffset != 0 || copy.getSize() != receive.getSize() || receive->getBlock() != copy->getBlock()) continue; @@ -285,7 +288,9 @@ static void forwardSingleConsumerReceiveCopies(func::FuncOp funcOp) { bool canSinkReceive = true; for (Operation* between = receive->getNextNode(); between != copy; between = between->getNextNode()) { - if (!isMemoryEffectFree(between)) { + // The receive writes a private allocation; intervening loads cannot + // observe it, and the destination is not written until the copy. + if (!isMemoryEffectFree(between) && !isa(between)) { canSinkReceive = false; break; } @@ -594,7 +599,6 @@ static LogicalResult verifyPimCopyAddressSpaces(ModuleOp moduleOp) { } static LogicalResult normalizePimMemory(ModuleOp moduleOp, func::FuncOp funcOp) { - forwardSingleConsumerReceiveCopies(funcOp); forwardSingleConsumerContiguousInputCopies(funcOp); forwardSingleConsumerPimOutputCopies(funcOp); @@ -676,6 +680,7 @@ static LogicalResult normalizePimMemory(ModuleOp moduleOp, func::FuncOp funcOp) moduleOp.emitError("failed to normalize Pim copy contiguity during bufferization"); return failure(); } + forwardSingleConsumerReceiveCopies(funcOp); annotateWeightsMemrefs(moduleOp, funcOp); dumpModule(moduleOp, "pim1_buff"); return success(); diff --git a/src/PIM/Dialect/Pim/Passes/Transforms/InstructionSelection/InstructionSelectionPass.cpp b/src/PIM/Dialect/Pim/Passes/Transforms/InstructionSelection/InstructionSelectionPass.cpp index b62e3bf..e716020 100644 --- a/src/PIM/Dialect/Pim/Passes/Transforms/InstructionSelection/InstructionSelectionPass.cpp +++ b/src/PIM/Dialect/Pim/Passes/Transforms/InstructionSelection/InstructionSelectionPass.cpp @@ -84,20 +84,7 @@ struct LowerTransposePattern final : OpRewritePattern { for (int64_t sourceDim : permutation) destinationShape.push_back(sourceShape[sourceDim]); - size_t runStart = rank; - int64_t runLength = 1; - int64_t sourceStride = 1; - for (size_t destinationDim = rank; destinationDim-- > 0;) { - if (destinationShape[destinationDim] == 1) - continue; - int64_t candidateStride = sourceStrides[permutation[destinationDim]]; - if (runStart != rank && candidateStride != runLength * sourceStride) - break; - runStart = destinationDim; - if (runLength == 1) - sourceStride = candidateStride; - runLength *= destinationShape[destinationDim]; - } + auto [runStart, runLength, sourceStride] = getTransposeRun(sourceShape, permutation); assert(runStart != rank && "non-storage-preserving transpose must have a non-unit dimension"); int64_t outerCount = sourceType.getNumElements() / runLength; diff --git a/src/PIM/Dialect/Spatial/CMakeLists.txt b/src/PIM/Dialect/Spatial/CMakeLists.txt index d46483f..e1389e2 100644 --- a/src/PIM/Dialect/Spatial/CMakeLists.txt +++ b/src/PIM/Dialect/Spatial/CMakeLists.txt @@ -34,8 +34,11 @@ add_pim_library(SpatialOps Passes/Transforms/MergeComputeNodes/ScheduledComputeVerification.cpp Passes/Transforms/MergeComputeNodes/SpatialDataflowCsvExporter.cpp Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp + Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.cpp Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp Passes/Transforms/MergeComputeNodes/Scheduling/PipelineScheduling.cpp + Passes/Transforms/MergeComputeNodes/Scheduling/PipelineWorkloadPreparation.cpp + Passes/Transforms/MergeComputeNodes/Scheduling/ScheduleEvaluation.cpp Passes/Transforms/TrivialGraphComputeMergePass.cpp EXCLUDE_FROM_OM_LIBS diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp index 712110c..cee807c 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationPlanning.cpp @@ -417,7 +417,8 @@ LogicalResult prepareSingleCpuInput(OpBuilder &, Location loc, Value input, Bloc && producer->instance.laneCount == static_cast(batch.getLaneCount())); if (auto available = availableValues.find(key); - hasCompleteValue && available != availableValues.end()) { + hasCompleteValue && available != availableValues.end() + && (!batch || available->second.hasGraphPhysicalOrder)) { plan.availableValue = available->second.payload; return success(); } diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp index aa19607..d1abb4e 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp @@ -387,8 +387,10 @@ FailureOr scheduleDeferredCommunication(func::FuncOp funcOp.emitOpError("global communication scheduler made no progress before IR mutation"); unsigned reported = 0; for (const TransferGroup& group : groups) { - if (group.scheduled || reported++ == 8) + if (group.scheduled) continue; + if (reported++ == 8) + break; diagnostic << "; blocked parent " << group.exchange->exchangeId; auto dependency = llvm::find_if(group.dependencies, [&](StreamThreshold item) { return streams[item.stream].completedStep < item.completedStep; diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp index 4b990a1..519f8af 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.cpp @@ -1,5 +1,6 @@ #include "ScheduledComputeMaterialization.hpp" #include "DeferredCommunicationPlanning.hpp" +#include "DeferredProjectionAnalysis.hpp" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" @@ -129,6 +130,36 @@ static void cloneComputeBody(OpBuilder &builder, Block &source, IRMapping &mappe yieldedValues.push_back(mapper.lookup(output)); } +} // namespace + +Value insertLocalFragment(OpBuilder &builder, Location loc, Value source, + Value destination, ArrayRef offsets, + ArrayRef sizes, + ArrayRef strides) { + auto type = cast(destination.getType()); + auto sourceType = cast(source.getType()); + bool complete = type.hasStaticShape() && sourceType.hasStaticShape() + && type.getRank() == sourceType.getRank() + 1 + && type.getDimSize(0) == 1 + && type.getElementType() == sourceType.getElementType() + && type.getEncoding() == sourceType.getEncoding() + && llvm::equal(type.getShape().drop_front(), sourceType.getShape()) + && offsets.size() == static_cast(type.getRank()) + && sizes.size() == offsets.size() && strides.size() == offsets.size(); + if (complete) + for (auto [dimension, offset, size, stride] : + llvm::zip_equal(type.getShape(), offsets, sizes, strides)) + complete &= getConstantIntValue(offset) == 0 + && getConstantIntValue(size) == dimension + && getConstantIntValue(stride) == 1; + if (complete) + return *addLeadingUnitTensorDimension(builder, loc, source); + return tensor::InsertSliceOp::create( + builder, loc, source, destination, offsets, sizes, strides); +} + +namespace { + static LogicalResult materializeResultfulBatchRun( PatternRewriter &rewriter, SpatComputeBatch batch, const ScheduledInstanceRun &run, ValueRange scheduledWeights, @@ -232,14 +263,11 @@ static LogicalResult materializeResultfulBatchRun( size_t resultIndex = outputIndexByArg.lookup(oldDest); SmallVector offsets = remapMixedOffsets(insert.getMixedOffsets(), mapper); offsets.front() = localLane; - current[resultIndex] = tensor::InsertSliceOp::create(builder, - insert.getLoc(), - mapper.lookup(insert.getSource()), - current[resultIndex], - offsets, - remapMixedOffsets(insert.getMixedSizes(), mapper), - remapMixedOffsets(insert.getMixedStrides(), mapper)) - .getResult(); + current[resultIndex] = insertLocalFragment( + builder, insert.getLoc(), mapper.lookup(insert.getSource()), + current[resultIndex], offsets, + remapMixedOffsets(insert.getMixedSizes(), mapper), + remapMixedOffsets(insert.getMixedStrides(), mapper)); } llvm::append_range(yielded, current); return success(); @@ -264,6 +292,7 @@ static LogicalResult materializeSingleCpuPeftClass( Block *block = createScheduledComputeBlock( rewriter, scheduled, instances.front().op->getLoc()); DenseMap availableValues; + GraphBatchPublicationCache publicationCache; DenseMap publicationIndices; for (const ScheduledPublication &publication : record.publications) publicationIndices[publication.producer] = publication.scheduledResultIndex; @@ -326,13 +355,23 @@ static LogicalResult materializeSingleCpuPeftClass( int64_t runLaneStart = instance.laneStart; const ComputeInstance &last = run.instances.back(); int64_t runLaneCount = last.laneStart + last.laneCount - runLaneStart; + bool hasGraphPhysicalOrder = false; + if (auto batch = dyn_cast(instance.op); + batch && runLaneStart == 0 && runLaneCount == batch.getLaneCount()) { + auto publication = getGraphBatchPublicationMap(batch, resultIndex, publicationCache); + if (failed(publication)) + return failure(); + hasGraphPhysicalOrder = llvm::all_of( + llvm::enumerate((*publication)->graphLaneToPhysicalSlot), + [](auto entry) { return entry.index() == static_cast(entry.value()); }); + } availableValues[{{instance.op, static_cast(runLaneStart), static_cast(runLaneCount)}, resultIndex}] = - {value, 0, runLaneCount}; + {value, 0, runLaneCount, hasGraphPhysicalOrder}; for (const ComputeInstance &member : run.instances) { ProducerValueKey key {member, resultIndex}; int64_t payloadLaneStart = member.laneStart - runLaneStart; - availableValues[key] = {value, payloadLaneStart, member.laneCount}; + availableValues[key] = {value, payloadLaneStart, member.laneCount, hasGraphPhysicalOrder}; Value published; if (auto it = publicationIndices.find(key); it != publicationIndices.end()) { @@ -654,14 +693,11 @@ static LogicalResult materializeMultiCpuPeftClass( size_t resultIndex = outputIndexByArg.lookup(oldDest); SmallVector offsets = remapMixedOffsets(insert.getMixedOffsets(), mapper); offsets.front() = innerLane; - current[resultIndex] = tensor::InsertSliceOp::create(builder, - insert.getLoc(), - mapper.lookup(insert.getSource()), - current[resultIndex], - offsets, - remapMixedOffsets(insert.getMixedSizes(), mapper), - remapMixedOffsets(insert.getMixedStrides(), mapper)) - .getResult(); + current[resultIndex] = insertLocalFragment( + builder, insert.getLoc(), mapper.lookup(insert.getSource()), + current[resultIndex], offsets, + remapMixedOffsets(insert.getMixedSizes(), mapper), + remapMixedOffsets(insert.getMixedStrides(), mapper)); } llvm::append_range(yielded, current); return success(); @@ -706,11 +742,9 @@ FailureOr materializeScheduledCompute(func::FuncOp funcOp, const MergeScheduleResult &schedule, PatternRewriter &rewriter) { - DenseMap graphIds; int64_t nextGraphId = 0; for (Operation &op : funcOp.getOps()) if (isa(op)) { - graphIds[&op] = nextGraphId; op.setAttr("scheduled.graph_id", rewriter.getI64IntegerAttr(nextGraphId++)); } @@ -726,13 +760,12 @@ materializeScheduledCompute(func::FuncOp funcOp, } for (auto &entry : peftClassPlans) { PeftClassPlan &peftClassPlan = entry.second; - llvm::sort(peftClassPlan.cpus); + // Keep class lanes in graph encounter order so physical placement does not + // turn regular source lanes into lookup tables. Core execution uses slots. for (size_t cpu : peftClassPlan.cpus) llvm::sort(peftClassPlan.instancesByCpu[cpu], [&](const ComputeInstance &lhs, const ComputeInstance &rhs) { - return std::tie(graphIds.find(lhs.op)->second, - schedule.computeToCpuSlotMap.find(lhs)->second) < - std::tie(graphIds.find(rhs.op)->second, - schedule.computeToCpuSlotMap.find(rhs)->second); + return schedule.computeToCpuSlotMap.lookup(lhs) + < schedule.computeToCpuSlotMap.lookup(rhs); }); if (failed(verifyPeftClassPlan(funcOp.getOperation(), peftClassPlan, schedule))) return failure(); diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.hpp index 61309d5..1c97cba 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.hpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeMaterialization.hpp @@ -17,6 +17,11 @@ struct ScheduledComputeMaterializationResult { std::vector processorStages; }; +Value insertLocalFragment(OpBuilder &builder, Location loc, Value source, + Value destination, ArrayRef offsets, + ArrayRef sizes, + ArrayRef strides); + FailureOr getBatchFragmentSpec(SpatComputeBatch batch, unsigned resultIndex, diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp index e5d330e..357227b 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputePlan.hpp @@ -94,6 +94,7 @@ struct MaterializedProducerRef { Value payload; int64_t payloadLaneStart = 0; int64_t payloadLaneCount = 1; + bool hasGraphPhysicalOrder = false; }; struct MaterializedStepValue { diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp index 8cdf867..58adf45 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputePlanning.cpp @@ -108,9 +108,15 @@ LogicalResult verifyPeftClassPlan(Operation *diagnosticAnchor, if (it == peftClassPlan.instancesByCpu.end()) return diagnosticAnchor->emitOpError("PEFT materialization class is missing a per-CPU schedule"); schedules.push_back(&it->second); - for (const ComputeInstance &instance : it->second) - if (!schedule.computeToCpuSlotMap.count(instance)) - return diagnosticAnchor->emitOpError("PEFT materialization class references a compute instance without a scheduler position"); + for (auto [slot, instance] : llvm::enumerate(it->second)) { + auto position = schedule.computeToCpuSlotMap.find(instance); + auto processor = schedule.computeToCpuMap.find(instance); + if (position == schedule.computeToCpuSlotMap.end() + || processor == schedule.computeToCpuMap.end() + || processor->second != cpu || position->second != slot) + return diagnosticAnchor->emitOpError( + "scheduled materialization must preserve each core's scheduler slot order"); + } } if (peftClassPlan.cpus.size() == 1) diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp index 4930a40..4486450 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledComputeReport.cpp @@ -112,6 +112,18 @@ void dumpScheduledComputeReport(ModuleOp moduleOp, func::FuncOp funcOp, const Me << " largest run: " << largestRun << "\n" << " instances compacted: " << instancesCompacted << "\n" << " compatible runs rejected: 0\n"; + const ScheduleCostSummary &costs = schedule.costs; + os << " estimated total serial work (ns): " << costs.totalWorkNs << "\n" + << " estimated maximum core work (ns): " << costs.maximumCoreWorkNs << "\n" + << " estimated dependency critical path (ns): " << costs.criticalPathNs << "\n" + << " summed edge communication time (ns): " << costs.communicationNs << "\n" + << " estimated single-inference completion (ns): " << costs.completionNs << "\n"; + for (auto [stage, span] : llvm::enumerate(costs.stageSpanNs)) + os << " estimated stage " << stage << " span (ns, excludes boundary synchronization/contention): " + << span << "\n"; + for (auto [stage, demand] : llvm::enumerate(costs.stageSerialDemandNs)) + os << " estimated stage " << stage << " maximum core compute/incident-transfer demand (ns): " + << demand << "\n"; llvm::MapVector> stageCores; for (auto [core, stage] : llvm::enumerate(schedule.processorStages)) stageCores[stage].push_back(core); diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp index 18ba528..b0a633d 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/ScheduledSpatialPasses.cpp @@ -19,6 +19,7 @@ namespace { static bool hasValidTarget(const SchedulingTarget& target) { return target.processorCount != 0 && target.residentWeightCapacity != 0 && target.transferWidthBytes != 0 + && target.matrixInputReadMultiplier != 0 && target.interProcessorLatencyNs.size() == target.processorCount * target.processorCount && (target.processorCount == 1 || target.averageInterProcessorLatencyNs != 0); } @@ -73,7 +74,12 @@ struct ScheduleAndRealizeSpatialPass final ComputeGraph scheduledGraph; MergeScheduleResult schedule; for (;;) { - scheduledGraph = buildComputeGraph(entryFunc, target); + auto graph = buildComputeGraph(entryFunc, target); + if (failed(graph)) { + signalPassFailure(); + return; + } + scheduledGraph = std::move(*graph); std::string pipelineError; if (pipelineStages > 1) { FailureOr preparation = @@ -87,10 +93,18 @@ struct ScheduleAndRealizeSpatialPass final if (*preparation == PipelineWorkloadPreparation::Changed) continue; } - MergeSchedulingAnalysis analysis( - scheduledGraph, target, entryFunc.getContext()); - schedule = std::move(analysis.getResult()); - if (succeeded(applyPipelineScheduling( + if (pipelineStages == 1) { + MergeSchedulingAnalysis analysis( + scheduledGraph, target, entryFunc.getContext(), pipelineError); + if (failed(analysis.getResult())) { + moduleOp.emitError() << pipelineError; + signalPassFailure(); + return; + } + schedule = std::move(*analysis.getResult()); + break; + } + if (succeeded(runPipelineScheduler( scheduledGraph, schedule, pipelineStages, target, pipelineError))) break; moduleOp.emitError() << pipelineError; diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp index 9beb2f3..6fc9e90 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp @@ -1,5 +1,6 @@ #include "mlir/Dialect/Affine/IR/AffineOps.h" #include "mlir/Dialect/Arith/IR/Arith.h" +#include "mlir/Dialect/Linalg/IR/Linalg.h" #include "mlir/Dialect/SCF/IR/SCF.h" #include "mlir/Dialect/Tensor/IR/Tensor.h" #include "mlir/IR/BuiltinAttributes.h" @@ -7,8 +8,12 @@ #include "mlir/IR/Operation.h" #include "mlir/IR/Unit.h" #include "mlir/IR/Value.h" +#include "mlir/Transforms/RegionUtils.h" #include "llvm/ADT/STLExtras.h" +#include "llvm/ADT/SmallBitVector.h" +#include "llvm/ADT/DenseSet.h" +#include "llvm/ADT/SetVector.h" #include "llvm/Support/Casting.h" #include @@ -21,6 +26,7 @@ #include "ComputeGraph.hpp" #include "ComputeInstanceUtils.hpp" +#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.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" @@ -43,18 +49,11 @@ TransferCost addTransferCosts(const TransferCost& lhs, checkedAdd(lhs.networkFlits, rhs.networkFlits)}; } -TransferCost scaleTransferCost(const TransferCost& cost, - Cost numerator, - Cost denominator = 1) { - assert(denominator > 0 && "transfer cost denominator must be positive"); - return {checkedMultiply(cost.fixed, numerator) / denominator, - checkedMultiply(cost.networkFlits, numerator) / denominator}; -} - struct SchedulerCostModel { static constexpr Cost kFallbackOperationCost = 1; const SchedulingTarget& target; + ComputeCostKind kind = ComputeCostKind::Execution; static Cost ceilDiv(Cost numerator, Cost denominator) { assert(denominator > 0 && "denominator must be positive"); @@ -110,7 +109,44 @@ struct SchedulerCostModel { } Cost getTensorMoveCost(Type type) const { - return getVectorReadWriteCost(getByteSize(type), 0); + Cost bytes = getByteSize(type, getComputeBitwidth(type)); + return getVectorReadWriteCost(bytes, bytes); + } + + Cost getTransposeCost(linalg::TransposeOp transpose) const { + if (kind == ComputeCostKind::VolumePriority) + return getTensorMoveCost(transpose.getInput().getType()); + auto type = cast(transpose.getInput().getType()); + if (!type.hasStaticShape()) + return getTensorMoveCost(type); + if (type.getNumElements() == 0) + return 0; + auto run = getTransposeRun(type.getShape(), transpose.getPermutation()); + Cost bytes = ceilDiv(checkedMultiply(static_cast(run.elements), getComputeBitwidth(type)), Cost{8}); + Cost cost = getVectorReadWriteCost(bytes, bytes); + if (run.sourceStride != 1) + cost = checkedAdd(cost, getVectorComputeCost(run.elements)); + return checkedMultiply(static_cast(type.getNumElements()) / run.elements, cost); + } + + Cost getConcatCost(ValueRange inputs, int64_t axis) const { + Cost cost = 0; + for (Value input : inputs) { + auto type = dyn_cast(input.getType()); + if (!type || !type.hasStaticShape()) + return kFallbackOperationCost; + Cost blocks = 1; + Cost elements = 1; + for (auto [dim, size] : llvm::enumerate(type.getShape())) { + if (static_cast(dim) < axis) + blocks = checkedMultiply(blocks, static_cast(size)); + else + elements = checkedMultiply(elements, static_cast(size)); + } + Cost bytes = ceilDiv(checkedMultiply(elements, getComputeBitwidth(type)), Cost{8}); + cost = checkedAdd(cost, checkedMultiply(blocks, getVectorReadWriteCost(bytes, bytes))); + } + return cost; } TransferCost getTransferCostFromBytes(Cost bytes) const { @@ -176,7 +212,7 @@ struct SchedulerCostModel { Cost inputBitwidth = getComputeBitwidth(inputType); Cost inputBytes = checkedMultiply(target.matrixRows, ceilDiv(inputBitwidth, static_cast(8))); - inputBytes = checkedMultiply(inputBytes, static_cast(8)); + inputBytes = checkedMultiply(inputBytes, target.matrixInputReadMultiplier); Cost outputBytes = getByteSize(outputType, getComputeBitwidth(outputType)); return checkedAdd(getVectorReadWriteCost(inputBytes, outputBytes), getMatrixComputeLatency(inputBitwidth)); } @@ -250,6 +286,10 @@ Cost getOperationCost(Operation& op, const SchedulerCostModel& costModel) { } if (auto insert = dyn_cast(&op)) return costModel.getTensorMoveCost(insert.getSource().getType()); + if (auto transpose = dyn_cast(&op)) + return costModel.getTransposeCost(transpose); + if (auto concat = dyn_cast(&op)) + return costModel.getConcatCost(concat.getInputs(), concat.getAxis()); Cost nestedCost = 0; for (Region& region : op.getRegions()) @@ -275,6 +315,30 @@ Cost getComputeBodyCost(Region& body, const SchedulerCostModel& costModel) { return getRegionCost(body, costModel); } +Cost getBlueprintAssemblyCost(const ComputeInstance& instance, + const SchedulerCostModel& costModel) { + Cost cost = 0; + for (Value input : getComputeInstanceInputs(instance)) { + auto blueprint = input.getDefiningOp(); + if (!blueprint || !isFragmentAssembly(blueprint.getMode())) + continue; + auto type = cast(input.getType()); + int64_t rank = type.getRank(); + ArrayRef sizes = blueprint.getFragmentSizes(); + if (sizes.empty()) + continue; + for (size_t fragment = 0; fragment < sizes.size() / rank; ++fragment) { + Cost elements = 1; + for (int64_t size : sizes.slice(fragment * rank, rank)) + elements = checkedMultiply(elements, static_cast(size)); + Cost bytes = SchedulerCostModel::ceilDiv( + checkedMultiply(elements, costModel.getComputeBitwidth(type)), Cost{8}); + cost = checkedAdd(cost, costModel.getVectorReadWriteCost(bytes, bytes)); + } + } + return cost; +} + uint64_t countOperationInstances(Operation& op) { if (auto loop = dyn_cast(&op)) { std::optional tripCount = getStaticTripCount(loop); @@ -320,12 +384,6 @@ static FailureOr evaluateIndexLike(OpFoldResult value, std::optional lane, Value laneArg); -bool isProjectedBatchOffset(OpFoldResult offset, Value laneArg) { - DenseMap bindings; - return succeeded(evaluateIndexLike(offset, bindings, /*lane=*/0, laneArg)) - && succeeded(evaluateIndexLike(offset, bindings, /*lane=*/1, laneArg)); -} - std::optional getConstantExtractLane(tensor::ExtractSliceOp extract) { if (extract.getMixedOffsets().empty()) return std::nullopt; @@ -348,65 +406,41 @@ std::optional getConstantExtractLane(tensor::ExtractSliceOp extract) { return std::nullopt; } -std::optional getBatchProjectedInputTransferCost( - SpatComputeBatch batch, Value input, - const SchedulerCostModel& costModel) { - auto inputIt = llvm::find(batch.getInputs(), input); - if (inputIt == batch.getInputs().end()) - return std::nullopt; - - size_t inputIndex = std::distance(batch.getInputs().begin(), inputIt); - std::optional inputArg = batch.getInputArgument(inputIndex); - std::optional laneArg = batch.getLaneArgument(); - if (!inputArg || !laneArg) - return std::nullopt; - - TransferCost projectedCost; - for (Operation* user : inputArg->getUsers()) { - auto extract = dyn_cast(user); - if (!extract || extract.getSource() != *inputArg) - return std::nullopt; - if (extract.getMixedOffsets().empty() || !isProjectedBatchOffset(extract.getMixedOffsets().front(), *laneArg)) - return std::nullopt; - - auto resultType = dyn_cast(extract.getResult().getType()); - if (!resultType || !resultType.hasStaticShape()) - return std::nullopt; - projectedCost = addTransferCosts( - projectedCost, - costModel.getTransferCostFromBytes( - costModel.getByteSize(resultType, costModel.getComputeBitwidth(resultType)))); - } - - if (projectedCost.fixed == 0 && projectedCost.networkFlits == 0) - return std::nullopt; - return projectedCost; -} - -static std::optional> -collectProjectedProducerValueRefs(SpatComputeBatch producer, - Value input, - const ComputeInstance& consumerInstance, - size_t processorCount) { +static FailureOr +collectProjectedEdges(SpatComputeBatch producer, Value input, + const ComputeInstance& consumerInstance, + size_t processorCount, size_t targetIndex, + const ComputeGraph& graph, + const SchedulerCostModel& costModel, + GraphBatchPublicationCache& publicationCache, + SmallVectorImpl& rawEdges) { auto consumer = dyn_cast(consumerInstance.op); if (!consumer) - return std::nullopt; + return false; auto inputIt = llvm::find(consumer.getInputs(), input); if (inputIt == consumer.getInputs().end()) - return std::nullopt; + return false; size_t inputIndex = std::distance(consumer.getInputs().begin(), inputIt); std::optional inputArg = consumer.getInputArgument(inputIndex); std::optional laneArg = consumer.getLaneArgument(); if (!inputArg || !laneArg) - return std::nullopt; + return false; + auto publication = getGraphBatchPublicationMap( + producer, cast(input).getResultNumber(), publicationCache); + if (failed(publication)) + return failure(); - SmallVector producers; + SmallVector edges; + llvm::SmallBitVector selectedSlots(producer.getLaneCount()); + TransferCost fragmentCost = costModel.getTransferCostFromBytes( + costModel.getByteSize((*publication)->publicationFragmentType, + costModel.getComputeBitwidth((*publication)->publicationFragmentType))); DenseMap bindings; for (Operation* user : inputArg->getUsers()) { auto extract = dyn_cast(user); if (!extract || extract.getSource() != *inputArg || extract.getMixedOffsets().empty() || extract.getMixedSizes().empty() || extract.getMixedStrides().empty()) - return std::nullopt; + return false; for (uint32_t lane = consumerInstance.laneStart; lane < consumerInstance.laneStart + consumerInstance.laneCount; ++lane) { @@ -417,142 +451,219 @@ collectProjectedProducerValueRefs(SpatComputeBatch producer, FailureOr stride = evaluateIndexLike(extract.getMixedStrides().front(), bindings, lane, *laneArg); if (failed(offset) || failed(size) || failed(stride) || *offset < 0 || *size <= 0 || *stride <= 0) - return std::nullopt; + return false; for (int64_t index = 0; index < *size; ++index) { - int64_t producerLane = *offset + index * *stride; - if (producerLane < 0 || producerLane >= producer.getLaneCount()) - return std::nullopt; + int64_t slot = *offset + index * *stride; + if (slot < 0 || slot >= producer.getLaneCount()) + return false; + if (selectedSlots.test(slot)) + continue; + selectedSlots.set(slot); + int64_t producerLane = (*publication)->physicalSlotToGraphLane[slot]; ComputeInstance instance = getBatchChunkForLane( producer, static_cast(producerLane), processorCount); - if (llvm::none_of(producers, [&](const ProducerValueRef& ref) { return ref.instance == instance; })) - producers.push_back({instance, 0}); + auto source = graph.instanceToIndex.find(instance); + if (source == graph.instanceToIndex.end()) + return producer.emitOpError("projected producer has no compute instance"), failure(); + edges.push_back({source->second, targetIndex, fragmentCost}); } } } - return producers; + llvm::append_range(rawEdges, edges); + return true; } -TransferCost getInputTransferCost(const ComputeInstance& consumerInstance, - Value input, +std::optional getProjectedInputTransferCost( + Value input, const ComputeInstance& instance, + const SchedulerCostModel& costModel) { + auto batch = dyn_cast(instance.op); + if (!batch) + return std::nullopt; + auto found = llvm::find(batch.getInputs(), input); + if (found == batch.getInputs().end()) + return std::nullopt; + auto argument = batch.getInputArgument(found - batch.getInputs().begin()); + if (!argument) + return std::nullopt; + TransferCost cost; + for (Operation* user : argument->getUsers()) { + auto extract = dyn_cast(user); + if (!extract || extract.getSource() != *argument) + return std::nullopt; + Cost count = instance.laneCount; + for (Operation* parent = extract->getParentOp(); parent != instance.op; + parent = parent->getParentOp()) { + auto loop = dyn_cast(parent); + auto trips = loop ? getStaticTripCount(loop) : std::nullopt; + if (!trips) + return std::nullopt; + count = checkedMultiply(count, static_cast(*trips)); + } + auto type = extract.getResult().getType(); + TransferCost fragment = costModel.getTransferCostFromBytes( + costModel.getByteSize(type, costModel.getComputeBitwidth(type))); + cost = addTransferCosts(cost, + {checkedMultiply(fragment.fixed, count), + checkedMultiply(fragment.networkFlits, count)}); + } + return cost; +} + +TransferCost getInputTransferCost(Value input, const SchedulerCostModel& costModel) { auto inputType = cast(input.getType()); - if (auto batch = dyn_cast(consumerInstance.op)) - if (std::optional projectedCost = - getBatchProjectedInputTransferCost(batch, input, costModel)) - return *projectedCost; return costModel.getTransferCostFromBytes( costModel.getByteSize(inputType, costModel.getComputeBitwidth(inputType))); } -uint32_t getLaneOverlapCount(const ComputeInstance& lhs, const ComputeInstance& rhs) { - uint32_t lhsEnd = lhs.laneStart + lhs.laneCount; - uint32_t rhsEnd = rhs.laneStart + rhs.laneCount; - return std::max(lhs.laneStart, rhs.laneStart) < std::min(lhsEnd, rhsEnd) - ? std::min(lhsEnd, rhsEnd) - std::max(lhs.laneStart, rhs.laneStart) - : 0; -} - -TransferCost scaleTransferCostByLaneCount( - const TransferCost& totalCost, uint32_t totalLaneCount, - uint32_t fragmentLaneCount) { - assert(totalLaneCount > 0 && "laneCount must be positive"); - assert(fragmentLaneCount > 0 && "fragmentLaneCount must be positive"); - if (fragmentLaneCount >= totalLaneCount) - return totalCost; - return scaleTransferCost(totalCost, static_cast(fragmentLaneCount), - static_cast(totalLaneCount)); -} - -SmallVector collectProducerValueRefs( - Value value, const ComputeInstance& consumerInstance, - size_t processorCount) { +FailureOr> collectProducerValueRefs( + Value value, + size_t processorCount, GraphBatchPublicationCache& publicationCache) { SmallVector producers; - Operation* op = value.getDefiningOp(); - if (!op) - return producers; - - while (auto extract = dyn_cast(op)) { - Value source = extract.getSource(); - auto batch = dyn_cast_or_null(source.getDefiningOp()); - if (batch && batch.getNumResults() != 0) { - if (std::optional lane = getConstantExtractLane(extract)) { - ComputeInstance instance = - getBatchChunkForLane(batch, *lane, processorCount); - producers.push_back({instance, 0}); - return producers; + SmallVector pending{value}; + llvm::SmallDenseSet visited; + llvm::SmallDenseSet, 8> indexed; + auto add = [&](ComputeInstance instance, size_t result) { + if (indexed.insert({instance, result}).second) + producers.push_back({instance, result}); + }; + while (!pending.empty()) { + Value current = pending.pop_back_val(); + Operation* op = current.getDefiningOp(); + if (!op || !visited.insert(current).second) + continue; + if (auto extract = dyn_cast(op)) { + Value source = extract.getSource(); + auto batch = source.getDefiningOp(); + if (batch && batch.getNumResults() != 0) { + if (auto slot = getConstantExtractLane(extract); + slot && getConstantIntValue(extract.getMixedSizes().front()) == 1) { + unsigned result = cast(source).getResultNumber(); + auto publication = getGraphBatchPublicationMap(batch, result, publicationCache); + if (failed(publication)) + return failure(); + if (*slot >= (*publication)->physicalSlotToGraphLane.size()) + return extract.emitOpError("physical fragment slot is out of range"), failure(); + add(getBatchChunkForLane(batch, + (*publication)->physicalSlotToGraphLane[*slot], processorCount), result); + continue; + } } - - for (ComputeInstance instance : - getBatchChunksForRange(batch, 0, - static_cast(batch.getLaneCount()), - processorCount)) - producers.push_back({instance, 0}); - return producers; + pending.push_back(source); + continue; } - - value = source; - op = value.getDefiningOp(); - if (!op) - return producers; - } - - if (auto compute = dyn_cast(op)) { - producers.push_back({ - ComputeInstance {compute.getOperation(), 0, 1}, - static_cast(cast(value).getResultNumber()) - }); - return producers; - } - - if (auto batch = dyn_cast(op)) { - if (batch.getNumResults() != 0) { - if (auto projected = collectProjectedProducerValueRefs( - batch, value, consumerInstance, processorCount)) - return *projected; - std::optional producer = - getProducerValueRef(value, &consumerInstance, processorCount); - if (!producer) - return producers; - for (ComputeInstance instance : - getBatchChunksForRange(batch, producer->instance.laneStart, - producer->instance.laneCount, - processorCount)) - producers.push_back({instance, 0}); - return producers; + if (auto compute = dyn_cast(op)) { + add({compute, 0, 1}, cast(current).getResultNumber()); + continue; + } + if (auto batch = dyn_cast(op)) { + unsigned result = cast(current).getResultNumber(); + for (ComputeInstance instance : getBatchChunksForRange( + batch, 0, batch.getLaneCount(), processorCount)) + add(instance, result); + continue; + } + // Unproven shaping retains every producer, including captured region values. + llvm::append_range(pending, op->getOperands()); + if (op->getNumRegions()) { + llvm::SetVector captures; + getUsedValuesDefinedAbove(op->getRegions(), captures); + llvm::append_range(pending, captures); } - - uint32_t lane = cast(value).getResultNumber(); - ComputeInstance instance = - getBatchChunkForLane(batch, lane, processorCount); - producers.push_back({instance, lane - instance.laneStart}); - return producers; } - return producers; } +TransferCost getProducerTransferCost(Value input, + const ComputeInstance& consumerInstance, const ProducerValueRef& producerRef, + const SchedulerCostModel& costModel); + +LogicalResult collectBlueprintEdges(Value input, size_t targetIndex, + const ComputeGraph& graph, size_t processorCount, + const SchedulerCostModel& costModel, + GraphBatchPublicationCache& publicationCache, + SmallVectorImpl& edges) { + auto blueprint = input.getDefiningOp(); + if (!blueprint || !isFragmentAssembly(blueprint.getMode())) + return success(); + auto operands = blueprint.getFragmentOperandIndices(); + auto slots = blueprint.getFragmentSourceSlots(); + auto type = cast(blueprint.getOutput().getType()); + if (!operands || !slots) + return blueprint.emitOpError("compute graph requires complete fragment publication metadata"); + SmallVector sources{blueprint.getInput()}; + llvm::append_range(sources, blueprint.getFragments()); + for (auto [fragment, operand] : llvm::enumerate(*operands)) { + Value source = sources[operand]; + auto result = dyn_cast(source); + auto batch = result ? dyn_cast(result.getOwner()) : SpatComputeBatch(); + if (!batch) { + auto producers = collectProducerValueRefs(source, processorCount, publicationCache); + if (failed(producers)) + return failure(); + for (const auto &producer : *producers) { + auto indexed = graph.instanceToIndex.find(producer.instance); + if (indexed == graph.instanceToIndex.end()) + return blueprint.emitOpError("fragment source has no compute instance"); + edges.push_back({indexed->second, targetIndex, + getProducerTransferCost(source, graph.nodes[targetIndex].instance, producer, costModel)}); + } + continue; + } + auto publication = getGraphBatchPublicationMap( + batch, result.getResultNumber(), publicationCache); + int64_t slot = (*slots)[fragment]; + if (failed(publication)) + return failure(); + if (slot < 0 + || slot >= static_cast((*publication)->physicalSlotToGraphLane.size())) + return blueprint.emitOpError("physical fragment slot is out of range"); + int64_t lane = (*publication)->physicalSlotToGraphLane[slot]; + ComputeInstance producer = getBatchChunkForLane(batch, lane, processorCount); + auto indexed = graph.instanceToIndex.find(producer); + if (indexed == graph.instanceToIndex.end()) + return blueprint.emitOpError("fragment producer has no compute instance"); + Cost elements = 1; + for (int64_t size : blueprint.getFragmentSizes().slice( + fragment * type.getRank(), type.getRank())) + elements = checkedMultiply(elements, static_cast(size)); + auto fragmentType = (*publication)->publicationFragmentType; + Cost bytes = costModel.getByteSize(fragmentType, costModel.getComputeBitwidth(fragmentType)); + TransferCost transfer = costModel.getTransferCostFromBytes(bytes); + auto sourceOffsets = blueprint.getFragmentSourceOffsets(); + if (!sourceOffsets) + return blueprint.emitOpError("compute graph requires fragment source offsets"); + // Full receives can write directly into the assembly destination. Its + // endpoint memory work is already included in the consumer's assembly cost. + if ((*sourceOffsets)[fragment] == 0 + && elements == static_cast(fragmentType.getNumElements())) + transfer.fixed = 0; + edges.push_back({indexed->second, targetIndex, transfer}); + } + return success(); +} + TransferCost getProducerTransferCost( Value input, const ComputeInstance& consumerInstance, const ProducerValueRef& producerRef, const SchedulerCostModel& costModel) { TransferCost transferCost = - getInputTransferCost(consumerInstance, input, costModel); + getInputTransferCost(input, costModel); auto producerBatch = dyn_cast(producerRef.instance.op); - if (!producerBatch || producerBatch.getNumResults() == 0) + if (!producerBatch || producerBatch.getNumResults() == 0) { + if (auto projected = getProjectedInputTransferCost(input, consumerInstance, costModel)) + return *projected; return transferCost; - - if (auto consumerBatch = dyn_cast(consumerInstance.op)) { - if (std::optional projectedCost = - getBatchProjectedInputTransferCost(consumerBatch, input, costModel)) { - uint32_t overlapLaneCount = getLaneOverlapCount(consumerInstance, producerRef.instance); - return scaleTransferCost( - *projectedCost, - static_cast(std::max(1, overlapLaneCount))); - } } - return scaleTransferCostByLaneCount( - transferCost, static_cast(producerBatch.getLaneCount()), producerRef.instance.laneCount); + auto type = cast(producerBatch.getResult(producerRef.resultIndex).getType()); + Cost elements = checkedMultiply( + static_cast(type.getNumElements()) / producerBatch.getLaneCount(), + static_cast(producerRef.instance.laneCount)); + // Round the selected chunk once; dividing rounded whole-result bytes loses sub-byte payloads. + Cost bytes = SchedulerCostModel::ceilDiv( + checkedMultiply(elements, costModel.getComputeBitwidth(type)), Cost{8}); + return costModel.getTransferCostFromBytes(bytes); } static ResidentWeight getOpaqueResidentWeight(Value value, std::optional lane) { @@ -819,13 +930,16 @@ ResidentWeightSet collectDistinctResidentWeights(Operation* owner, std::optional return usage; } -Cost getComputeInstanceCost(const ComputeInstance& instance, const SchedulingTarget& target) { - SchedulerCostModel costModel {target}; +Cost getComputeInstanceCost(const ComputeInstance& instance, const SchedulingTarget& target, + ComputeCostKind kind) { + SchedulerCostModel costModel {target, kind}; + Cost assembly = getBlueprintAssemblyCost(instance, costModel); if (auto spatCompute = dyn_cast(instance.op)) - return getComputeBodyCost(spatCompute.getBody(), costModel); + return checkedAdd(getComputeBodyCost(spatCompute.getBody(), costModel), assembly); auto batch = cast(instance.op); return checkedMultiply( - getComputeBodyCost(batch.getBody(), costModel), static_cast(instance.laneCount)); + checkedAdd(getComputeBodyCost(batch.getBody(), costModel), assembly), + static_cast(instance.laneCount)); } bool containsResidentWeight(ArrayRef usage, const ResidentWeight& weight) { @@ -893,7 +1007,7 @@ ResidentWeightSet getComputeInstanceResidentWeights(const ComputeInstance& insta return tiled; } -ComputeGraph buildComputeGraph(Operation* entryOp, +FailureOr buildComputeGraph(Operation* entryOp, const SchedulingTarget& target, size_t computePartitionCount) { ComputeGraph graph; @@ -921,12 +1035,13 @@ ComputeGraph buildComputeGraph(Operation* entryOp, continue; size_t chunkCount = getBatchChunkTargetCount(batch, computePartitionCount); + Cost laneCost = getComputeInstanceCost({batch, 0, 1}, target); for (size_t chunkIndex = 0; chunkIndex < chunkCount; ++chunkIndex) { ComputeInstance instance = getBatchChunkForIndex( batch, chunkIndex, computePartitionCount); size_t index = graph.nodes.size(); graph.nodes.push_back({instance, - getComputeInstanceCost(instance, target), + checkedMultiply(laneCost, static_cast(instance.laneCount)), getComputeInstanceResidentWeights(instance, target), index}); graph.instanceToIndex[instance] = index; @@ -937,15 +1052,32 @@ ComputeGraph buildComputeGraph(Operation* entryOp, } llvm::SmallVector rawEdges; + GraphBatchPublicationCache publicationCache; for (const auto& [targetIndex, node] : llvm::enumerate(graph.nodes)) { llvm::SmallVector inputs = getComputeInstanceInputs(node.instance); for (Value input : inputs) { - for (const ProducerValueRef& producerRef : - collectProducerValueRefs(input, node.instance, - computePartitionCount)) { + if (auto producer = input.getDefiningOp()) { + auto projected = collectProjectedEdges(producer, input, node.instance, + computePartitionCount, targetIndex, graph, costModel, publicationCache, rawEdges); + if (failed(projected)) + return failure(); + if (*projected) + continue; + } + if (auto blueprint = input.getDefiningOp(); + blueprint && isFragmentAssembly(blueprint.getMode())) { + if (failed(collectBlueprintEdges(input, targetIndex, graph, computePartitionCount, + costModel, publicationCache, rawEdges))) + return failure(); + continue; + } + auto producers = collectProducerValueRefs(input, computePartitionCount, publicationCache); + if (failed(producers)) + return failure(); + for (const ProducerValueRef& producerRef : *producers) { auto producerIt = graph.instanceToIndex.find(producerRef.instance); if (producerIt == graph.instanceToIndex.end()) - continue; + return node.instance.op->emitOpError("input producer has no compute instance"), failure(); rawEdges.push_back({producerIt->second, targetIndex, getProducerTransferCost(input, node.instance, producerRef, costModel)}); diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp index 3d65845..ff61a55 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.hpp @@ -37,16 +37,26 @@ namespace spatial { struct ComputeGraphNode { ComputeInstance instance; - Cost cost = 0; + Cost cost = 0; // Per-instance serial execution estimate, in nanoseconds. ResidentWeightSet residentWeights; size_t originalOrder = 0; }; struct TransferCost { - Cost fixed = 0; - Cost networkFlits = 0; + Cost fixed = 0; // Endpoint memory time, in nanoseconds. + Cost networkFlits = 0; // Multiplied by the target's physical link time. }; +inline Time getTransferTime(const TransferCost& transferCost, + size_t sourceProcessor, + size_t targetProcessor, + const SchedulingTarget& target) { + if (sourceProcessor == targetProcessor) + return 0; + return checkedAdd(transferCost.fixed, + checkedMultiply(transferCost.networkFlits, target.getInterProcessorLatencyNs(sourceProcessor, targetProcessor))); +} + struct ComputeGraphEdge { size_t source = 0; size_t target = 0; @@ -61,7 +71,7 @@ struct ComputeGraph { llvm::DenseMap instanceToIndex; }; -ComputeGraph buildComputeGraph(mlir::Operation* entryOp, +mlir::FailureOr buildComputeGraph(mlir::Operation* entryOp, const SchedulingTarget& target, size_t computePartitionCount = 0); bool verifyAcyclic(const ComputeGraph& graph); @@ -70,7 +80,14 @@ TransferCost getTransferCostFromBytes(Cost bytes, const SchedulingTarget& target); uint64_t countComputeBodyInstructions(mlir::Region& body); uint64_t countComputeBodyOperationInstances(mlir::Region& body); -Cost getComputeInstanceCost(const ComputeInstance& instance, const SchedulingTarget& target); +enum class ComputeCostKind { + Execution, + // A nominal volume-based scheduling weight, not predicted execution time. + VolumePriority, +}; + +Cost getComputeInstanceCost(const ComputeInstance& instance, const SchedulingTarget& target, + ComputeCostKind kind = ComputeCostKind::Execution); ResidentWeightSet collectDistinctResidentWeights(mlir::Operation* owner, std::optional lane = std::nullopt); ResidentWeightSet getComputeInstanceResidentWeights(const ComputeInstance& instance); diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.cpp new file mode 100644 index 0000000..d40d7e1 --- /dev/null +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.cpp @@ -0,0 +1,144 @@ +#include "MergeSchedule.hpp" +#include "ComputeGraph.hpp" +#include "llvm/ADT/STLExtras.h" +#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp" + +namespace onnx_mlir::spatial { +using namespace mlir; +namespace { +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 || !isFragmentAssembly(blueprint.getMode())) + 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 + +void populateScheduleEquivalenceClasses( + const ComputeGraph &graph, + const std::vector> &tasksByProcessor, + MergeScheduleResult &schedule) { + auto &equivalentClass = schedule.equivalentClass; + equivalentClass.clear(); + size_t processorCount = tasksByProcessor.size(); + GraphBatchPublicationCache publicationCache; + llvm::DenseMap publicationSignatures; + llvm::DenseSet invalidPublicationSignatures; + auto haveCompatiblePublications = [&](const ComputeInstance& lhs, const ComputeInstance& rhs) { + auto batch = dyn_cast_or_null(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; ++currentProcessor) { + if (tasksByProcessor[currentProcessor].empty()) + continue; + for (size_t controlProcessor = currentProcessor + 1; controlProcessor < processorCount; ++controlProcessor) { + if (!schedule.processorStages.empty() + && schedule.processorStages[currentProcessor] != schedule.processorStages[controlProcessor]) + continue; + if (tasksByProcessor[currentProcessor].size() != tasksByProcessor[controlProcessor].size()) + continue; + auto& currentTasks = tasksByProcessor[currentProcessor]; + auto& controlTasks = tasksByProcessor[controlProcessor]; + bool equalSchedule = true; + + for (auto [currentTask, controlTask] : llvm::zip(currentTasks, controlTasks)) { + const ComputeInstance currentComputeInstance = graph.nodes[currentTask].instance; + const ComputeInstance controlComputeInstance = graph.nodes[controlTask].instance; + if (!currentComputeInstance.op || currentComputeInstance.op != controlComputeInstance.op + || currentComputeInstance.laneCount != controlComputeInstance.laneCount + || !haveCompatiblePublications(currentComputeInstance, controlComputeInstance)) { + equalSchedule = false; + break; + } + } + + if (equalSchedule) { + equivalentClass[currentProcessor].push_back(controlProcessor); + equivalentClass[controlProcessor].push_back(currentProcessor); + } + } + } + +} + +} // namespace onnx_mlir::spatial diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.hpp index c792d83..8e3c8bb 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.hpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedule.hpp @@ -12,7 +12,18 @@ namespace onnx_mlir { namespace spatial { +struct ScheduleCostSummary { + uint64_t totalWorkNs = 0; + uint64_t maximumCoreWorkNs = 0; + uint64_t criticalPathNs = 0; + uint64_t communicationNs = 0; + uint64_t completionNs = 0; + std::vector stageSpanNs; + std::vector stageSerialDemandNs; +}; + struct MergeScheduleResult { + ScheduleCostSummary costs; size_t processorCount = 0; std::vector processorStages; std::vector dominanceOrderCompute; @@ -24,5 +35,11 @@ struct MergeScheduleResult { llvm::DenseMap> equivalentClass; }; +struct ComputeGraph; +void populateScheduleEquivalenceClasses( + const ComputeGraph &graph, + const std::vector> &tasksByProcessor, + MergeScheduleResult &schedule); + } // namespace spatial } // namespace onnx_mlir diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp index 7f635fc..f62849a 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.cpp @@ -75,7 +75,7 @@ void verifySchedule(const ComputeGraph& graph, Time earliestTargetStart = addOrMax(sourceStart, graph.nodes[edge.source].cost); if (sourceCpu != targetCpu) earliestTargetStart = addOrMax( - earliestTargetStart, getPeftTransferTime(edge.transferCost, sourceCpu, targetCpu, target)); + earliestTargetStart, getTransferTime(edge.transferCost, sourceCpu, targetCpu, target)); if (targetStart < earliestTargetStart) { std::string message = llvm::formatv("merge scheduling: dependency legality failed between tasks {0} and {1}", graph.nodes[edge.source].originalOrder, @@ -90,20 +90,24 @@ void verifySchedule(const ComputeGraph& graph, MergeSchedulingAnalysis::MergeSchedulingAnalysis(const ComputeGraph& computeGraph, const SchedulingTarget& schedulingTarget, - mlir::MLIRContext* context) + mlir::MLIRContext* context, std::string& error) : context(context), target(schedulingTarget), graph(computeGraph) { - result = run(); + result = run(error); } -MergeScheduleResult MergeSchedulingAnalysis::run() { - if (!verifyAcyclic(graph)) - llvm::report_fatal_error("merge scheduling: compute graph is cyclic"); +mlir::FailureOr MergeSchedulingAnalysis::run(std::string& error) { + if (!verifyAcyclic(graph)) { + error = "merge scheduling: compute graph is cyclic"; + return mlir::failure(); + } - MergeScheduleResult schedule = runPeftScheduler( + auto schedule = runPeftScheduler( graph, PeftScheduleOptions { target, - context}); - verifySchedule(graph, schedule, target); + context}, error); + if (mlir::failed(schedule)) + return mlir::failure(); + verifySchedule(graph, *schedule, target); return schedule; } diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.hpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.hpp index f5fba6b..2735930 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.hpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/MergeSchedulingAnalysis.hpp @@ -1,5 +1,7 @@ #pragma once +#include + #include "mlir/IR/Operation.h" #include "MergeSchedule.hpp" @@ -13,16 +15,16 @@ class MergeSchedulingAnalysis { public: MergeSchedulingAnalysis(const ComputeGraph& graph, const SchedulingTarget& target, - mlir::MLIRContext* context); - MergeScheduleResult& getResult() { return result; } + mlir::MLIRContext* context, std::string& error); + mlir::FailureOr& getResult() { return result; } private: mlir::MLIRContext* context = nullptr; const SchedulingTarget& target; const ComputeGraph& graph; - MergeScheduleResult result; + mlir::FailureOr result; - MergeScheduleResult run(); + mlir::FailureOr run(std::string& error); }; } // namespace spatial diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp index 6a28900..24f325b 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp @@ -1,6 +1,5 @@ #include "mlir/IR/Threading.h" -#include "llvm/ADT/SmallSet.h" #include "llvm/Support/ErrorHandling.h" #include "llvm/Support/FormatVariadic.h" @@ -11,7 +10,8 @@ #include #include "PeftScheduler.hpp" -#include "src/Accelerators/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredProjectionAnalysis.hpp" +#include "ResidentWeightPacking.hpp" +#include "ScheduleEvaluation.hpp" namespace onnx_mlir { namespace spatial { @@ -32,12 +32,18 @@ struct ScheduledTask { struct TopologyModel { const SchedulingTarget& target; + std::vector centerDistances; + + explicit TopologyModel(const SchedulingTarget& target) + : target(target), centerDistances(target.processorCount) { + for (size_t processor = 0; processor < target.processorCount; ++processor) + for (size_t other = 0; other < target.processorCount; ++other) + centerDistances[processor] = checkedAdd(centerDistances[processor], + static_cast(target.getInterProcessorLatencyNs(processor, other))); + } size_t getCenterDistance(size_t processor) const { - Cost total = 0; - for (size_t other = 0; other < target.processorCount; ++other) - total = checkedAdd(total, target.getInterProcessorLatencyNs(processor, other)); - return static_cast(total); + return centerDistances[processor]; } }; @@ -93,18 +99,20 @@ bool hasHighResidentWeightPressure(const ComputeGraph& graph, size_t processorCo return false; } -std::vector planResidentWeightReservations(const ComputeGraph& graph, +mlir::FailureOr> planResidentWeightReservations(const ComputeGraph& graph, size_t processorCount, size_t residentWeightCapacity, const TopologyModel& topology, bool preferWeightReuse, - bool preserveTaskOrder = true) { + std::string& error, + bool largestWeightsFirst = false) { std::vector weightedTasks; for (size_t task = 0; task < graph.nodes.size(); ++task) if (!graph.nodes[task].residentWeights.empty()) weightedTasks.push_back(task); llvm::sort(weightedTasks, [&](size_t lhs, size_t rhs) { - if ((!preserveTaskOrder || preferWeightReuse) && graph.nodes[lhs].residentWeights.size() != graph.nodes[rhs].residentWeights.size()) + if ((preferWeightReuse || largestWeightsFirst) + && graph.nodes[lhs].residentWeights.size() != graph.nodes[rhs].residentWeights.size()) return graph.nodes[lhs].residentWeights.size() > graph.nodes[rhs].residentWeights.size(); return graph.nodes[lhs].originalOrder < graph.nodes[rhs].originalOrder; }); @@ -125,11 +133,11 @@ std::vector planResidentWeightReservations(const ComputeGraph Time communication = 0; for (const auto &[predecessor, cost] : graph.predecessors[task]) if (reservedCpus[predecessor] != processorCount) - communication = addOrMax(communication, getPeftTransferTime( + communication = addOrMax(communication, getTransferTime( cost, reservedCpus[predecessor], processor, topology.target)); for (const auto &[successor, cost] : graph.successors[task]) if (reservedCpus[successor] != processorCount) - communication = addOrMax(communication, getPeftTransferTime( + communication = addOrMax(communication, getTransferTime( cost, processor, reservedCpus[successor], topology.target)); Time load = addOrMax(reservedLoad[processor], communication); ReservationScore score {preferWeightReuse ? addedWeights : load, @@ -143,19 +151,38 @@ std::vector planResidentWeightReservations(const ComputeGraph } } if (bestProcessor == std::numeric_limits::max()) { - // Task order improves locality but can fragment the available crossbars. - if (preserveTaskOrder && !preferWeightReuse) + // Preserve load-aware reservations before falling back to feasibility-only packing. + if (!preferWeightReuse && !largestWeightsFirst) return planResidentWeightReservations(graph, processorCount, - residentWeightCapacity, topology, preferWeightReuse, false); - std::string message = - llvm::formatv("PEFT reservation planner: cannot place task {0} with {1} distinct weights in {2} " + residentWeightCapacity, topology, preferWeightReuse, error, true); + IndexedResidentWeights indexed(graph); + std::vector costs; + for (const ComputeGraphNode& node : graph.nodes) + costs.push_back(node.cost); + auto assignment = indexed.pack( + weightedTasks, costs, residentWeightCapacity, processorCount); + if (assignment) { + reservations.assign(processorCount, {}); + for (auto [index, weightedTask] : llvm::enumerate(weightedTasks)) + insertResidentWeights(reservations[(*assignment)[index]], + graph.nodes[weightedTask].residentWeights); + return reservations; + } + size_t lowerBound = indexed.getPackCountLowerBound( + weightedTasks, residentWeightCapacity); + if (lowerBound > processorCount) { + error = "PEFT reservation planner: resident-weight lower bound exceeds physical capacity"; + return mlir::failure(); + } + error = + llvm::formatv("PEFT resident-weight heuristic search exhausted its bounded repair: could not place task {0} with {1} distinct weights in {2} " "processors of capacity {3}", graph.nodes[task].originalOrder, graph.nodes[task].residentWeights.size(), processorCount, residentWeightCapacity) .str(); - llvm::report_fatal_error(llvm::StringRef(message)); + return mlir::failure(); } insertResidentWeights(reservations[bestProcessor], graph.nodes[task].residentWeights); reservedCpus[task] = bestProcessor; @@ -164,89 +191,24 @@ std::vector planResidentWeightReservations(const ComputeGraph 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 || !isFragmentAssembly(blueprint.getMode())) - 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 -MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftScheduleOptions& options) { +mlir::FailureOr runPeftScheduler(const ComputeGraph& graph, + const PeftScheduleOptions& options, std::string& error) { const size_t nodeCount = graph.nodes.size(); const size_t processorCount = options.target.processorCount; - if (processorCount == 0) - llvm::report_fatal_error("PEFT scheduler: processor count must be positive"); + if (processorCount == 0 || options.target.residentWeightCapacity == 0) { + error = "PEFT scheduler: processor count and resident-weight capacity must be positive"; + return mlir::failure(); + } TopologyModel topology {options.target}; const bool preferWeightReuse = hasHighResidentWeightPressure(graph, processorCount, options.target.residentWeightCapacity); - std::vector capacityReservations = planResidentWeightReservations( - graph, processorCount, options.target.residentWeightCapacity, topology, preferWeightReuse); + auto reservations = planResidentWeightReservations( + graph, processorCount, options.target.residentWeightCapacity, topology, preferWeightReuse, error); + if (mlir::failed(reservations)) + return mlir::failure(); + std::vector capacityReservations = std::move(*reservations); std::vector> reverseLevels = buildReverseLevels(graph); @@ -268,30 +230,11 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu computeNodeOct(i); } - struct RankEntry { - long double rank = 0.0L; - size_t node = 0; - size_t originalOrder = 0; - }; - std::vector ranks(nodeCount); - auto computeRank = [&](size_t node) { - long double rank = static_cast(oct[node]); - ranks[node] = {rank, node, graph.nodes[node].originalOrder}; - }; - - if (options.context != nullptr) - mlir::parallelFor(options.context, 0, nodeCount, computeRank); - else - for (size_t node = 0; node < nodeCount; ++node) - computeRank(node); - auto readyCompare = [&](size_t lhs, size_t rhs) { - const RankEntry& lhsRank = ranks[lhs]; - const RankEntry& rhsRank = ranks[rhs]; - if (lhsRank.rank != rhsRank.rank) - return lhsRank.rank < rhsRank.rank; - if (lhsRank.originalOrder != rhsRank.originalOrder) - return lhsRank.originalOrder > rhsRank.originalOrder; + if (oct[lhs] != oct[rhs]) + return oct[lhs] < oct[rhs]; + if (graph.nodes[lhs].originalOrder != graph.nodes[rhs].originalOrder) + return graph.nodes[lhs].originalOrder > graph.nodes[rhs].originalOrder; return lhs > rhs; }; @@ -309,8 +252,6 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu std::vector schedules(nodeCount); std::vector> tasksByProcessor(processorCount); std::vector> timelineByProcessor(processorCount); - std::vector taskOrder; - taskOrder.reserve(nodeCount); size_t scheduledCount = 0; while (!readyQueue.empty()) { @@ -318,15 +259,12 @@ MergeScheduleResult runPeftScheduler(const ComputeGraph& graph, const PeftSchedu readyQueue.pop(); if (scheduled[task]) continue; - taskOrder.push_back(task); size_t bestProcessor = std::numeric_limits::max(); Time bestEst = 0; Time bestEft = 0; - Time bestOeft = std::numeric_limits