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

[Jacques Pienaar]

https://github.com/jpienaar updated https://github.com/llvm/llvm-project/pull/123642

>From fe81a79ae1a278e2f9a209dc6929c73184f564a6 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 1/3] 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 | 22 +++++++++++++-------
 1 file changed, 14 insertions(+), 8 deletions(-)

diff --git a/mlir/lib/Interfaces/SideEffectInterfaces.cpp b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
index c9feb001a19844..9862c6c0ce7a6b 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;
 
@@ -41,10 +42,17 @@ bool mlir::isOpTriviallyDead(Operation *op) {
 /// allows for marking region operations as trivially dead without always being
 /// 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);
+  // The set of operation intervals (end-exclusive) to consider when checking
+  // for side effects.
+  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();
+    if (it == end) {
+      effectingOps.pop_back();
+    }
+    Block::iterator &it = effectingOps.back().first;
+    Block::iterator end = effectingOps.back().second;
+    mlir::Operation *op = &*(it++);
 
     // If the operation has recursive effects, push all of the nested operations
     // on to the stack to consider.
@@ -53,8 +61,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 +93,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;
 

>From d3a683380febbbc5ec4b54f12d5c607dcb6952f3 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jacques+gh at japienaar.info>
Date: Fri, 24 Jan 2025 13:07:02 -0800
Subject: [PATCH 2/3] Fix definition ordering

---
 mlir/lib/Interfaces/SideEffectInterfaces.cpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Interfaces/SideEffectInterfaces.cpp b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
index 9862c6c0ce7a6b..88303cd68e5f61 100644
--- a/mlir/lib/Interfaces/SideEffectInterfaces.cpp
+++ b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
@@ -47,11 +47,11 @@ static bool wouldOpBeTriviallyDeadImpl(Operation *rootOp) {
   SmallVector<std::pair<Block::iterator, Block::iterator>, 1> effectingOps = {
       std::make_pair(Block::iterator(rootOp), ++Block::iterator(rootOp))};
   while (!effectingOps.empty()) {
+    Block::iterator &it = effectingOps.back().first;
+    Block::iterator end = effectingOps.back().second;
     if (it == end) {
       effectingOps.pop_back();
     }
-    Block::iterator &it = effectingOps.back().first;
-    Block::iterator end = effectingOps.back().second;
     mlir::Operation *op = &*(it++);
 
     // If the operation has recursive effects, push all of the nested operations

>From 3065545130f14aa2c7d7c43217c39b03a6f91d77 Mon Sep 17 00:00:00 2001
From: Jacques Pienaar <jpienaar at google.com>
Date: Sat, 25 Jan 2025 14:17:33 +0000
Subject: [PATCH 3/3] Missing continue

---
 mlir/lib/Interfaces/SideEffectInterfaces.cpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/mlir/lib/Interfaces/SideEffectInterfaces.cpp b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
index 88303cd68e5f61..59fd19310cea54 100644
--- a/mlir/lib/Interfaces/SideEffectInterfaces.cpp
+++ b/mlir/lib/Interfaces/SideEffectInterfaces.cpp
@@ -51,6 +51,7 @@ static bool wouldOpBeTriviallyDeadImpl(Operation *rootOp) {
     Block::iterator end = effectingOps.back().second;
     if (it == end) {
       effectingOps.pop_back();
+      continue;
     }
     mlir::Operation *op = &*(it++);