[Mlir-commits] [mlir] ba833b4 - [mlir][OpenMP] Keep loop allocas in the entry block after nested splits (#231270)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Oct 14 00:06:37 PDT 2026
Author: Sairudra More
Date: 2026-10-14T12:36:31+05:30
New Revision: ba833b45f892897663c09e91c1e77cf3d396ce19
URL: https://github.com/llvm/llvm-project/commit/ba833b45f892897663c09e91c1e77cf3d396ce19
DIFF: https://github.com/llvm/llvm-project/commit/ba833b45f892897663c09e91c1e77cf3d396ce19.diff
LOG: [mlir][OpenMP] Keep loop allocas in the entry block after nested splits (#231270)
#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`.
Assisted-by: GitHub Copilot
Added:
mlir/test/Target/LLVMIR/openmp-nested-loop-alloca.mlir
Modified:
mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
Removed:
################################################################################
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index e8613060fee9e..6d8185ed441bd 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -74,17 +74,17 @@ 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:
MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(OpenMPAllocStackFrame)
explicit OpenMPAllocStackFrame(
- llvm::OpenMPIRBuilder::InsertPointTy allocaIP,
+ llvm::BasicBlock *allocBlock,
llvm::ArrayRef<llvm::BasicBlock *> deallocBlocks)
- : allocInsertPoint(allocaIP), deallocBlocks(deallocBlocks) {}
- llvm::OpenMPIRBuilder::InsertPointTy allocInsertPoint;
+ : allocBlock(allocBlock), 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();
});
@@ -2670,7 +2670,7 @@ convertOmpTeams(omp::TeamsOp op, llvm::IRBuilderBase &builder,
auto bodyCB = [&](InsertPointTy allocaIP, InsertPointTy codegenIP,
llvm::ArrayRef<llvm::BasicBlock *> deallocBlocks) {
LLVM::ModuleTranslation::SaveStack<OpenMPAllocStackFrame> frame(
- moduleTranslation, allocaIP, deallocBlocks);
+ moduleTranslation, allocaIP.getNodeParent(), deallocBlocks);
builder.restoreIP(codegenIP);
return convertOmpOpRegions(op.getRegion(), "omp.teams.region", builder,
moduleTranslation)
@@ -3374,7 +3374,7 @@ convertOmpTaskOp(omp::TaskOp taskOp, llvm::IRBuilderBase &builder,
// Save the alloca insertion point on ModuleTranslation stack for use in
// nested regions.
LLVM::ModuleTranslation::SaveStack<OpenMPAllocStackFrame> frame(
- moduleTranslation, allocaIP, deallocBlocks);
+ moduleTranslation, allocaIP.getNodeParent(), deallocBlocks);
// Allocate and initialize private variables
builder.SetInsertPoint(initBlock->getTerminator());
@@ -3465,7 +3465,7 @@ convertOmpTaskOp(omp::TaskOp taskOp, llvm::IRBuilderBase &builder,
// Save the alloca insertion point on ModuleTranslation stack for use in
// nested regions.
LLVM::ModuleTranslation::SaveStack<OpenMPAllocStackFrame> frame(
- moduleTranslation, allocaIP, deallocBlocks);
+ moduleTranslation, allocaIP.getNodeParent(), deallocBlocks);
// translate the body of the task:
builder.restoreIP(codegenIP);
@@ -3856,7 +3856,7 @@ convertOmpTaskloopContextOp(omp::TaskloopContextOp contextOp,
splitBB(builder, /*CreateBranch=*/true, "omp.private.init");
LLVM::ModuleTranslation::SaveStack<OpenMPAllocStackFrame> frame(
- moduleTranslation, allocaIP, deallocBlocks);
+ moduleTranslation, allocaIP.getNodeParent(), deallocBlocks);
// Allocate and initialize private variables
builder.SetInsertPoint(initBlock->getTerminator());
@@ -3984,7 +3984,7 @@ convertOmpTaskloopContextOp(omp::TaskloopContextOp contextOp,
// Save the alloca insertion point on ModuleTranslation stack for use in
// nested regions.
LLVM::ModuleTranslation::SaveStack<OpenMPAllocStackFrame> frame(
- moduleTranslation, allocaIP, deallocBlocks);
+ moduleTranslation, allocaIP.getNodeParent(), deallocBlocks);
// translate the body of the taskloop:
builder.restoreIP(codegenIP);
@@ -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,
@@ -5046,7 +5049,7 @@ convertOmpParallel(omp::ParallelOp opInst, llvm::IRBuilderBase &builder,
// Save the alloca insertion point on ModuleTranslation stack for use in
// nested regions.
LLVM::ModuleTranslation::SaveStack<OpenMPAllocStackFrame> frame(
- moduleTranslation, allocaIP, deallocBlocks);
+ moduleTranslation, allocaIP.getNodeParent(), deallocBlocks);
// ParallelOp has only one region associated with it.
llvm::Expected<llvm::BasicBlock *> regionBlock = convertOmpOpRegions(
@@ -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);
@@ -8835,7 +8839,7 @@ convertOmpDistribute(Operation &opInst, llvm::IRBuilderBase &builder,
// nested regions. This must be after allocatePrivateVars, which splits
// the alloca block and updates allocaIP.
LLVM::ModuleTranslation::SaveStack<OpenMPAllocStackFrame> frame(
- moduleTranslation, allocaIP, deallocBlocks);
+ moduleTranslation, allocaIP.getNodeParent(), deallocBlocks);
if (handleError(initPrivateVars(builder, moduleTranslation, privVarsInfo),
opInst)
@@ -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,
@@ -9764,7 +9770,7 @@ convertOmpTarget(Operation &opInst, llvm::IRBuilderBase &builder,
}
LLVM::ModuleTranslation::SaveStack<OpenMPAllocStackFrame> frame(
- moduleTranslation, allocaIP, deallocBlocks);
+ moduleTranslation, allocaIP.getNodeParent(), deallocBlocks);
llvm::Expected<llvm::BasicBlock *> exitBlock = convertOmpOpRegions(
targetRegion, "omp.target", builder, moduleTranslation);
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 0000000000000..7ab0e288ab16c
--- /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