[Mlir-commits] [mlir] [mlir][gpu] Allocate kernel argument storage in the entry block (PR #234399)

[Jared Hoberock]

https://github.com/jaredhoberock created https://github.com/llvm/llvm-project/pull/234399

In the LLVM translation of `gpu.launch_func`, hoist allocas for the kernel argument struct and pointer array into the enclosing function's entry block instead of emitting them at the launch site. Allocas outside the entry block are dynamic, so a kernel launch in a loop grows the stack on every iteration and can overflow the stack for a large loop.

The driver copies arguments at the launch call, so reusing storage is safe.

Add a test that checks that a launch's allocas end up in the entry block.

Assisted-by: Claude

>From 343b5801ea7b65a1d3da18812e3da73203df6d59 Mon Sep 17 00:00:00 2001
From: Jared Hoberock <jaredhoberock at gmail.com>
Date: Tue, 20 Oct 2026 12:46:48 -0500
Subject: [PATCH] [mlir][gpu] Allocate kernel argument storage in the entry
 block

The LLVM IR translation of `gpu.launch_func` (through
`#gpu.select_object`) emits two allocas at the launch: a struct containing
the kernel arguments and an array of pointers to its fields. When the
launch is in a loop, these are dynamic allocas executed on every
iteration, so the stack grows with the trip count until the function
returns.

This patch emits the two allocas at the start of the function's entry
block, where they are static allocations in the frame. The stores of the
arguments stay at the launch. Reusing the storage across launches is safe
because the launch call reads the arguments when it is made; an
asynchronous launch that outlives the function's frame already relies on
that today.

Assisted-by: Claude
---
 .../LLVMIR/Dialect/GPU/SelectObjectAttr.cpp   |  8 +++++--
 mlir/test/Target/LLVMIR/gpu.mlir              | 23 +++++++++++++++++++
 2 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp b/mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp
index d40a2789a9feb1..cff300c18fb48c 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp
@@ -384,8 +384,12 @@ llvm::LaunchKernel::createKernelArgArray(mlir::gpu::LaunchFuncOp op) {
     structTypes[i] = arg->getType();
 
   Type *structTy = StructType::create(module.getContext(), structTypes);
-  Value *argStruct = builder.CreateAlloca(structTy, 0u);
-  Value *argArray = builder.CreateAlloca(
+  // Allocate in the entry block, so that a launch in a loop does not grow the
+  // stack on every iteration.
+  BasicBlock &entry = builder.GetInsertBlock()->getParent()->getEntryBlock();
+  IRBuilder<> allocaBuilder(&entry, entry.getFirstInsertionPt());
+  Value *argStruct = allocaBuilder.CreateAlloca(structTy, 0u);
+  Value *argArray = allocaBuilder.CreateAlloca(
       ptrTy, ConstantInt::get(intPtrTy, structTypes.size()));
 
   for (auto [i, arg] : enumerate(args)) {
diff --git a/mlir/test/Target/LLVMIR/gpu.mlir b/mlir/test/Target/LLVMIR/gpu.mlir
index 140fabca444daa..13b528428e48a1 100644
--- a/mlir/test/Target/LLVMIR/gpu.mlir
+++ b/mlir/test/Target/LLVMIR/gpu.mlir
@@ -43,6 +43,29 @@ module attributes {gpu.container_module} {
 
 // -----
 
+// Checking that the argument storage of a launch in a loop is allocated in the
+// entry block.
+module attributes {gpu.container_module} {
+  gpu.binary @kernel_module  [#gpu.object<#nvvm.target, "BLOB">]
+  // CHECK-LABEL: define void @launch_in_loop(
+  // CHECK-NEXT: alloca %{{.*}}, align 8
+  // CHECK-NEXT: [[ARGS_ARRAY:%.*]] = alloca ptr, i64 1, align 8
+  // CHECK-NEXT: br label
+  // CHECK-NOT: alloca
+  // CHECK: call void @mgpuLaunchKernel({{.*}}, ptr [[ARGS_ARRAY]], ptr null, i64 1)
+  llvm.func @launch_in_loop(%cond : i1) {
+    %0 = llvm.mlir.constant(8 : i64) : i64
+    llvm.br ^bb1
+  ^bb1:
+    gpu.launch_func @kernel_module::@kernel blocks in (%0, %0, %0) threads in (%0, %0, %0) : i64 args(%0 : i64)
+    llvm.cond_br %cond, ^bb1, ^bb2
+  ^bb2:
+    llvm.return
+  }
+}
+
+// -----
+
 // Checking the correct selection of the second object using an index as a selector.
 module {
   // CHECK: @kernel_module_binary = internal constant [1 x i8] c"1", align 8