mirror of
https://github.com/ROCm/ROCm.git
synced 2026-04-05 03:01:17 -04:00
[BACKEND] Refactor RemoveLayoutConversion pass (#2181)
Significant changes to the pass logic. Move away from greedy rewrites and use more global analysis instead. The pass is now bocken down into 2 main phases. First forward propagation of layout starting from ops that we don't want to change. Propagate to all the nodes. If there is a single layout needed for the op then we can rewrite the op, if there are multiple layout required based on dependency we need a tie break. The second phase is backward propgation that gets a backward slice of operations starting from the convert and if all the operations in the slice can be rematerialized rewrite the slice. This backward phase now supports going through loop arguments. This will allow more complex logic in the future to add a cost model to decide which convert to leave and which to fold
This commit is contained in:
@@ -473,9 +473,11 @@ struct DFSState {
|
||||
SmallVector<Operation *, 16> topologicalCounts;
|
||||
DenseSet<Operation *> seen;
|
||||
|
||||
/// We mark each op as ready if all its operands are seen. If an op is ready,
|
||||
/// we add it to the queue. Otherwise, we keep adding its operands to the
|
||||
/// ancestors set.
|
||||
/// We mark each op as ready if all its operands and parents ops are seen. If
|
||||
/// an op is ready, we add it to the queue. Otherwise, we keep adding its
|
||||
/// operands to the ancestors set.
|
||||
/// We always want an op to be scheduled after all its parents to handle
|
||||
/// correctly cases with scf operations.
|
||||
void addToReadyQueue(Operation *op, DFSSubgraphState &subGraph,
|
||||
SmallVector<Operation *, 4> &readyQueue) {
|
||||
bool ready = true;
|
||||
@@ -486,6 +488,14 @@ struct DFSState {
|
||||
ready = false;
|
||||
}
|
||||
}
|
||||
Operation *parent = op->getParentOp();
|
||||
while (parent) {
|
||||
if (!seen.count(parent)) {
|
||||
subGraph.push_back(parent);
|
||||
ready = false;
|
||||
}
|
||||
parent = parent->getParentOp();
|
||||
}
|
||||
if (ready)
|
||||
readyQueue.push_back(op);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user