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

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Jared Hoberock (jaredhoberock)

<details>
<summary>Changes</summary>

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

---
Full diff: https://github.com/llvm/llvm-project/pull/234399.diff


2 Files Affected:

- (modified) mlir/lib/Target/LLVMIR/Dialect/GPU/SelectObjectAttr.cpp (+6-2) 
- (modified) mlir/test/Target/LLVMIR/gpu.mlir (+23) 


``````````diff
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

``````````

</details>


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