[llvm] [mlir] [OpenMP][OMPIRBuilder] Fix lastiter assertion in target workshare loop (PR #214996)
via llvm-commits
llvm-commits at lists.llvm.org
Tue Aug 11 00:13:24 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/2] [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/2] [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]]
More information about the llvm-commits
mailing list