[Mlir-commits] [flang] [mlir] Add operands to worklist when only used by deleted op (PR #86990)

llvmlistbot at llvm.org llvmlistbot at llvm.org
Fri Mar 29 00:36:53 PDT 2024


https://github.com/mlevesquedion updated https://github.com/llvm/llvm-project/pull/86990

>From 27500e1e09bcebed3413d0114b829d2acff45643 Mon Sep 17 00:00:00 2001
From: Michael Levesque-Dion <mlevesquedion at google.com>
Date: Thu, 28 Mar 2024 11:53:07 -0700
Subject: [PATCH 1/7] Add operands to worklist when only used by deleted op

Fixes #86765
---
 .../Utils/GreedyPatternRewriteDriver.cpp      | 24 ++++----
 ...edy-pattern-rewrite-driver-bottom-up.mlir} |  0
 ...reedy-pattern-rewrite-driver-top-down.mlir | 60 +++++++++++++++++++
 3 files changed, 73 insertions(+), 11 deletions(-)
 rename mlir/test/IR/{greedy-pattern-rewriter-driver.mlir => greedy-pattern-rewrite-driver-bottom-up.mlir} (100%)
 create mode 100644 mlir/test/IR/greedy-pattern-rewrite-driver-top-down.mlir

diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index 6cb5635e68c922..2f2a15a0a28e51 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -377,7 +377,7 @@ class GreedyPatternRewriteDriver : public PatternRewriter,
   /// be re-added to the worklist. This function should be called when an
   /// operation is modified or removed, as it may trigger further
   /// simplifications.
