Hold a queue of iterator ranges (not operations) in wouldOpBeTriviallyDead

[apaszke]

Ranges let us push the whole blocks onto the queue in constant time. If one of the first ops in the block is side-effecting we'll be able to provide the answer quickly. The previous implementation had to walk the block and queue all the operations only to start traversing them again, which was a considerable slowdown for compile times of large MLIR programs in our benchmarks.

[llvmbot]

@llvm/pr-subscribers-mlir

Author: Adam Paszke (apaszke)

Changes

Ranges let us push the whole blocks onto the queue in constant time. If one of the first ops in the block is side-effecting we'll be able to provide the answer quickly. The previous implementation had to walk the block and queue all the operations only to start traversing them again, which was a considerable slowdown for compile times of large MLIR programs in our benchmarks.

---

Full diff: https://github.com/llvm/llvm-project/pull/123642.diff

1 Files Affected:

• (modified) mlir/lib/Interfaces/SideEffectInterfaces.cpp (+12-7)

diff --git a/mlir/lib/Interfaces/SideEffectInterfaces.cpp b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
index c9feb001a19844..97fda71d120fe0 100644
--- a/mlir/lib/Interfaces/SideEffectInterfaces.cpp
+++ b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
@@ -10,6 +10,7 @@
 
 #include "mlir/IR/SymbolTable.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include <utility>
 
 using namespace mlir;
 
@@ -42,9 +43,15 @@ bool mlir::isOpTriviallyDead(Operation *op) {
 /// conservative of terminators.
 static bool wouldOpBeTriviallyDeadImpl(Operation *rootOp) {
   // The set of operations to consider when checking for side effects.
-  SmallVector<Operation *, 1> effectingOps(1, rootOp);
+  SmallVector<std::pair<Block::iterator, Block::iterator>, 1> effectingOps = {
+      std::make_pair(Block::iterator(rootOp), ++Block::iterator(rootOp))};
   while (!effectingOps.empty()) {
-    Operation *op = effectingOps.pop_back_val();
+    Block::iterator &it = effectingOps.back().first;
+    Block::iterator end = effectingOps.back().second;
+    mlir::Operation *op = &*(it++);
+    if (it == end) {
+      effectingOps.pop_back();
+    }
 
     // If the operation has recursive effects, push all of the nested operations
     // on to the stack to consider.
@@ -53,8 +60,7 @@ static bool wouldOpBeTriviallyDeadImpl(Operation *rootOp) {
     if (hasRecursiveEffects) {
       for (Region &region : op->getRegions()) {
         for (auto &block : region) {
-          for (auto &nestedOp : block)
-            effectingOps.push_back(&nestedOp);
+          effectingOps.push_back(std::make_pair(block.begin(), block.end()));
         }
       }
     }
@@ -86,10 +92,9 @@ static bool wouldOpBeTriviallyDeadImpl(Operation *rootOp) {
         return false;
       }
       continue;
-
-      // Otherwise, if the op has recursive side effects we can treat the
-      // operation itself as having no effects.
     }
+    // Otherwise, if the op only has recursive side effects we can treat the
+    // operation itself as having no effects. We will visit its children next.
     if (hasRecursiveEffects)
       continue;
 

[jpienaar]

Typo? This should probably be below their declaration below.