[Mlir-commits] [mlir] 4b9e2f8 - [MLIR][Affine] Make fusion helper check method significantly more efficient

[Uday Bondhugula]

Author: Uday Bondhugula
Date: 2022-12-29T01:50:29+05:30
New Revision: 4b9e2f8fe0fcb1d104fd10a149372aaf7ca99af2

URL: https://github.com/llvm/llvm-project/commit/4b9e2f8fe0fcb1d104fd10a149372aaf7ca99af2
DIFF: https://github.com/llvm/llvm-project/commit/4b9e2f8fe0fcb1d104fd10a149372aaf7ca99af2.diff

LOG: [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

Added: 
    

Modified: 
    mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
index 90e2f755ba03d..5628b82d9d41d 100644
--- a/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Dialect/Affine/Transforms/LoopFusion.cpp
@@ -328,11 +328,13 @@ struct MemRefDependenceGraph {
   }
 
   // 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: <node-id, next-output-edge-index-to-visit>
     SmallVector<std::pair<unsigned, unsigned>, 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 @@ struct MemRefDependenceGraph {
       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;
   }