[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:
Thomas
2023-08-28 19:05:16 -07:00
committed by GitHub
parent 1465b573e8
commit d4644d6cb3
7 changed files with 1248 additions and 814 deletions

View File

@@ -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);
}