Skip to content

Commit 819cb78

Browse files
authored
[Pipelining] make schedule.dump deterministic (#7294)
Lately I've been working on AMD's pipelining and I've been looking at schedule dumps. They're great but unfortunately (due to `DenseMap`) not deterministic. This PR makes `opToStageAndCluster` a `MapVector` instead (which produces deterministic dumps).
1 parent 4a119a3 commit 819cb78

File tree

2 files changed

+7
-6
lines changed

2 files changed

+7
-6
lines changed

include/triton/Dialect/TritonGPU/Transforms/Schedule.h

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,7 +85,7 @@ class CoarseSchedule {
8585
using Cluster = ClusterList::iterator;
8686
using ClusterHash = size_t;
8787

88-
DenseMap<Operation *, std::pair<int, Cluster>> opToStageAndCluster;
88+
llvm::MapVector<Operation *, std::pair<int, Cluster>> opToStageAndCluster;
8989

9090
void setNumStages(int numStages) { this->numStages = numStages; }
9191
int getNumStages() const { return numStages; }

lib/Dialect/TritonGPU/Transforms/Pipeliner/Schedule.cpp

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -126,8 +126,8 @@ tt::CoarseSchedule::splitClusterBefore(Operation *op, scf::ForOp forOp) {
126126
bool tt::CoarseSchedule::isOpBefore(Operation *a, Operation *b) const {
127127
assert(opToStageAndCluster.count(a) && opToStageAndCluster.count(b) &&
128128
"Operations must be in the schedule");
129-
auto [aStage, aCluster] = opToStageAndCluster.at(a);
130-
auto [bStage, bCluster] = opToStageAndCluster.at(b);
129+
auto [aStage, aCluster] = opToStageAndCluster.lookup(a);
130+
auto [bStage, bCluster] = opToStageAndCluster.lookup(b);
131131
if (aStage != bStage) {
132132
return aStage < bStage;
133133
}
@@ -141,14 +141,15 @@ bool tt::CoarseSchedule::isOpInEarlierCluster(Operation *a,
141141
Operation *b) const {
142142
assert(opToStageAndCluster.count(a) && opToStageAndCluster.count(b) &&
143143
"Operations must be in the schedule");
144-
return clusters.isBefore(opToStageAndCluster.at(a).second,
145-
opToStageAndCluster.at(b).second);
144+
return clusters.isBefore(opToStageAndCluster.lookup(a).second,
145+
opToStageAndCluster.lookup(b).second);
146146
}
147147

148148
bool tt::CoarseSchedule::isOpInSameCluster(Operation *a, Operation *b) const {
149149
assert(opToStageAndCluster.count(a) && opToStageAndCluster.count(b) &&
150150
"Operations must be in the schedule");
151-
return opToStageAndCluster.at(a).second == opToStageAndCluster.at(b).second;
151+
return opToStageAndCluster.lookup(a).second ==
152+
opToStageAndCluster.lookup(b).second;
152153
}
153154

154155
SmallVector<std::tuple<Operation *, int, tt::CoarseSchedule::Cluster>>

0 commit comments

Comments
 (0)