[mlir][dialect-conversion] Fix OOB crash in convertFuncOpTypes for funcs with extra block args (llvm#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 llvm#184744

Assisted-by: Claude Code

Author: joker-eph

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); });
 

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
+  }
+}