[Mlir-commits] [mlir] [mlir] Make eraseFunctionArguments fail instead of crash on used args (PR #203367)

[llvmlistbot at llvm.org]

https://github.com/lijinpei-amd created https://github.com/llvm/llvm-project/pull/203367

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

Assisted-by: Claude (Claude Code)

>From a071dc3b97e3354c11550c0a9e4f6e8db9804ec3 Mon Sep 17 00:00:00 2001
From: Li Jinpei <jinpli at amd.com>
Date: Fri, 12 Jun 2026 02:43:40 +0800
Subject: [PATCH] [mlir] Make eraseFunctionArguments fail instead of crash on
 used args

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

Assisted-by: Claude (Claude Code)
---
 mlir/include/mlir/Interfaces/FunctionInterfaces.h |  8 +++++---
 .../include/mlir/Interfaces/FunctionInterfaces.td |  6 +++---
 mlir/lib/Interfaces/FunctionInterfaces.cpp        | 10 +++++++++-
 mlir/test/IR/test-func-erase-arg-error.mlir       | 15 +++++++++++++++
 4 files changed, 32 insertions(+), 7 deletions(-)
 create mode 100644 mlir/test/IR/test-func-erase-arg-error.mlir

diff --git a/mlir/include/mlir/Interfaces/FunctionInterfaces.h b/mlir/include/mlir/Interfaces/FunctionInterfaces.h
index e10e9bd342702..dd3d8605e44ff 100644
--- a/mlir/include/mlir/Interfaces/FunctionInterfaces.h
+++ b/mlir/include/mlir/Interfaces/FunctionInterfaces.h
@@ -67,9 +67,11 @@ void insertFunctionResults(FunctionOpInterface op,
                            ArrayRef<DictionaryAttr> resultAttrs,
                            unsigned originalNumResults, Type newType);
 
-/// Erase the specified arguments and update the function type attribute.
-void eraseFunctionArguments(FunctionOpInterface op, const BitVector &argIndices,
-                            Type newType);
+/// Erase the specified arguments and update the function type attribute. Returns
+/// failure if the function cannot be updated to have the new signature, e.g.
+/// because one of the erased arguments still has uses.
+LogicalResult eraseFunctionArguments(FunctionOpInterface op,
+                                     const BitVector &argIndices, Type newType);
 
 /// Erase the specified results and update the function type attribute.
 void eraseFunctionResults(FunctionOpInterface op,
diff --git a/mlir/include/mlir/Interfaces/FunctionInterfaces.td b/mlir/include/mlir/Interfaces/FunctionInterfaces.td
index f701e828ed641..371d9319e25e5 100644
--- a/mlir/include/mlir/Interfaces/FunctionInterfaces.td
+++ b/mlir/include/mlir/Interfaces/FunctionInterfaces.td
@@ -320,14 +320,14 @@ def FunctionOpInterface : OpInterface<"FunctionOpInterface", [
     }
 
     /// Erases the arguments listed in `argIndices`. Returns failure if the
-    /// function cannot be updated to have the new signature.
+    /// function cannot be updated to have the new signature, e.g. because one
+    /// of the erased arguments still has uses.
     ::llvm::LogicalResult eraseArguments(const ::llvm::BitVector &argIndices) {
       ::mlir::Type newType = $_op.getTypeWithoutArgs(argIndices);
       if (!newType)
         return ::llvm::failure();
-      ::mlir::function_interface_impl::eraseFunctionArguments(
+      return ::mlir::function_interface_impl::eraseFunctionArguments(
         $_op, argIndices, newType);
-      return ::llvm::success();
     }
 
     /// Erase a single result at `resultIndex`. Returns failure if the function
diff --git a/mlir/lib/Interfaces/FunctionInterfaces.cpp b/mlir/lib/Interfaces/FunctionInterfaces.cpp
index e0f1135e992ac..a096f8c642961 100644
--- a/mlir/lib/Interfaces/FunctionInterfaces.cpp
+++ b/mlir/lib/Interfaces/FunctionInterfaces.cpp
@@ -278,13 +278,20 @@ void function_interface_impl::insertFunctionResults(
   op.setFunctionTypeAttr(TypeAttr::get(newType));
 }
 
-void function_interface_impl::eraseFunctionArguments(
+LogicalResult function_interface_impl::eraseFunctionArguments(
     FunctionOpInterface op, const BitVector &argIndices, Type newType) {
   // There are 3 things that need to be updated:
   // - Function type.
   // - Arg attrs.
   // - Block arguments of entry block, if not empty.
 
+  if (!op.isExternal()) {
+    Block &entry = op->getRegion(0).front();
+    for (unsigned i : argIndices.set_bits())
+      if (!entry.getArgument(i).use_empty())
+        return failure();
+  }
+
   // Update the argument attributes of the function.
   if (ArrayAttr argAttrs = op.getArgAttrsAttr()) {
     SmallVector<DictionaryAttr, 4> newArgAttrs;
@@ -303,6 +310,7 @@ void function_interface_impl::eraseFunctionArguments(
     Block &entry = op->getRegion(0).front();
     entry.eraseArguments(argIndices);
   }
+  return success();
 }
 
 void function_interface_impl::eraseFunctionResults(
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..4c2ecb762fe57
--- /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 an argument that still has uses must fail gracefully instead of
+// crashing (see https://github.com/llvm/llvm-project/issues/203218).
+
+// expected-error @below {{failed to erase arguments}}
+func.func @f(%arg0: f32 {test.erase_this_arg}) -> f32 {
+  return %arg0 : f32
+}
+
+// -----
+
+// An external function has no body, so there are no argument uses to check and
+// erasure must succeed without crashing.
+func.func private @ext(f32 {test.erase_this_arg})