[Mlir-commits] [mlir] [mlir][Transforms] Replace dead function argument uses with poison (PR #208881)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: jpwang (jjppp)

<details>
<summary>Changes</summary>

Dropping all uses of a dead function argument can leave null operands in surviving side-effecting operations in unreachable code.
Subsequent canonicalization may then crash while folding the invalid operation.

This PR replaces the remaining uses with `ub.poison` before erasing the argument and add a regression test.

Fixes #<!-- -->206920, #<!-- -->203226

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


2 Files Affected:

- (modified) mlir/lib/Transforms/RemoveDeadValues.cpp (+14-3) 
- (modified) mlir/test/Transforms/remove-dead-values.mlir (+49) 


``````````diff
diff --git a/mlir/lib/Transforms/RemoveDeadValues.cpp b/mlir/lib/Transforms/RemoveDeadValues.cpp
index 6e55bc390be23..247c66aacef3a 100644
--- a/mlir/lib/Transforms/RemoveDeadValues.cpp
+++ b/mlir/lib/Transforms/RemoveDeadValues.cpp
@@ -640,9 +640,20 @@ static void cleanUpDeadVals(MLIRContext *ctx, RDVFinalCleanupList &list) {
       llvm::interleaveComma(f.nonLiveRets.set_bits(), os);
       os << "]";
     });
-    // Drop all uses of the dead arguments.
-    for (auto deadIdx : f.nonLiveArgs.set_bits())
-      f.funcOp.getArgument(deadIdx).dropAllUses();
+    // Replace remaining uses of the dead arguments with poison values. Simply
+    // dropping the uses would leave null operands behind in ops that survive
+    // the pass (e.g. a side-effecting op in an unreachable function, whose
+    // values are initialized as dead by the liveness analysis).
+    for (auto deadIdx : f.nonLiveArgs.set_bits()) {
+      BlockArgument arg = f.funcOp.getArgument(deadIdx);
+      // Avoid creating an unused poison value if there are no uses to replace.
+      if (arg.use_empty())
+        continue;
+      rewriter.setInsertionPointToStart(arg.getOwner());
+      Value poison =
+          ub::PoisonOp::create(rewriter, arg.getLoc(), arg.getType());
+      rewriter.replaceAllUsesWith(arg, poison);
+    }
     // Some functions may not allow erasing arguments or results. These calls
     // return failure in such cases without modifying the function, so it's okay
     // to proceed.
diff --git a/mlir/test/Transforms/remove-dead-values.mlir b/mlir/test/Transforms/remove-dead-values.mlir
index 390a448060b7f..207ce48cba4f5 100644
--- a/mlir/test/Transforms/remove-dead-values.mlir
+++ b/mlir/test/Transforms/remove-dead-values.mlir
@@ -895,3 +895,52 @@ module @func_with_non_call_users {
   }
   spirv.EntryPoint "GLCompute" @callee
 }
+
+// -----
+
+// Verify that the upper bound of a for loop inside an unreachable
+// function is correctly replaced with a poison value.
+// CHECK-LABEL: module @unreachable_func_with_for_loops
+// CHECK-NEXT:    func.func private @main_func() {
+// CHECK-NEXT:      %[[ARG3:.*]] = ub.poison : index
+// CHECK-NEXT:      %[[C0:.*]] = ub.poison : index
+// CHECK-NEXT:      %[[C1:.*]] = ub.poison : index
+// CHECK-NEXT:      %[[C10:.*]] = ub.poison : index
+// CHECK-NEXT:      scf.for %{{.*}} = %[[C0]] to %[[C10]] step %[[C1]] {
+// CHECK-NEXT:        gpu.barrier
+// CHECK-NEXT:      }
+// CHECK-NEXT:      scf.for %{{.*}} = %[[C0]] to %[[ARG3]] step %[[C1]] {
+// CHECK-NEXT:        gpu.barrier
+// CHECK-NEXT:      }
+// CHECK-NEXT:      return
+
+// CHECK-CANONICALIZE-LABEL: module @unreachable_func_with_for_loops
+// CHECK-CANONICALIZE-NEXT:    func.func private @main_func() {
+// CHECK-CANONICALIZE-NEXT:      %[[ARG3:.*]] = ub.poison : index
+// CHECK-CANONICALIZE-NEXT:      %[[C0:.*]] = ub.poison : index
+// CHECK-CANONICALIZE-NEXT:      %[[C1:.*]] = ub.poison : index
+// CHECK-CANONICALIZE-NEXT:      %[[C10:.*]] = ub.poison : index
+// CHECK-CANONICALIZE-NEXT:      scf.for %{{.*}} = %[[C0]] to %[[C10]] step %[[C1]] {
+// CHECK-CANONICALIZE-NEXT:        gpu.barrier
+// CHECK-CANONICALIZE-NEXT:      }
+// CHECK-CANONICALIZE-NEXT:      scf.for %{{.*}} = %[[C0]] to %[[ARG3]] step %[[C1]] {
+// CHECK-CANONICALIZE-NEXT:        gpu.barrier
+// CHECK-CANONICALIZE-NEXT:      }
+// CHECK-CANONICALIZE-NEXT:      return
+module @unreachable_func_with_for_loops {
+  func.func private @main_func(%arg0: index, %arg1: memref<?xf32>, %arg2: index, %arg3: index) {
+    %cst = arith.constant 0.000000e+00 : f32
+    %c0 = arith.constant 0 : index
+    %c1 = arith.constant 1 : index
+    %c10 = arith.constant 10 : index
+    scf.for %arg4 = %c0 to %c10 step %c1 {
+      %0 = memref.load %arg1[%arg4] : memref<?xf32>
+      gpu.barrier
+    }
+    scf.for %arg4 = %c0 to %arg3 step %c1 {
+      %0 = memref.load %arg1[%arg4] : memref<?xf32>
+      gpu.barrier
+    }
+    return
+  }
+}

``````````

</details>


https://github.com/llvm/llvm-project/pull/208881