-  void addOperandsToWorklist(ValueRange operands);
+  void addOperandsToWorklist(Operation* op);
 
   /// Notify the driver that the given block was inserted.
   void notifyBlockInserted(Block *block, Region *previous,
@@ -688,15 +688,17 @@ void GreedyPatternRewriteDriver::notifyOperationModified(Operation *op) {
   addToWorklist(op);
 }
 
-void GreedyPatternRewriteDriver::addOperandsToWorklist(ValueRange operands) {
-  for (Value operand : operands) {
-    // If the use count of this operand is now < 2, we re-add the defining
-    // operation to the worklist.
-    // TODO: This is based on the fact that zero use operations
-    // may be deleted, and that single use values often have more
-    // canonicalization opportunities.
-    if (!operand || (!operand.use_empty() && !operand.hasOneUse()))
-      continue;
+void GreedyPatternRewriteDriver::addOperandsToWorklist(Operation* op) {
+  for (Value operand : op->getOperands()) {
+    // If this operand was only used by the op under consideration, we re-add
+    // the operation that defined it to the worklist. Indeed, if the op is about
+    // to be deleted and it was the sole user of the operand, the operand may
+    // also be deleted.
+    // TODO: if the operand has a single use besides the op under consideration,
+    // there may be further canonicalization opportunities, so it should be
+    // added to the worklist.
+    if (!operand) continue;
+    if (!llvm::all_of(operand.getUsers(), [&op](auto u) { return u == op; })) continue;
     if (auto *defOp = operand.getDefiningOp())
       addToWorklist(defOp);
   }
@@ -722,7 +724,7 @@ void GreedyPatternRewriteDriver::notifyOperationErased(Operation *op) {
   if (config.listener)
     config.listener->notifyOperationErased(op);
 
-  addOperandsToWorklist(op->getOperands());
+  addOperandsToWorklist(op);
   worklist.remove(op);
 
   if (config.strictMode != GreedyRewriteStrictness::AnyOp)
diff --git a/mlir/test/IR/greedy-pattern-rewriter-driver.mlir b/mlir/test/IR/greedy-pattern-rewrite-driver-bottom-up.mlir
similarity index 100%
rename from mlir/test/IR/greedy-pattern-rewriter-driver.mlir
rename to mlir/test/IR/greedy-pattern-rewrite-driver-bottom-up.mlir
diff --git a/mlir/test/IR/greedy-pattern-rewrite-driver-top-down.mlir b/mlir/test/IR/greedy-pattern-rewrite-driver-top-down.mlir
new file mode 100644
index 00000000000000..5a0d58c62845bf
--- /dev/null
+++ b/mlir/test/IR/greedy-pattern-rewrite-driver-top-down.mlir
@@ -0,0 +1,60 @@
+// RUN: mlir-opt %s -test-patterns="max-iterations=1 top-down=true" \
+// RUN:     -allow-unregistered-dialect --split-input-file | FileCheck %s
+
+// Tests for https://github.com/llvm/llvm-project/issues/86765. Ensure
+// that operands of a dead op are added to the worklist even if the same value
+// appears multiple times as an operand.
+
+// -----
+
+// 2 uses of the same operand
+
+// CHECK:       func.func @f(%arg0: i1) {
+// CHECK-NEXT:    return
+// CHECK-NEXT:  }
+func.func @f(%arg0: i1) {
+  %0 = arith.constant 0 : i32
+  %if = scf.if %arg0 -> (i32) {
+    scf.yield %0 : i32
+  } else {
+    scf.yield %0 : i32
+  }
+  %dead_leaf = arith.addi %if, %if : i32
+  return
+}
+
+// -----
+
+// 3 uses of the same operand
+
+// CHECK:       func.func @f() {
+// CHECK-NEXT:    return
+// CHECK-NEXT:  }
+func.func @f() {
+  %0 = arith.constant 0 : i1
+  %if = scf.if %0 -> (i1) {
+    scf.yield %0 : i1
+  } else {
+    scf.yield %0 : i1
+  }
+  %dead_leaf = arith.select %if, %if, %if : i1
+  return
+}
+
+// -----
+
+// 2 uses of the same operand, op has 3 operands
+
+// CHECK:       func.func @f() {
+// CHECK-NEXT:    return
+// CHECK-NEXT:  }
+func.func @f() {
+  %0 = arith.constant 0 : i1
+  %if = scf.if %0 -> (i1) {
+    scf.yield %0 : i1
+  } else {
+    scf.yield %0 : i1
+  }
+  %dead_leaf = arith.select %0, %if, %if : i1
+  return
+}

>From e2fd08df237d418e6c5b431b6cb0f05380957a7a Mon Sep 17 00:00:00 2001
From: Michael Levesque-Dion <mlevesquedion at google.com>
Date: Thu, 28 Mar 2024 13:28:31 -0700
Subject: [PATCH 2/7] Fix formatting

---
 .../Transforms/Utils/GreedyPatternRewriteDriver.cpp    | 10 ++++++----
 1 file changed, 6 insertions(+), 4 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index 2f2a15a0a28e51..e7a94fab82af43 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -377,7 +377,7 @@ class GreedyPatternRewriteDriver : public PatternRewriter,
   /// be re-added to the worklist. This function should be called when an
   /// operation is modified or removed, as it may trigger further
   /// simplifications.
-  void addOperandsToWorklist(Operation* op);
+  void addOperandsToWorklist(Operation *op);
 
   /// Notify the driver that the given block was inserted.
   void notifyBlockInserted(Block *block, Region *previous,
@@ -688,7 +688,7 @@ void GreedyPatternRewriteDriver::notifyOperationModified(Operation *op) {
   addToWorklist(op);
 }
 
-void GreedyPatternRewriteDriver::addOperandsToWorklist(Operation* op) {
+void GreedyPatternRewriteDriver::addOperandsToWorklist(Operation *op) {
   for (Value operand : op->getOperands()) {
     // If this operand was only used by the op under consideration, we re-add
     // the operation that defined it to the worklist. Indeed, if the op is about
@@ -697,8 +697,10 @@ void GreedyPatternRewriteDriver::addOperandsToWorklist(Operation* op) {
     // TODO: if the operand has a single use besides the op under consideration,
     // there may be further canonicalization opportunities, so it should be
     // added to the worklist.
-    if (!operand) continue;
-    if (!llvm::all_of(operand.getUsers(), [&op](auto u) { return u == op; })) continue;
+    if (!operand)
+      continue;
+    if (!llvm::all_of(operand.getUsers(), [&op](auto u) { return u == op; }))
+      continue;
     if (auto *defOp = operand.getDefiningOp())
       addToWorklist(defOp);
   }

>From 28e0df4f35acb7ecdd429f692a4083ea389793ec Mon Sep 17 00:00:00 2001
From: Michael Levesque-Dion <mlevesquedion at google.com>
Date: Thu, 28 Mar 2024 13:43:29 -0700
Subject: [PATCH 3/7] Spell out auto

---
 mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index e7a94fab82af43..e4e87a3c7f5778 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -699,7 +699,8 @@ void GreedyPatternRewriteDriver::addOperandsToWorklist(Operation *op) {
     // added to the worklist.
     if (!operand)
       continue;
-    if (!llvm::all_of(operand.getUsers(), [&op](auto u) { return u == op; }))
+    if (!llvm::all_of(operand.getUsers(),
+                      [&op](Operation *u) { return u == op; }))
       continue;
     if (auto *defOp = operand.getDefiningOp())
       addToWorklist(defOp);

>From a8dd0551eeb0cb67de04d18fbf194c8d7b339e7c Mon Sep 17 00:00:00 2001
From: Michael Levesque-Dion <mlevesquedion at google.com>
Date: Thu, 28 Mar 2024 15:45:02 -0700
Subject: [PATCH 4/7] Address review comments

---
 .../Utils/GreedyPatternRewriteDriver.cpp      | 21 +++++++++++--------
 ...reedy-pattern-rewrite-driver-top-down.mlir |  2 +-
 2 files changed, 13 insertions(+), 10 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index e4e87a3c7f5778..e32d0b8252e9a0 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -690,17 +690,20 @@ void GreedyPatternRewriteDriver::notifyOperationModified(Operation *op) {
 
 void GreedyPatternRewriteDriver::addOperandsToWorklist(Operation *op) {
   for (Value operand : op->getOperands()) {
-    // If this operand was only used by the op under consideration, we re-add
-    // the operation that defined it to the worklist. Indeed, if the op is about
-    // to be deleted and it was the sole user of the operand, the operand may
-    // also be deleted.
-    // TODO: if the operand has a single use besides the op under consideration,
-    // there may be further canonicalization opportunities, so it should be
-    // added to the worklist.
+    // If this operand currently has at most 2 users, add its defining op to the
+    // worklist. Indeed, after the op is deleted, then the operand will have at
+    // most 1 user left. If it has 0 users left, it can be deleted too,
+    // and if it has 1 user left, there may be further canonicalization
+    // opportunities.
     if (!operand)
       continue;
-    if (!llvm::all_of(operand.getUsers(),
-                      [&op](Operation *u) { return u == op; }))
+    Operation *otherUser = nullptr;
+    if (!llvm::all_of(operand.getUsers(), [&](Operation *user) {
+      if (user == op) return true;
+      if (otherUser && user != otherUser) return false;
+      otherUser = user;
+      return true;
+    }))
       continue;
     if (auto *defOp = operand.getDefiningOp())
       addToWorklist(defOp);
diff --git a/mlir/test/IR/greedy-pattern-rewrite-driver-top-down.mlir b/mlir/test/IR/greedy-pattern-rewrite-driver-top-down.mlir
index 5a0d58c62845bf..03b7cc817a3f18 100644
--- a/mlir/test/IR/greedy-pattern-rewrite-driver-top-down.mlir
+++ b/mlir/test/IR/greedy-pattern-rewrite-driver-top-down.mlir
@@ -1,5 +1,5 @@
 // RUN: mlir-opt %s -test-patterns="max-iterations=1 top-down=true" \
-// RUN:     -allow-unregistered-dialect --split-input-file | FileCheck %s
+// RUN:     --split-input-file | FileCheck %s
 
 // Tests for https://github.com/llvm/llvm-project/issues/86765. Ensure
 // that operands of a dead op are added to the worklist even if the same value

>From 46ec146c636b37be2465446cad540f8540500cf0 Mon Sep 17 00:00:00 2001
From: Michael Levesque-Dion <mlevesquedion at google.com>
Date: Thu, 28 Mar 2024 16:13:57 -0700
Subject: [PATCH 5/7] Address review comments and format code

---
 .../Utils/GreedyPatternRewriteDriver.cpp      | 21 ++++++++++++-------
 1 file changed, 14 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index e32d0b8252e9a0..2aae77fbf5ac9f 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -697,16 +697,23 @@ void GreedyPatternRewriteDriver::addOperandsToWorklist(Operation *op) {
     // opportunities.
     if (!operand)
       continue;
+
+    auto *defOp = operand.getDefiningOp();
+    if (!defOp)
+      continue;
+
     Operation *otherUser = nullptr;
     if (!llvm::all_of(operand.getUsers(), [&](Operation *user) {
-      if (user == op) return true;
-      if (otherUser && user != otherUser) return false;
-      otherUser = user;
-      return true;
-    }))
+          if (user == op)
+            return true;
+          if (otherUser && user != otherUser)
+            return false;
+          otherUser = user;
+          return true;
+        }))
       continue;
-    if (auto *defOp = operand.getDefiningOp())
-      addToWorklist(defOp);
+
+    addToWorklist(defOp);
   }
 }
 

>From 14faca9d10be002252a8ea6740cb3c24ca1ce8ea Mon Sep 17 00:00:00 2001
From: Michael Levesque-Dion <mlevesquedion at google.com>
Date: Thu, 28 Mar 2024 17:14:56 -0700
Subject: [PATCH 6/7] Fix flang test failures -- a few more canonicalizations
 are happening

---
 flang/test/Transforms/stack-arrays.fir | 9 +++------
 1 file changed, 3 insertions(+), 6 deletions(-)

diff --git a/flang/test/Transforms/stack-arrays.fir b/flang/test/Transforms/stack-arrays.fir
index f4fe737e88d785..a2ffe555091ebc 100644
--- a/flang/test/Transforms/stack-arrays.fir
+++ b/flang/test/Transforms/stack-arrays.fir
@@ -127,9 +127,7 @@ func.func @placement1() {
   return
 }
 // CHECK:      func.func @placement1() {
-// CHECK-NEXT:   %[[ONE:.*]] = arith.constant 1 : index
-// CHECK-NEXT:   %[[TWO:.*]] = arith.constant 2 : index
-// CHECK-NEXT:   %[[ARG:.*]] = arith.addi %[[ONE]], %[[TWO]] : index
+// CHECK-NEXT:   %[[ARG:.*]] = arith.constant 3 : index
 // CHECK-NEXT:   %[[MEM:.*]] = fir.alloca !fir.array<?xi32>, %[[ARG]]
 // CHECK-NEXT:   return
 // CHECK-NEXT: }
@@ -204,13 +202,12 @@ func.func @placement4(%arg0 : i1) {
 // CHECK:      func.func @placement4(%arg0: i1) {
 // CHECK-NEXT:   %[[C1:.*]] = arith.constant 1 : index
 // CHECK-NEXT:   %[[C1_I32:.*]] = fir.convert %[[C1]] : (index) -> i32
-// CHECK-NEXT:   %[[C2:.*]] = arith.constant 2 : index
 // CHECK-NEXT:   %[[C10:.*]] = arith.constant 10 : index
 // CHECK-NEXT:   cf.br ^bb1
 // CHECK-NEXT: ^bb1:
-// CHECK-NEXT:   %[[SUM:.*]] = arith.addi %[[C1]], %[[C2]] : index
+// CHECK-NEXT:   %[[C3:.*]] = arith.constant 3 : index
 // CHECK-NEXT:   %[[SP:.*]] = fir.call @llvm.stacksave.p0() : () -> !fir.ref<i8>
-// CHECK-NEXT:   %[[MEM:.*]] = fir.alloca !fir.array<?xi32>, %[[SUM]]
+// CHECK-NEXT:   %[[MEM:.*]] = fir.alloca !fir.array<?xi32>, %[[C3]]
 // CHECK-NEXT:   fir.call @llvm.stackrestore.p0(%[[SP]]) : (!fir.ref<i8>) -> ()
 // CHECK-NEXT:   cf.cond_br %arg0, ^bb1, ^bb2
 // CHECK-NEXT: ^bb2:

>From ebf195a4f3696d0514f42e6868d464ab3d6cdbce Mon Sep 17 00:00:00 2001
From: Michael Levesque-Dion <mlevesquedion at google.com>
Date: Fri, 29 Mar 2024 00:36:38 -0700
Subject: [PATCH 7/7] Address review comments

---
 .../Utils/GreedyPatternRewriteDriver.cpp      | 20 +++++++++++--------
 ...reedy-pattern-rewrite-driver-top-down.mlir |  2 --
 2 files changed, 12 insertions(+), 10 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
index 2aae77fbf5ac9f..bbecbdb8566935 100644
--- a/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
+++ b/mlir/lib/Transforms/Utils/GreedyPatternRewriteDriver.cpp
@@ -703,14 +703,18 @@ void GreedyPatternRewriteDriver::addOperandsToWorklist(Operation *op) {
       continue;
 
     Operation *otherUser = nullptr;
-    if (!llvm::all_of(operand.getUsers(), [&](Operation *user) {
-          if (user == op)
-            return true;
-          if (otherUser && user != otherUser)
-            return false;
-          otherUser = user;
-          return true;
-        }))
+    bool hasMoreThanTwoUses = false;
+    for (auto user : operand.getUsers()) {
+      if (user == op || user == otherUser)
+        continue;
+      if (!otherUser) {
+        otherUser = user;
+        continue;
+      }
+      hasMoreThanTwoUses = true;
+      break;
+    }
+    if (hasMoreThanTwoUses)
       continue;
 
     addToWorklist(defOp);
diff --git a/mlir/test/IR/greedy-pattern-rewrite-driver-top-down.mlir b/mlir/test/IR/greedy-pattern-rewrite-driver-top-down.mlir
index 03b7cc817a3f18..a362d6f99b9478 100644
--- a/mlir/test/IR/greedy-pattern-rewrite-driver-top-down.mlir
+++ b/mlir/test/IR/greedy-pattern-rewrite-driver-top-down.mlir
@@ -5,8 +5,6 @@
 // that operands of a dead op are added to the worklist even if the same value
 // appears multiple times as an operand.
 
-// -----
-
 // 2 uses of the same operand
 
 // CHECK:       func.func @f(%arg0: i1) {



More information about the Mlir-commits mailing list