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

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/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

>From 5c6a1a48647a172d20e9f79763dec4eb1899f24e Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Thu, 5 Mar 2026 02:50:28 -0800
Subject: [PATCH] [mlir][dialect-conversion] Fix OOB crash in
 convertFuncOpTypes for funcs with extra block args

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
---
 .../Transforms/Utils/DialectConversion.cpp    | 37 +++++++++++++++----
 .../test-legalize-type-conversion.mlir        | 19 ++++++++++
 2 files changed, 49 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/Transforms/Utils/DialectConversion.cpp b/mlir/lib/Transforms/Utils/DialectConversion.cpp
index 1dfa2103f57a7..c80f4917799c0 100644
--- a/mlir/lib/Transforms/Utils/DialectConversion.cpp
+++ b/mlir/lib/Transforms/Utils/DialectConversion.cpp
@@ -3843,19 +3843,42 @@ static LogicalResult convertFuncOpTypes(FunctionOpInterface funcOp,
   if (!type)
     return failure();
 
+  // Some function ops may have extra block arguments beyond the function type
+  // inputs (e.g., workgroup memory arguments in gpu.func). Size the
+  // SignatureConversion to cover all entry block arguments so that
+  // applySignatureConversion does not access out-of-bounds mappings.
+  unsigned numFuncTypeInputs = type.getNumInputs();
+  Block *entryBlock = !funcOp.getFunctionBody().empty()
+                          ? &funcOp.getFunctionBody().front()
+                          : nullptr;
+  unsigned numEntryBlockArgs =
+      entryBlock ? entryBlock->getNumArguments() : numFuncTypeInputs;
+
   // Convert the original function types.
-  TypeConverter::SignatureConversion result(type.getNumInputs());
+  TypeConverter::SignatureConversion result(numEntryBlockArgs);
   SmallVector<Type, 1> newResults;
   if (failed(typeConverter.convertSignatureArgs(type.getInputs(), result)) ||
       failed(typeConverter.convertTypes(type.getResults(), newResults)))
     return failure();
-  if (!funcOp.getFunctionBody().empty())
-    rewriter.applySignatureConversion(&funcOp.getFunctionBody().front(), result,
-                                      &typeConverter);
 
-  // Update the function signature in-place.
-  auto newType = FunctionType::get(rewriter.getContext(),
-                                   result.getConvertedTypes(), newResults);
+  // Track how many converted types correspond to the function signature inputs
+  // before adding identity mappings for any extra block arguments.
+  unsigned numConvertedFuncArgTypes = result.getConvertedTypes().size();
+
+  if (entryBlock) {
+    // Add identity mappings for any extra block arguments beyond the function
+    // type inputs. These arguments are preserved as-is in the new block.
+    for (unsigned i = numFuncTypeInputs; i < numEntryBlockArgs; ++i)
+      result.addInputs(i, entryBlock->getArgument(i).getType());
+    rewriter.applySignatureConversion(entryBlock, result, &typeConverter);
+  }
+
+  // Update the function signature in-place, using only the converted types for
+  // the function type inputs (not the extra block arguments).
+  auto newType = FunctionType::get(
+      rewriter.getContext(),
+      result.getConvertedTypes().take_front(numConvertedFuncArgTypes),
+      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
+  }
+}