[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 updated https://github.com/llvm/llvm-project/pull/185060

>From 4357898cd6337126488a4620a1d85c7d81ce0d40 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    | 36 ++++++++++++++-----
 .../test-legalize-type-conversion.mlir        | 19 ++++++++++
 2 files changed, 47 insertions(+), 8 deletions(-)

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