[Mlir-commits] [llvm] [mlir] [OpenMP][OMPIRBuilder] Fix lastiter assertion in target workshare loop (PR #214996)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Aug 14 06:51:25 PDT 2026
https://github.com/RohithPariki updated https://github.com/llvm/llvm-project/pull/214996
>From c282049a8f9d2b63a3aecef9262e1ddf6087095b Mon Sep 17 00:00:00 2001
From: Rohith Pariki <rohithpariki at gmail.com>
Date: Sat, 8 Aug 2026 18:36:57 +0530
Subject: [PATCH 1/3] [OpenMP][OMPIRBuilder] Fix lastiter assertion in target
workshare loop
When translating an omp.wsloop with a linear clause for a target device, applyWorkshareLoopTarget failed to set the lastiter value in CanonicalLoopInfo, resulting in an assertion failure when finalizing linear variables. This patch allocates and sets lastiter to true, as target workshare loops execute synchronously on the device and complete entirely before returning.
---
llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 9 +++++++++
1 file changed, 9 insertions(+)
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 5a363d0ac3dbd..173cbddf2a4d6 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -6577,6 +6577,15 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::applyWorkshareLoopTarget(
}
Value *Ident = getOrCreateIdent(SrcLocStr, SrcLocStrSize, Flag);
+ // Allocate p.lastiter and set it to 1 (true). The target workshare loop
+ // executes synchronously on the device and completely finishes before
+ // returning, so the thread executing after it is effectively the one that
+ // executed the last iteration, and needs to do the linear variable updates.
+ Builder.restoreIP(AllocaIP);
+ Value *PLastIter = Builder.CreateAlloca(Builder.getInt32Ty(), nullptr, "p.lastiter");
+ Builder.CreateStore(Builder.getInt32(1), PLastIter);
+ CLI->setLastIter(PLastIter);
+
auto OI = std::make_unique<OutlineInfo>();
OI->OuterAllocBB = CLI->getPreheader();
Function *OuterFn = CLI->getPreheader()->getParent();
>From a91887941045df51abefe5312d67a8b0486a65d9 Mon Sep 17 00:00:00 2001
From: Rohith Pariki <rohithpariki at gmail.com>
Date: Tue, 11 Aug 2026 12:40:47 +0530
Subject: [PATCH 2/3] [OpenMP][OMPIRBuilder] Fix lastiter computation for
target workshare loops
Address reviewer feedback by ensuring lastiter is evaluated based on the iteration counter within the outlined loop body rather than unconditionally setting it to 1. This prevents race conditions where multiple threads executing the target kernel would otherwise attempt to perform linear variable updates.
Added a test using the reproducer from #213905.
---
llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 18 ++++++++---
.../LLVMIR/openmp-target-wsloop-linear.mlir | 32 +++++++++++++++++++
2 files changed, 45 insertions(+), 5 deletions(-)
create mode 100644 mlir/test/Target/LLVMIR/openmp-target-wsloop-linear.mlir
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 173cbddf2a4d6..a1c04c3c57ed1 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -6577,13 +6577,12 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::applyWorkshareLoopTarget(
}
Value *Ident = getOrCreateIdent(SrcLocStr, SrcLocStrSize, Flag);
- // Allocate p.lastiter and set it to 1 (true). The target workshare loop
- // executes synchronously on the device and completely finishes before
- // returning, so the thread executing after it is effectively the one that
- // executed the last iteration, and needs to do the linear variable updates.
+ // Allocate p.lastiter and initialize it to 0.
+ // The actual value will be set inside the loop body if the current iteration
+ // is the last one.
Builder.restoreIP(AllocaIP);
Value *PLastIter = Builder.CreateAlloca(Builder.getInt32Ty(), nullptr, "p.lastiter");
- Builder.CreateStore(Builder.getInt32(1), PLastIter);
+ Builder.CreateStore(Builder.getInt32(0), PLastIter);
CLI->setLastIter(PLastIter);
auto OI = std::make_unique<OutlineInfo>();
@@ -6614,6 +6613,15 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::applyWorkshareLoopTarget(
ToBeDeleted.push_back(NewLoopCntLoad);
ToBeDeleted.push_back(NewLoopCnt);
+ // Set p.lastiter to 1 if the current iteration is the last one.
+ Builder.restoreIP(CLI->getBody(), CLI->getBody()->getFirstInsertionPt());
+ Value *IsLast = Builder.CreateICmpEQ(
+ NewLoopCntLoad,
+ Builder.CreateSub(CLI->getTripCount(),
+ ConstantInt::get(CLI->getTripCount()->getType(), 1)));
+ Value *IsLastExt = Builder.CreateZExt(IsLast, Builder.getInt32Ty());
+ Builder.CreateStore(IsLastExt, PLastIter);
+
// Analyse loop body region. Find all input variables which are used inside
// loop body region.
SmallPtrSet<BasicBlock *, 32> ParallelRegionBlockSet;
diff --git a/mlir/test/Target/LLVMIR/openmp-target-wsloop-linear.mlir b/mlir/test/Target/LLVMIR/openmp-target-wsloop-linear.mlir
new file mode 100644
index 0000000000000..81574e9089f1e
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/openmp-target-wsloop-linear.mlir
@@ -0,0 +1,32 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+module attributes {llvm.target_triple = "amdgcn-amd-amdhsa",
+ omp.is_gpu = true, omp.is_target_device = true} {
+ llvm.func @wsloop_linear_target(%x : !llvm.ptr) attributes {
+ omp.declare_target = #omp.declaretarget<device_type = (any), capture_clause = (to)>
+ } {
+ %lb = llvm.mlir.constant(0 : i32) : i32
+ %ub = llvm.mlir.constant(9 : i32) : i32
+ %step = llvm.mlir.constant(1 : i32) : i32
+ omp.wsloop linear(%x : !llvm.ptr = %step : i32) {
+ omp.loop_nest (%iv) : i32 = (%lb) to (%ub) inclusive step (%step) {
+ omp.yield
+ }
+ } {linear_var_types = [i32]}
+ llvm.return
+ }
+}
+
+// CHECK-LABEL: define {{(protected )?}}void @wsloop_linear_target
+// CHECK: %[[P_LASTITER:.*]] = alloca i32
+// CHECK: store i32 0, ptr %[[P_LASTITER]]
+// CHECK: call void @__kmpc_for_static_loop_4u({{.*}}, ptr @wsloop_linear_target..omp_wsloop
+// CHECK: %[[LASTITER_VAL:.*]] = load i32, ptr %[[P_LASTITER]]
+// CHECK: %[[IS_LASTITER:.*]] = icmp ne i32 %[[LASTITER_VAL]], 0
+// CHECK: br i1 %[[IS_LASTITER]], label %[[LINEAR_UPDATE_BLOCK:.*]], label %[[LINEAR_SKIP_BLOCK:.*]]
+
+// CHECK-LABEL: define internal void @wsloop_linear_target..omp_wsloop(
+// CHECK-SAME: ptr %[[LASTITER_ARG:.*]], i32 %[[LOOP_CNT_ARG:.*]])
+// CHECK: %[[CMP:.*]] = icmp eq i32 %[[LOOP_CNT_ARG]], 9
+// CHECK: %[[ZEXT:.*]] = zext i1 %[[CMP]] to i32
+// CHECK: store i32 %[[ZEXT]], ptr %[[LASTITER_ARG]]
>From f25769d522b653825d78cda9654cc77be5308da7 Mon Sep 17 00:00:00 2001
From: Rohith Pariki <rohithpariki at gmail.com>
Date: Fri, 14 Aug 2026 18:40:10 +0530
Subject: [PATCH 3/3] [OpenMP][OMPIRBuilder] Conditionally emit lastiter for
target workshare loops
---
.../llvm/Frontend/OpenMP/OMPIRBuilder.h | 5 ++-
llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp | 43 +++++++++++--------
.../OpenMP/OpenMPToLLVMIRTranslation.cpp | 3 +-
3 files changed, 30 insertions(+), 21 deletions(-)
diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
index 1965f7b983805..9e9fd4cc0dd0d 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
@@ -1178,7 +1178,8 @@ class OpenMPIRBuilder {
InsertPointTy applyWorkshareLoopTarget(DebugLoc DL, CanonicalLoopInfo *CLI,
InsertPointTy AllocaIP,
omp::WorksharingLoopType LoopType,
- bool NoLoop);
+ bool NoLoop,
+ bool HasLastiterClause = false);
/// Modifies the canonical loop to be a statically-scheduled workshare loop.
///
@@ -1350,7 +1351,7 @@ class OpenMPIRBuilder {
omp::WorksharingLoopType LoopType =
omp::WorksharingLoopType::ForStaticLoop,
bool NoLoop = false, bool HasDistSchedule = false,
- Value *DistScheduleChunkSize = nullptr);
+ Value *DistScheduleChunkSize = nullptr, bool HasLastiterClause = false);
/// Tile a loop nest.
///
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index a1c04c3c57ed1..9a9ebd6dc72e3 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -6560,7 +6560,7 @@ static void workshareLoopTargetCallback(
OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::applyWorkshareLoopTarget(
DebugLoc DL, CanonicalLoopInfo *CLI, InsertPointTy AllocaIP,
- WorksharingLoopType LoopType, bool NoLoop) {
+ WorksharingLoopType LoopType, bool NoLoop, bool HasLastiterClause) {
uint32_t SrcLocStrSize;
Constant *SrcLocStr = getOrCreateSrcLocStr(DL, SrcLocStrSize);
IdentFlag Flag = IdentFlag(0);
@@ -6577,13 +6577,17 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::applyWorkshareLoopTarget(
}
Value *Ident = getOrCreateIdent(SrcLocStr, SrcLocStrSize, Flag);
- // Allocate p.lastiter and initialize it to 0.
- // The actual value will be set inside the loop body if the current iteration
- // is the last one.
- Builder.restoreIP(AllocaIP);
- Value *PLastIter = Builder.CreateAlloca(Builder.getInt32Ty(), nullptr, "p.lastiter");
- Builder.CreateStore(Builder.getInt32(0), PLastIter);
- CLI->setLastIter(PLastIter);
+ Value *PLastIter = nullptr;
+ if (HasLastiterClause) {
+ // Allocate p.lastiter and initialize it to 0.
+ // The actual value will be set inside the loop body if the current
+ // iteration is the last one.
+ Builder.restoreIP(AllocaIP);
+ PLastIter =
+ Builder.CreateAlloca(Builder.getInt32Ty(), nullptr, "p.lastiter");
+ Builder.CreateStore(Builder.getInt32(0), PLastIter);
+ CLI->setLastIter(PLastIter);
+ }
auto OI = std::make_unique<OutlineInfo>();
OI->OuterAllocBB = CLI->getPreheader();
@@ -6613,14 +6617,16 @@ OpenMPIRBuilder::InsertPointTy OpenMPIRBuilder::applyWorkshareLoopTarget(
ToBeDeleted.push_back(NewLoopCntLoad);
ToBeDeleted.push_back(NewLoopCnt);
- // Set p.lastiter to 1 if the current iteration is the last one.
- Builder.restoreIP(CLI->getBody(), CLI->getBody()->getFirstInsertionPt());
- Value *IsLast = Builder.CreateICmpEQ(
- NewLoopCntLoad,
- Builder.CreateSub(CLI->getTripCount(),
- ConstantInt::get(CLI->getTripCount()->getType(), 1)));
- Value *IsLastExt = Builder.CreateZExt(IsLast, Builder.getInt32Ty());
- Builder.CreateStore(IsLastExt, PLastIter);
+ if (HasLastiterClause) {
+ // Set p.lastiter to 1 if the current iteration is the last one.
+ Builder.restoreIP({CLI->getBody(), CLI->getBody()->getFirstInsertionPt()});
+ Value *IsLast = Builder.CreateICmpEQ(
+ NewLoopCntLoad,
+ Builder.CreateSub(CLI->getTripCount(),
+ ConstantInt::get(CLI->getTripCount()->getType(), 1)));
+ Value *IsLastExt = Builder.CreateZExt(IsLast, Builder.getInt32Ty());
+ Builder.CreateStore(IsLastExt, PLastIter);
+ }
// Analyse loop body region. Find all input variables which are used inside
// loop body region.
@@ -6685,9 +6691,10 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::applyWorkshareLoop(
bool HasSimdModifier, bool HasMonotonicModifier,
bool HasNonmonotonicModifier, bool HasOrderedClause,
WorksharingLoopType LoopType, bool NoLoop, bool HasDistSchedule,
- Value *DistScheduleChunkSize) {
+ Value *DistScheduleChunkSize, bool HasLastiterClause) {
if (Config.isTargetDevice())
- return applyWorkshareLoopTarget(DL, CLI, AllocaIP, LoopType, NoLoop);
+ return applyWorkshareLoopTarget(DL, CLI, AllocaIP, LoopType, NoLoop,
+ HasLastiterClause);
OMPScheduleType EffectiveScheduleType = computeOpenMPScheduleType(
SchedKind, ChunkSize, HasSimdModifier, HasMonotonicModifier,
HasNonmonotonicModifier, HasOrderedClause, DistScheduleChunkSize);
diff --git a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
index 2c81c9231a9b7..9f3378560cf62 100644
--- a/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
+++ b/mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp
@@ -4811,7 +4811,8 @@ convertOmpWsloop(Operation &opInst, llvm::IRBuilderBase &builder,
convertToScheduleKind(schedule), chunk, isSimd,
scheduleMod == omp::ScheduleModifier::monotonic,
scheduleMod == omp::ScheduleModifier::nonmonotonic, isOrdered,
- workshareLoopType, noLoopMode, hasDistSchedule, distScheduleChunk);
+ workshareLoopType, noLoopMode, hasDistSchedule, distScheduleChunk,
+ /*HasLastiterClause=*/!wsloopOp.getLinearVars().empty());
if (failed(handleError(wsloopIP, opInst)))
return failure();
More information about the Mlir-commits
mailing list