From 4b9e2f8fe0fcb1d104fd10a149372aaf7ca99af2 Mon Sep 17 00:00:00 2001 From: Uday Bondhugula Date: Thu, 29 Dec 2022 01:36:21 +0530 Subject: [PATCH] [MLIR][Affine] Make fusion helper check method significantly more efficient The `hasDependencePath` method in affine fusion is quite inefficient as it does a DFS on the complete graph for what is a small part of the checks before fusion can be performed. Make this efficient by using the fact that the nodes involved are all at the top-level of the same block. With this change, for large graphs with about 10,000 nodes, the check runs in a few seconds instead of not terminating even in a few hours. This is NFC from a functionality standpoint; it only leads to an improvement in pass running time on large IR. Differential Revision: https://reviews.llvm.org/D140522 --- mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp index 90e2f75..5628b82 100644 --- a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp +++ b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp @@ -328,11 +328,13 @@ public: } // Returns true if there is a path in the dependence graph from node 'srcId' - // to node 'dstId'. Returns false otherwise. + // to node 'dstId'. Returns false otherwise. `srcId`, `dstId`, and the + // operations that the edges connected are expected to be from the same block. bool hasDependencePath(unsigned srcId, unsigned dstId) { // Worklist state is: SmallVector, 4> worklist; worklist.push_back({srcId, 0}); + Operation *dstOp = getNode(dstId)->op; // Run DFS traversal to see if 'dstId' is reachable from 'srcId'. while (!worklist.empty()) { auto &idAndIndex = worklist.back(); @@ -350,8 +352,12 @@ public: Edge edge = outEdges[idAndIndex.first][idAndIndex.second]; // Increment next output edge index for 'idAndIndex'. ++idAndIndex.second; - // Add node at 'edge.id' to worklist. - worklist.push_back({edge.id, 0}); + // Add node at 'edge.id' to the worklist. We don't need to consider + // nodes that are "after" dstId in the containing block; one can't have a + // path to `dstId` from any of those nodes. + bool afterDst = dstOp->isBeforeInBlock(getNode(edge.id)->op); + if (!afterDst && edge.id != idAndIndex.first) + worklist.push_back({edge.id, 0}); } return false; } -- 2.7.4