[Mlir-commits] [mlir] [mlir][OpenMP] Keep loop allocas in the entry block after nested splits (PR #231270)
Sairudra More
llvmlistbot at llvm.org
Tue Oct 13 05:59:21 PDT 2026
https://github.com/Saieiei updated https://github.com/llvm/llvm-project/pull/231270
>From 359cd08015c0f9526f1b24f2a3f7c5e72604e92c Mon Sep 17 00:00:00 2001
From: Sairudra More <sairudra60 at gmail.com>
Date: Mon, 12 Oct 2026 14:43:13 -0500
Subject: [PATCH] [mlir][OpenMP] Keep loop allocas in the entry block after
nested splits
#228117 made IRBuilder insertion points plain instruction iterators. The
OpenMP translation saves an alloca insertion point and keeps using it after
nested regions are translated. When a nested construct such as `simd` or
`wsloop` allocates its private variables, it splits the alloca block, and
the saved iterator moves with its instruction into the new block. Allocas
emitted later through that point, such as the `distribute` loop body's
argument struct or a nested `simd` private variable, end up outside the
entry block. Those are no longer static allocas, and at -O0 the device code
for `target teams` with a nested `distribute simd` faults at run time on
gfx906 and gfx908.
Store the alloca block in `OpenMPAllocStackFrame` and take its first
insertion point when it is needed, so later splits cannot move it.
`convertOmpWsloop` and `convertOmpDistribute` also refresh their local
`allocaIP` the same way before calling `applyWorkshareLoop`.
---
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 14 ++-
.../LLVMIR/openmp-nested-loop-alloca.mlir | 99 +++++++++++++++++++
2 files changed, 109 insertions(+), 4 deletions(-)
create mode 100644 mlir/test/Target/LLVMIR/openmp-nested-loop-alloca.mlir
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index e8613060fee9ed..4ca0fb24722cb4 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -74,7 +74,7 @@ convertToScheduleKind(std::optional<omp::ClauseScheduleKind> schedKind) {
}
/// ModuleTranslation stack frame for OpenMP operations. This keeps track of the
-/// insertion points for allocas.
+/// blocks in which allocas are emitted.
class OpenMPAllocStackFrame
: public StateStackFrameBase<OpenMPAllocStackFrame> {
public:
@@ -83,8 +83,8 @@ class OpenMPAllocStackFrame
explicit OpenMPAllocStackFrame(
llvm::OpenMPIRBuilder::InsertPointTy allocaIP,
llvm::ArrayRef<llvm::BasicBlock *> deallocBlocks)
- : allocInsertPoint(allocaIP), deallocBlocks(deallocBlocks) {}
- llvm::OpenMPIRBuilder::InsertPointTy allocInsertPoint;
+ : allocBlock(allocaIP.getNodeParent()), deallocBlocks(deallocBlocks) {}
+ llvm::BasicBlock *allocBlock;
llvm::SmallVector<llvm::BasicBlock *> deallocBlocks;
};
@@ -666,7 +666,7 @@ static llvm::OpenMPIRBuilder::InsertPointTy findAllocInsertPoints(
llvm::ArrayRef<llvm::BasicBlock *> deallocInsertPoints;
WalkResult walkResult = moduleTranslation.stackWalk<OpenMPAllocStackFrame>(
[&](OpenMPAllocStackFrame &frame) {
- allocInsertPoint = frame.allocInsertPoint;
+ allocInsertPoint = frame.allocBlock->getFirstInsertionPt();
deallocInsertPoints = frame.deallocBlocks;
return WalkResult::interrupt();
});
@@ -4774,6 +4774,7 @@ convertOmpWsloop(Operation &opInst, llvm::IRBuilderBase &builder,
llvm::OpenMPIRBuilder::InsertPointTy allocaIP =
findAllocInsertPoints(builder, moduleTranslation);
+ llvm::BasicBlock *allocaBlock = allocaIP.getNodeParent();
SmallVector<llvm::Value *> privateReductionVariables(
wsloopOp.getNumReductionVars());
@@ -4907,6 +4908,8 @@ convertOmpWsloop(Operation &opInst, llvm::IRBuilderBase &builder,
linearClauseProcessor.rewriteInPlace(builder, loopInfo->getBody(),
loopInfo->getLatch(), index);
+ // Nested regions can move the instruction referenced by allocaIP.
+ allocaIP = allocaBlock->getFirstInsertionPt();
llvm::OpenMPIRBuilder::InsertPointOrErrorTy wsloopIP =
ompBuilder->applyWorkshareLoop(
ompLoc.DL, loopInfo, allocaIP, loopNeedsBarrier,
@@ -8822,6 +8825,7 @@ convertOmpDistribute(Operation &opInst, llvm::IRBuilderBase &builder,
auto bodyGenCB =
[&](InsertPointTy allocaIP, InsertPointTy codeGenIP,
llvm::ArrayRef<llvm::BasicBlock *> deallocBlocks) -> llvm::Error {
+ llvm::BasicBlock *allocaBlock = allocaIP.getNodeParent();
// DistributeOp has only one region associated with it.
builder.restoreIP(codeGenIP);
PrivateVarsInfo privVarsInfo(distributeOp);
@@ -8877,6 +8881,8 @@ convertOmpDistribute(Operation &opInst, llvm::IRBuilderBase &builder,
distributeOp.getDistScheduleChunkSize());
llvm::CanonicalLoopInfo *loopInfo =
findCurrentLoopInfo(moduleTranslation);
+ // Nested regions can move the instruction referenced by allocaIP.
+ allocaIP = allocaBlock->getFirstInsertionPt();
llvm::OpenMPIRBuilder::InsertPointOrErrorTy wsloopIP =
ompBuilder->applyWorkshareLoop(
ompLoc.DL, loopInfo, allocaIP, loopNeedsBarrier,
diff --git a/mlir/test/Target/LLVMIR/openmp-nested-loop-alloca.mlir b/mlir/test/Target/LLVMIR/openmp-nested-loop-alloca.mlir
new file mode 100644
index 00000000000000..7ab0e288ab16c1
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-nested-loop-alloca.mlir
@@ -0,0 +1,99 @@
+// RUN: mlir-translate --mlir-to-llvmir --split-input-file %s | FileCheck %s
+
+// Keep private variables and loop callback arguments in the entry block when
+// nested SIMD regions split the allocation block.
+
+module attributes {
+ dlti.dl_spec = #dlti.dl_spec<"dlti.alloca_memory_space" = 5 : ui64>,
+ llvm.target_triple = "amdgcn-amd-amdhsa",
+ omp.is_gpu = true,
+ omp.is_target_device = true
+} {
+ omp.private {type = private} @private_i : i32
+
+ llvm.func @distribute_simd(%array: !llvm.ptr, %n: i32) attributes {
+ omp.declare_target = #omp.declaretarget<device_type = nohost, capture_clause = to>
+ } {
+ %one = llvm.mlir.constant(1 : i32) : i32
+ %i = llvm.alloca %one x i32 : (i32) -> !llvm.ptr<5>
+ %i_cast = llvm.addrspacecast %i : !llvm.ptr<5> to !llvm.ptr
+ omp.teams {
+ omp.distribute {
+ omp.simd private(@private_i %i_cast -> %private_i : !llvm.ptr) {
+ omp.loop_nest (%iv) : i32 = (%one) to (%n) inclusive step (%one) {
+ llvm.store %iv, %private_i : i32, !llvm.ptr
+ %offset = llvm.load %private_i : !llvm.ptr -> i32
+ %element = llvm.getelementptr %array[%offset] : (!llvm.ptr, i32) -> !llvm.ptr, i32
+ llvm.store %iv, %element : i32, !llvm.ptr
+ omp.yield
+ }
+ } {omp.composite}
+ } {omp.composite}
+ omp.terminator
+ }
+ llvm.return
+ }
+}
+
+// CHECK-LABEL: define internal void @distribute_simd..omp_par.1(
+// CHECK: distribute.alloca:
+// CHECK-NOT: br
+// CHECK: %[[ARGS:.*]] = alloca { ptr, ptr }, align 8, addrspace(5)
+// CHECK: %[[ARGS_CAST:.*]] = addrspacecast ptr addrspace(5) %[[ARGS]] to ptr
+// CHECK: br label %omp.region.after_alloca
+// CHECK-NOT: = alloca {
+// CHECK: call void @__kmpc_distribute_static_loop_4u({{.*}}, ptr %[[ARGS_CAST]],
+// CHECK: ret void
+
+// -----
+
+module attributes {
+ dlti.dl_spec = #dlti.dl_spec<"dlti.alloca_memory_space" = 5 : ui64>,
+ llvm.target_triple = "amdgcn-amd-amdhsa",
+ omp.is_gpu = true,
+ omp.is_target_device = true
+} {
+ omp.private {type = private} @private_i : i32
+
+ llvm.func @distribute_parallel_simd(%array: !llvm.ptr, %n: i32) attributes {
+ omp.declare_target = #omp.declaretarget<device_type = nohost, capture_clause = to>
+ } {
+ %one = llvm.mlir.constant(1 : i32) : i32
+ %i = llvm.alloca %one x i32 : (i32) -> !llvm.ptr<5>
+ %i_cast = llvm.addrspacecast %i : !llvm.ptr<5> to !llvm.ptr
+ omp.teams {
+ omp.parallel {
+ omp.distribute {
+ omp.wsloop {
+ omp.simd private(@private_i %i_cast -> %private_i : !llvm.ptr) {
+ omp.loop_nest (%iv) : i32 = (%one) to (%n) inclusive step (%one) {
+ llvm.store %iv, %private_i : i32, !llvm.ptr
+ %offset = llvm.load %private_i : !llvm.ptr -> i32
+ %element = llvm.getelementptr %array[%offset] : (!llvm.ptr, i32) -> !llvm.ptr, i32
+ llvm.store %iv, %element : i32, !llvm.ptr
+ omp.yield
+ }
+ } {omp.composite}
+ } {omp.composite}
+ } {omp.composite}
+ omp.terminator
+ } {omp.composite}
+ omp.terminator
+ }
+ llvm.return
+ }
+}
+
+// CHECK-LABEL: define internal void @distribute_parallel_simd..omp_par.1(
+// CHECK: distribute.alloca:
+// CHECK-NOT: br
+// CHECK: %[[ARGS:.*]] = alloca { ptr, ptr }, align 8, addrspace(5)
+// CHECK: %[[ARGS_CAST:.*]] = addrspacecast ptr addrspace(5) %[[ARGS]] to ptr
+// CHECK-NOT: br
+// CHECK: %[[PRIVATE:.*]] = alloca i32, align 4, addrspace(5)
+// CHECK: %[[PRIVATE_CAST:.*]] = addrspacecast ptr addrspace(5) %[[PRIVATE]] to ptr
+// CHECK: br label %omp.region.after_alloca
+// CHECK-NOT: = alloca
+// CHECK: store ptr %[[PRIVATE_CAST]], ptr addrspace(5)
+// CHECK: call void @__kmpc_distribute_for_static_loop_4u({{.*}}, ptr %[[ARGS_CAST]],
+// CHECK: ret void
More information about the Mlir-commits
mailing list