[Mlir-commits] [llvm] [mlir] [OpenMP][OMPIRBuilder] Fix device shared-memory struct alloc/free imb… (PR #229616)
Sunil Shrestha
llvmlistbot at llvm.org
Tue Oct 13 14:48:51 PDT 2026
https://github.com/sshrestha-aa updated https://github.com/llvm/llvm-project/pull/229616
>From e0727529e10616ec5eeed464f2fb99d4e49a24ba Mon Sep 17 00:00:00 2001
From: Sunil Shrestha <sunil.shrestha at hpe.com>
Date: Tue, 29 Sep 2026 15:59:00 -0500
Subject: [PATCH] [OpenMP][OMPIRBuilder] Fix device shared-memory struct
alloc/free imbalance for parallel in a loop
When outlining an `omp.parallel` for a generic-mode target kernel, the
captured-args aggregate (`structArg`) was allocated with `__kmpc_alloc_shared`
in the outer function's alloca block but freed with `__kmpc_free_shared` at the
parallel region exit. When the parallel is nested in a loop, the allocation runs
once while the free runs every iteration, corrupting the device shared-memory
stack and faulting at -O0 (latent bug at higher opt levels).
For the device shared-memory path, allocate the aggregate at the parallel
region's entry block instead of the outer function's alloca block, paired with
the free at the region exit, so the alloc/free runtime calls stay in the same
scope and balanced on every iteration. The host path is unchanged since it uses
a stack alloca reclaimed implicitly at function return.
Fixes #226611
Assisted-by: Copilot
---
llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 14 ++---
.../omptarget-debug-shared-alloc-loc.mlir | 4 +-
...omptarget-parallel-in-loop-shared-mem.mlir | 51 +++++++++++++++++++
.../LLVMIR/omptarget-parallel-llvm.mlir | 2 +-
.../openmp-target-private-shared-mem.mlir | 4 +-
5 files changed, 64 insertions(+), 11 deletions(-)
create mode 100644 mlir/test/Target/LLVMIR/omptarget-parallel-in-loop-shared-mem.mlir
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index b6a11bc3c805e..b17880a8a0f38 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -2115,11 +2115,15 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createParallel(
}
OI->FixUpNonEntryAllocas = true;
- OI->OuterAllocBB = OuterAllocaBlock;
+ // With device shared memory the aggregate arguments struct is allocated and
+ // freed with explicit runtime calls, so its allocation and deallocation must
+ // stay in the same scope to remain balanced when the region is nested in a
+ // loop. Allocate it at the parallel region's (caller-side) entry block,
+ // paired with the free at the region exit. Host uses a stack alloca reclaimed
+ // implicitly, so it keeps the function entry block.
+ OI->OuterAllocBB = UsesDeviceSharedMemory ? EntryBB : OuterAllocaBlock;
OI->EntryBB = PRegEntryBB;
OI->ExitBB = PRegExitBB;
- OI->OuterDeallocBBs.reserve(OuterDeallocBlocks.size());
- copy(OuterDeallocBlocks, OI->OuterDeallocBBs.end());
SmallPtrSet<BasicBlock *, 32> ParallelRegionBlockSet;
SmallVector<BasicBlock *, 32> Blocks;
@@ -12385,9 +12389,7 @@ std::unique_ptr<CodeExtractor> DeviceSharedMemOutlineInfo::createCodeExtractor(
/* AllowVarArgs */ true,
/* AllowAlloca */ true,
/* AllocationBlock*/ OuterAllocBB,
- /* DeallocationBlocks */ OuterDeallocBBs.empty()
- ? SmallVector<BasicBlock *>{ExitBB}
- : OuterDeallocBBs,
+ /* DeallocationBlocks */ SmallVector<BasicBlock *>{ExitBB},
/* Suffix */ Suffix.str(), ArgsInZeroAddressSpace);
}
diff --git a/mlir/test/Target/LLVMIR/omptarget-debug-shared-alloc-loc.mlir b/mlir/test/Target/LLVMIR/omptarget-debug-shared-alloc-loc.mlir
index 19f7fe573b4be..6b5f2b6094928 100644
--- a/mlir/test/Target/LLVMIR/omptarget-debug-shared-alloc-loc.mlir
+++ b/mlir/test/Target/LLVMIR/omptarget-debug-shared-alloc-loc.mlir
@@ -35,8 +35,8 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memo
// must carry a debug location, scoped to the correct function.
// CHECK: define {{.*}}@__omp_offloading_{{.*}} !dbg ![[SP:[0-9]+]] {
-// CHECK: call {{.*}}@__kmpc_alloc_shared(i64 16), !dbg ![[LOC:[0-9]+]]
-// CHECK: call {{.*}}@__kmpc_alloc_shared(i64 4), !dbg ![[LOC]]
+// CHECK: call {{.*}}@__kmpc_alloc_shared(i64 4), !dbg ![[LOC:[0-9]+]]
+// CHECK: call {{.*}}@__kmpc_alloc_shared(i64 16), !dbg ![[LOC]]
// CHECK: call void @__kmpc_free_shared(ptr {{.*}}, i64 16), !dbg ![[LOC]]
// CHECK: call void @__kmpc_free_shared(ptr {{.*}}, i64 4), !dbg ![[LOC]]
// CHECK-DAG: ![[SP]] = distinct !DISubprogram(name: "__omp_offloading_target"
diff --git a/mlir/test/Target/LLVMIR/omptarget-parallel-in-loop-shared-mem.mlir b/mlir/test/Target/LLVMIR/omptarget-parallel-in-loop-shared-mem.mlir
new file mode 100644
index 0000000000000..776676fb666b9
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/omptarget-parallel-in-loop-shared-mem.mlir
@@ -0,0 +1,51 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// Check that for a device generic-mode `omp.parallel` nested in a loop, its
+// captured-arguments aggregate is allocated in shared memory at the call site
+// (inside the loop body), balancing the free on the region-exit path, rather
+// than being hoisted to the function entry block.
+
+module attributes {omp.is_gpu = true, omp.is_target_device = true} {
+ llvm.func @parallel_in_loop(%arg0: !llvm.ptr) attributes {omp.declare_target = #omp.declaretarget<device_type = host, capture_clause = to>} {
+ %0 = omp.map.info var_ptr(%arg0 : !llvm.ptr, i32) map_clauses(from) capture(ByRef) name("d") -> !llvm.ptr
+ omp.target kernel_type(generic) map_entries(%0 -> %arg2 : !llvm.ptr) {
+ %c0 = llvm.mlir.constant(0 : i32) : i32
+ %c2 = llvm.mlir.constant(2 : i32) : i32
+ %c1 = llvm.mlir.constant(1 : i32) : i32
+ llvm.br ^bb1(%c0 : i32)
+ ^bb1(%iv: i32):
+ %cond = llvm.icmp "slt" %iv, %c2 : i32
+ llvm.cond_br %cond, ^bb2, ^bb3
+ ^bb2:
+ omp.parallel {
+ llvm.store %iv, %arg2 : i32, !llvm.ptr
+ omp.terminator
+ }
+ %next = llvm.add %iv, %c1 : i32
+ llvm.br ^bb1(%next : i32)
+ ^bb3:
+ omp.terminator
+ }
+ llvm.return
+ }
+}
+
+// CHECK-LABEL: define {{.*}} @__omp_offloading_{{.*}}parallel_in_loop
+// CHECK: user_code.entry:
+// The aggregate holding the parallel's arguments must NOT be allocated in the
+// function entry block (outside the loop); only the smaller captured buffers
+// are. The CHECK-NOT is bounded by the next block label so it guards exactly
+// the user_code.entry block.
+// CHECK-NOT: call align {{.*}} ptr @__kmpc_alloc_shared(i64 16)
+// CHECK: omp.region.after_alloca:
+// It is allocated at the parallel region's entry, inside the loop body...
+// CHECK: %[[STRUCTARG:.*]] = call align {{.*}} ptr @__kmpc_alloc_shared(i64 16)
+// CHECK: call void @__kmpc_parallel_60(
+// ...and freed on the region-exit path, also inside the loop body: the exit
+// branches back to the loop header (which performs the loop test), so the
+// alloc/free pair sits on the loop backedge and repeats every iteration.
+// CHECK: omp.par.exit:
+// CHECK: call void @__kmpc_free_shared(ptr %[[STRUCTARG]], i64 16)
+// CHECK: br label %[[LOOP_HEADER:.*]]
+// CHECK: [[LOOP_HEADER]]:
+// CHECK: icmp slt i32 {{.*}}, 2
diff --git a/mlir/test/Target/LLVMIR/omptarget-parallel-llvm.mlir b/mlir/test/Target/LLVMIR/omptarget-parallel-llvm.mlir
index db220b222f6e1..b41355830f9c2 100644
--- a/mlir/test/Target/LLVMIR/omptarget-parallel-llvm.mlir
+++ b/mlir/test/Target/LLVMIR/omptarget-parallel-llvm.mlir
@@ -59,10 +59,10 @@ module attributes {dlti.dl_spec = #dlti.dl_spec<#dlti.dl_entry<"dlti.alloca_memo
// CHECK: %[[EXEC_USER_CODE:.*]] = icmp eq i32 %[[TMP4]], -1
// CHECK: br i1 %[[EXEC_USER_CODE]], label %[[USER_CODE_ENTRY:.*]], label %[[WORKER_EXIT:.*]]
// CHECK: %[[TMP5:.*]] = addrspacecast ptr addrspace(5) %[[TMP1]] to ptr
-// CHECK: %[[STRUCTARG:.*]] = call align 8 ptr @__kmpc_alloc_shared(i64 8)
// CHECK: %[[TMP2:.*]] = call align 8 ptr @__kmpc_alloc_shared(i64 8)
// CHECK: store ptr %[[TMP0]], ptr %[[TMP2]], align 8
// CHECK: %[[TMP6:.*]] = load ptr, ptr %[[TMP2]], align 8
+// CHECK: %[[STRUCTARG:.*]] = call align 8 ptr @__kmpc_alloc_shared(i64 8)
// CHECK: %[[OMP_GLOBAL_THREAD_NUM:.*]] = call i32 @__kmpc_global_thread_num(ptr addrspacecast (ptr addrspace(1) @[[GLOB1:[0-9]+]] to ptr))
// CHECK: %[[GEP_:.*]] = getelementptr { ptr }, ptr %[[STRUCTARG]], i32 0, i32 0
// CHECK: store ptr %[[TMP6]], ptr %[[GEP_]], align 8
diff --git a/mlir/test/Target/LLVMIR/openmp-target-private-shared-mem.mlir b/mlir/test/Target/LLVMIR/openmp-target-private-shared-mem.mlir
index 775ae98ae0b51..c6085ed51b398 100644
--- a/mlir/test/Target/LLVMIR/openmp-target-private-shared-mem.mlir
+++ b/mlir/test/Target/LLVMIR/openmp-target-private-shared-mem.mlir
@@ -34,8 +34,8 @@ module attributes {omp.is_target_device = true, llvm.target_triple = "amdgcn-amd
// CHECK-LABEL: define {{.*}} void @__omp_offloading_{{.*}}target_map_single_shared_mem_private{{.*}}({{.*}})
// CHECK: call i32 @__kmpc_target_init
- // CHECK: %[[ALLOC_ARGS0:.*]] = call align 8 ptr @__kmpc_alloc_shared(i64 8)
// CHECK: %[[ALLOC1:.*]] = call align 8 ptr @__kmpc_alloc_shared(i64 4)
+ // CHECK: %[[ALLOC_ARGS0:.*]] = call align 8 ptr @__kmpc_alloc_shared(i64 8)
// CHECK: %[[GEP0:.*]] = getelementptr { ptr }, ptr %[[ALLOC_ARGS0]], i32 0, i32 0
// CHECK: store ptr %[[ALLOC1]], ptr %[[GEP0]], align 8
// CHECK: %[[GEP1:.*]] = getelementptr inbounds [1 x ptr], ptr %[[PAR_ARGS0:.*]], i64 0, i64 0
@@ -54,8 +54,8 @@ module attributes {omp.is_target_device = true, llvm.target_triple = "amdgcn-amd
// CHECK-LABEL: define {{.*}} void @__omp_offloading_{{.*}}target_map_single_shared_mem_private{{.*}}({{.*}})
// CHECK: call i32 @__kmpc_target_init
- // CHECK: %[[ALLOC_ARGS1:.*]] = call align 8 ptr @__kmpc_alloc_shared(i64 8)
// CHECK: %[[ALLOC2:.*]] = call align 8 ptr @__kmpc_alloc_shared(i64 4)
+ // CHECK: %[[ALLOC_ARGS1:.*]] = call align 8 ptr @__kmpc_alloc_shared(i64 8)
// CHECK: %[[GEP2:.*]] = getelementptr { ptr }, ptr %[[ALLOC_ARGS1]], i32 0, i32 0
// CHECK: store ptr %[[ALLOC2]], ptr %[[GEP2]], align 8
// CHECK: %[[GEP3:.*]] = getelementptr inbounds [1 x ptr], ptr %[[PAR_ARGS1:.*]], i64 0, i64 0
More information about the Mlir-commits
mailing list