[Mlir-commits] [mlir] b0cc322 - [mlir] Check for argument uses in test-func-erase-arg pass (#203367)

[llvmlistbot at llvm.org]

Author: lijinpei-amd
Date: 2026-06-12T15:57:19+02:00
New Revision: b0cc3225b2f618b5760552b1dc7ce0ea26225c8f

URL: https://github.com/llvm/llvm-project/commit/b0cc3225b2f618b5760552b1dc7ce0ea26225c8f
DIFF: https://github.com/llvm/llvm-project/commit/b0cc3225b2f618b5760552b1dc7ce0ea26225c8f.diff

LOG: [mlir] Check for argument uses in test-func-erase-arg pass (#203367)

The -test-func-erase-arg pass crashed when erasing arguments that still
had uses. Diagnose every such argument and fail the pass without
erasing.

Fixes https://github.com/llvm/llvm-project/issues/203218

Assisted-by: Claude (Claude Code)

Added: 
    mlir/test/IR/test-func-erase-arg-error.mlir

Modified: 
    mlir/test/lib/IR/TestFunc.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/test/IR/test-func-erase-arg-error.mlir b/mlir/test/IR/test-func-erase-arg-error.mlir
new file mode 100644
index 0000000000000..bf826517c0c4d
--- /dev/null
+++ b/mlir/test/IR/test-func-erase-arg-error.mlir
@@ -0,0 +1,15 @@
+// RUN: mlir-opt %s -test-func-erase-arg -split-input-file -verify-diagnostics
+
+// Erasing arguments that still have uses must fail (see #203218).
+
+// expected-error @below {{cannot erase argument 0 which still has uses}}
+// expected-error @below {{cannot erase argument 1 which still has uses}}
+func.func @f(%arg0: f32 {test.erase_this_arg},
+             %arg1: f32 {test.erase_this_arg}) -> (f32, f32) {
+  return %arg0, %arg1 : f32, f32
+}
+
+// -----
+
+// An external function has no body, so erasure must succeed.
+func.func private @ext(f32 {test.erase_this_arg})

diff  --git a/mlir/test/lib/IR/TestFunc.cpp b/mlir/test/lib/IR/TestFunc.cpp
index 94a4610365863..2d4050d3484d0 100644
--- a/mlir/test/lib/IR/TestFunc.cpp
+++ b/mlir/test/lib/IR/TestFunc.cpp
@@ -106,9 +106,19 @@ struct TestFuncEraseArg
 
     for (auto func : module.getOps<FunctionOpInterface>()) {
       BitVector indicesToErase(func.getNumArguments());
-      for (auto argIndex : llvm::seq<int>(0, func.getNumArguments()))
-        if (func.getArgAttr(argIndex, "test.erase_this_arg"))
-          indicesToErase.set(argIndex);
+      bool hasUsedArg = false;
+      for (auto argIndex : llvm::seq<int>(0, func.getNumArguments())) {
+        if (!func.getArgAttr(argIndex, "test.erase_this_arg"))
+          continue;
+        indicesToErase.set(argIndex);
+        if (!func.isExternal() && !func.getArgument(argIndex).use_empty()) {
+          emitError(func->getLoc()) << "cannot erase argument " << argIndex
+                                    << " which still has uses";
+          hasUsedArg = true;
+        }
+      }
+      if (hasUsedArg)
+        return signalPassFailure();
       if (succeeded(func.eraseArguments(indicesToErase)))
         continue;
       emitError(func->getLoc()) << "failed to erase arguments";