[Mlir-commits] [mlir] d8f3be7 - [mlir][dialect-conversion] Fix OOB crash in convertFuncOpTypes for funcs with extra block args (#185060)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-11T14:25:03+01:00
New Revision: d8f3be726e0dfc869c62dfd5f9c486c1270d3d25

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

LOG: [mlir][dialect-conversion] Fix OOB crash in convertFuncOpTypes for funcs with extra block args (#185060)

Some function ops (e.g., gpu.func with workgroup memory arguments) have
more entry block arguments than their FunctionType has inputs. The
workgroup memory arguments are not part of the public function signature
but are present as additional block arguments.

`convertFuncOpTypes` previously created a `SignatureConversion` sized
only for `type.getNumInputs()`, then called `applySignatureConversion`
on the entry block. When the block had more arguments (e.g., workgroup
args), the loop in `applySignatureConversion` would call
`getInputMapping(i)` with out-of-bounds indices, causing an assertion
failure in `SmallVector::operator[]`.

Fix this by:
1. Sizing the `SignatureConversion` for all entry block arguments.
2. Adding identity mappings for extra block args beyond the function
type inputs.
3. Using only the converted function-type-input types when updating the
FunctionType (so extra block arg types are not included in the
signature).

Fixes #184744

Assisted-by: Claude Code

Added: 
    

Modified: 
    mlir/lib/Transforms/Utils/DialectConversion.cpp
    mlir/test/Transforms/test-legalize-type-conversion.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 1dfa2103f57a7..ad3ee5d20a534 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -3843,19 +3843,39 @@ static LogicalResult convertFuncOpTypes(FunctionOpInterface funcOp,
   if (!type)
     return failure();
 
-  // Convert the original function types.
-  TypeConverter::SignatureConversion result(type.getNumInputs());
+  // Convert the function signature (inputs and results).
+  TypeConverter::SignatureConversion funcConversion(type.getNumInputs());
   SmallVector<Type, 1> newResults;
-  if (failed(typeConverter.convertSignatureArgs(type.getInputs(), result)) ||
+  if (failed(typeConverter.convertSignatureArgs(type.getInputs(),
+                                                funcConversion)) ||
       failed(typeConverter.convertTypes(type.getResults(), newResults)))
     return failure();
-  if (!funcOp.getFunctionBody().empty())
-    rewriter.applySignatureConversion(&funcOp.getFunctionBody().front(), result,
+
+  // If the function has a body, apply a separate signature conversion to the
+  // entry block. Some function ops (e.g., gpu.func) have extra block arguments
+  // beyond the function type inputs (e.g., workgroup memory arguments that are
+  // not part of the public signature). Use a distinct conversion sized for all
+  // entry block arguments so that applySignatureConversion does not access
+  // out-of-bounds mappings.
+  if (!funcOp.getFunctionBody().empty()) {
+    Block *entryBlock = &funcOp.getFunctionBody().front();
+    unsigned numEntryBlockArgs = entryBlock->getNumArguments();
+    unsigned numFuncTypeInputs = type.getNumInputs();
+    TypeConverter::SignatureConversion blockConversion(numEntryBlockArgs);
+    // Convert the function-type inputs the same way as for the function type.
+    if (failed(typeConverter.convertSignatureArgs(type.getInputs(),
+                                                  blockConversion)))
+      return failure();
+    // Add identity mappings for extra block args beyond the function type
+    // inputs. These arguments are preserved as-is.
+    for (unsigned i = numFuncTypeInputs; i < numEntryBlockArgs; ++i)
+      blockConversion.addInputs(i, entryBlock->getArgument(i).getType());
+    rewriter.applySignatureConversion(entryBlock, blockConversion,
                                       &typeConverter);
+  }
 
-  // Update the function signature in-place.
-  auto newType = FunctionType::get(rewriter.getContext(),
-                                   result.getConvertedTypes(), newResults);
+  auto newType = FunctionType::get(
+      rewriter.getContext(), funcConversion.getConvertedTypes(), newResults);
 
   rewriter.modifyOpInPlace(funcOp, [&] { funcOp.setType(newType); });
 

diff  --git a/mlir/test/Transforms/test-legalize-type-conversion.mlir b/mlir/test/Transforms/test-legalize-type-conversion.mlir
index 91f83a0afaeef..ed675c36752b0 100644
--- a/mlir/test/Transforms/test-legalize-type-conversion.mlir
+++ b/mlir/test/Transforms/test-legalize-type-conversion.mlir
@@ -165,3 +165,22 @@ func.func @test_unstructured_cf_conversion(%arg0: f32, %c: i1) {
   "test.bar"(%arg2) : (f32) -> ()
   return
 }
+
+// -----
+
+// Test that gpu.func with workgroup arguments (which are extra block arguments
+// beyond the function type inputs) does not crash during signature conversion.
+// See https://github.com/llvm/llvm-project/issues/184744
+
+// CHECK-LABEL: gpu.func @func_with_workgroup_args
+// CHECK-SAME: (%{{.*}}: memref<?xi32>, %{{.*}}: i32) workgroup(%{{.*}} : memref<512xi32>)
+gpu.module @cuda_events {
+  gpu.func @func_with_workgroup_args(%arg0: memref<?xi32>, %arg1: i32)
+      workgroup(%workgroup_mem: memref<512xi32>)
+      kernel {
+    %idx = arith.constant 0 : index
+    %val = memref.load %arg0[%idx] : memref<?xi32>
+    memref.store %val, %arg0[%idx] : memref<?xi32>
+    gpu.return
+  }
+}