[Mlir-commits] [mlir] Hold a queue of iterator ranges (not operations) in wouldOpBeTriviallyDead (PR #123642)

[Adam Paszke]

https://github.com/apaszke created https://github.com/llvm/llvm-project/pull/123642

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.

>From 55267457296c50938750e0014e5ce83c9afd52bb Mon Sep 17 00:00:00 2001
From: Adam Paszke <adam.paszke at gmail.com>
Date: Mon, 20 Jan 2025 16:24:15 +0000
Subject: [PATCH] Hold a queue of iterator ranges (not operations) in
 wouldOpBeTriviallyDead

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.
---
 mlir/lib/Interfaces/SideEffectInterfaces.cpp | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/Interfaces/SideEffectInterfaces.cpp b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
index c9feb001a198446..97fda71d120fe04 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;