diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp index fafce40..0701b7b 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredBoundaryPlanning.cpp @@ -277,10 +277,10 @@ static LogicalResult assignPipelineSynchronization( llvm::sort(transfers.stageZeroCores); transfers.stageZeroCores.erase( llvm::unique(transfers.stageZeroCores), transfers.stageZeroCores.end()); - llvm::sort(transfers.downstreamCores); - transfers.downstreamCores.erase( - llvm::unique(transfers.downstreamCores), - transfers.downstreamCores.end()); + DenseSet downstreamSeen; + llvm::erase_if(transfers.downstreamCores, [&](int64_t core) { + return !downstreamSeen.insert(core).second; + }); llvm::erase_if(transfers.downstreamCores, [&](int64_t core) { return llvm::is_contained(transfers.stageZeroCores, core); }); diff --git a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp index 47c0790..aa19607 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/DeferredCommunicationScheduling.cpp @@ -36,9 +36,9 @@ struct TransferGroup { SmallVector dependencies; unsigned unsatisfied = 0; bool scheduled = false; - std::tuple originalPriority; + std::tuple originalPriority; std::tuple> staticPriority; + std::tuple> staticPriority; }; struct StreamProgress { @@ -189,8 +189,12 @@ static SmallVector buildGroups(DeferredTransferPlan& plan) { unsigned source = family.sourceStreams.valueAt(0); thresholdByStream[source] = std::max(thresholdByStream.lookup(source), family.requirement->producer->step + 1); minSource = std::min(minSource, source); - for (size_t index = 0; index < family.targetStreams.size(); ++index) - minTarget = std::min(minTarget, family.targetStreams.valueAt(index)); + for (size_t index = 0; index < family.targetStreams.size(); ++index) { + unsigned target = family.targetStreams.valueAt(index); + thresholdByStream[target] = std::max( + thresholdByStream.lookup(target), exchange->consumerStep); + minTarget = std::min(minTarget, target); + } } for (LocalAvailabilityFamily& local : exchange->local) for (LaneInterval interval : local.targetLanes.intervals()) @@ -206,7 +210,7 @@ static SmallVector buildGroups(DeferredTransferPlan& plan) { group.dependencies.push_back({stream, thresholdByStream.lookup(stream)}); group.unsatisfied = llvm::count_if(group.dependencies, [](StreamThreshold dependency) { return dependency.completedStep != 0; }); - group.originalPriority = {exchange->consumerStep, minSource, minTarget, exchange->exchangeId}; + group.originalPriority = {exchange->exchangeId, exchange->consumerStep, minSource, minTarget}; orderGroupSlices(group); groups.push_back(std::move(group)); } 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 85c6f27..9beb2f3 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/ComputeGraph.cpp @@ -24,6 +24,7 @@ #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/Accelerators/PIM/Common/IR/ShapeUtils.hpp" #include "src/Support/TypeUtilities.hpp" namespace onnx_mlir { @@ -233,8 +234,20 @@ Cost getOperationCost(Operation& op, const SchedulerCostModel& costModel) { costModel.getUnaryVectorCost(softmax.getInput().getType(), softmax.getOutput().getType()); return checkedMultiply(unary, static_cast(4)); } - if (auto extract = dyn_cast(&op)) + if (auto extract = dyn_cast(&op)) { + bool weightOnly = !extract.getResult().use_empty() + && llvm::all_of(extract.getResult().getUses(), [&](OpOperand &use) { + auto vmm = dyn_cast(use.getOwner()); + return vmm && vmm.getWeight() == extract.getResult(); + }); + if (weightOnly) + return 0; + if (extract.getSourceType().hasStaticShape() && isContiguousSubviewWithDynamicOffsets( + extract.getSourceType().getShape(), extract.getMixedOffsets(), + extract.getStaticSizes(), extract.getStaticStrides())) + return 0; return costModel.getTensorMoveCost(extract.getResult().getType()); + } if (auto insert = dyn_cast(&op)) return costModel.getTensorMoveCost(insert.getSource().getType()); 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 a298731..6a28900 100644 --- a/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp +++ b/src/PIM/Dialect/Spatial/Passes/Transforms/MergeComputeNodes/Scheduling/PeftScheduler.cpp @@ -97,22 +97,24 @@ std::vector planResidentWeightReservations(const ComputeGraph size_t processorCount, size_t residentWeightCapacity, const TopologyModel& topology, - bool preferWeightReuse) { + bool preferWeightReuse, + bool preserveTaskOrder = true) { 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 (graph.nodes[lhs].residentWeights.size() != graph.nodes[rhs].residentWeights.size()) + if ((!preserveTaskOrder || preferWeightReuse) && 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; }); std::vector reservations(processorCount); std::vector