[llvm] [mlir] [OpenMPIRBuilder] Serialize lanes around critical on the device (PR #215009)
via llvm-commits
llvm-commits at lists.llvm.org
Sat Aug 8 12:15:16 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-llvm
Author: Spencer Bryngelson (sbryngelson)
<details>
<summary>Changes</summary>
`critical` inside a target region does not serialize the lanes of a wavefront in flang, so a conforming program silently gets wrong results. The C equivalent is correct on the same GPU. Fixes #<!-- -->214965.
`setCriticalLock` in the DeviceRTL takes the lock on the lowest active lane of a wavefront only. That is sound if exactly one lane per wavefront reaches it, which clang arranges: `CGOpenMPRuntimeGPU::emitCriticalRegion` wraps the region in a loop over `__kmpc_get_hardware_num_threads_in_block()`, lets only the matching thread enter, and calls `__kmpc_syncwarp` between turns, before delegating to the generic emission for the lock itself. `OpenMPIRBuilder::createCritical`, which flang reaches through `convertOmpCritical`, has no device path and emits `__kmpc_critical` / body / `__kmpc_end_critical` directly, so every active lane enters at once and all but one update is lost.
This adds the same turn loop to `createCritical`:
```
mask = __kmpc_warp_active_thread_mask();
for (i = 0; i < __kmpc_get_hardware_num_threads_in_block(); ++i) {
if (__kmpc_get_hardware_thread_id_in_block() == i)
<__kmpc_critical / body / __kmpc_end_critical>
__kmpc_syncwarp(mask);
}
```
Two details worth calling out for review.
The guard is `Config.IsGPU` rather than `Config.isTargetDevice()`. `IsGPU` is set only for AMDGPU and NVPTX, which is the scope clang applies `CGOpenMPRuntimeGPU` to, and the loop is only meaningful where lanes share a program counter. It is read as `IsGPU.value_or(false)` because the getter asserts when the optional is unset, which is the case for callers that do not configure the builder.
`EmitOMPInlinedRegion` relocates the exit call but not the entry call: `emitCommonDirectiveExit` moves `ExitCall` into the finalization block, while `emitCommonDirectiveEntry` returns immediately when `Conditional` is false and leaves `EntryCall` where it was built. The entry call therefore has to be moved into the region block explicitly. Without that the lock is taken once before the loop and released on every turn, which serializes partially and gets worse as the number of wavefronts grows.
### Testing
A `target parallel do` that increments a mapped scalar inside `critical`, on gfx90a. The count is the number of wavefronts before the patch and the number of threads after, matching both the `atomic update` control in the same construct and the C equivalent.
| threads | before | after |
|---|---|---|
| 64 | 1 | 64 |
| 128 | 2 | 128 |
| 256 | 4 | 256 |
| 512 | 8 | 512 |
| 1024 | 16 | 1024 |
New test `mlir/test/Target/LLVMIR/omptarget-critical-device.mlir` checks the shape of the loop and that the lock is acquired and released inside the region a single thread enters.
`llvm/unittests/Frontend`, `mlir/test/Target/LLVMIR`, `mlir/test/Dialect/OpenMP`, `flang/test/Lower/OpenMP`, `flang/test/Integration/OpenMP` and `clang/test/OpenMP` all pass.
The root cause analysis, the reduced test cases and this fix were produced with Claude; I reviewed and verified them.
---
Full diff: https://github.com/llvm/llvm-project/pull/215009.diff
3 Files Affected:
- (modified) llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h (+7)
- (modified) llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp (+87)
- (added) mlir/test/Target/LLVMIR/omptarget-critical-device.mlir (+46)
``````````diff
diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
index 1965f7b983805..0d33c9baaa903 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPIRBuilder.h
@@ -3272,6 +3272,13 @@ class OpenMPIRBuilder {
/// \param HintInst Hint Instruction for hint clause associated with critical
///
/// \returns The insertion position *after* the critical.
+ /// Emit a critical region once per thread of the block on a target device, so
+ /// only one lane of a wavefront is inside it at a time.
+ InsertPointOrErrorTy emitDeviceSerializedCritical(Instruction *EntryCall,
+ Instruction *ExitCall,
+ BodyGenCallbackTy BodyGenCB,
+ FinalizeCallbackTy FiniCB);
+
LLVM_ABI InsertPointOrErrorTy createCritical(const LocationDescription &Loc,
BodyGenCallbackTy BodyGenCB,
FinalizeCallbackTy FiniCB,
diff --git a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
index 5a363d0ac3dbd..b46298300f129 100644
--- a/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
+++ b/llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp
@@ -8050,10 +8050,97 @@ OpenMPIRBuilder::InsertPointOrErrorTy OpenMPIRBuilder::createCritical(
getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_end_critical);
Instruction *ExitCall = createRuntimeFunctionCall(ExitRTLFn, Args);
+ // On a GPU the runtime lock is acquired by a single lane of a wavefront, so
+ // every other lane would enter the region unsynchronized. Give each thread of
+ // the block its own turn, as clang does in
+ // CGOpenMPRuntimeGPU::emitCriticalRegion, and let the lock serialize the
+ // wavefronts against each other.
+ if (Config.IsGPU.value_or(false))
+ return emitDeviceSerializedCritical(EntryCall, ExitCall, BodyGenCB, FiniCB);
+
return EmitOMPInlinedRegion(OMPD, EntryCall, ExitCall, BodyGenCB, FiniCB,
/*Conditional*/ false, /*hasFinalize*/ true);
}
+/// Emit the critical region once per thread of the block, so only one lane of a
+/// wavefront is inside it at a time:
+///
+/// mask = __kmpc_warp_active_thread_mask();
+/// for (i = 0; i < __kmpc_get_hardware_num_threads_in_block(); ++i) {
+/// if (__kmpc_get_hardware_thread_id_in_block() == i)
+/// <critical region>
+/// __kmpc_syncwarp(mask);
+/// }
+OpenMPIRBuilder::InsertPointOrErrorTy
+OpenMPIRBuilder::emitDeviceSerializedCritical(Instruction *EntryCall,
+ Instruction *ExitCall,
+ BodyGenCallbackTy BodyGenCB,
+ FinalizeCallbackTy FiniCB) {
+ Function *CurFn = Builder.GetInsertBlock()->getParent();
+ LLVMContext &Ctx = CurFn->getContext();
+ Type *I32 = Type::getInt32Ty(Ctx);
+
+ // The mask must describe the lanes that are active here; an all-ones mask
+ // makes the reconvergence wait on lanes that are not present.
+ Value *Mask = createRuntimeFunctionCall(
+ getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_warp_active_thread_mask), {});
+ Value *TId = createRuntimeFunctionCall(
+ getOrCreateRuntimeFunctionPtr(
+ OMPRTL___kmpc_get_hardware_thread_id_in_block),
+ {});
+ Value *NumThreads = createRuntimeFunctionCall(
+ getOrCreateRuntimeFunctionPtr(
+ OMPRTL___kmpc_get_hardware_num_threads_in_block),
+ {});
+
+ // The insertion block has no terminator yet, so use the OMPIRBuilder
+ // splitter rather than BasicBlock::splitBasicBlock.
+ BasicBlock *ExitBB =
+ splitBB(Builder, /*CreateBranch=*/false, "omp.critical.serial.exit");
+ BasicBlock *EntryBB = Builder.GetInsertBlock();
+ BasicBlock *HeaderBB =
+ BasicBlock::Create(Ctx, "omp.critical.serial.header", CurFn, ExitBB);
+ BasicBlock *TurnBB =
+ BasicBlock::Create(Ctx, "omp.critical.serial.turn", CurFn, ExitBB);
+ BasicBlock *RegionBB =
+ BasicBlock::Create(Ctx, "omp.critical.serial.region", CurFn, ExitBB);
+ BasicBlock *SyncBB =
+ BasicBlock::Create(Ctx, "omp.critical.serial.sync", CurFn, ExitBB);
+
+ Builder.SetInsertPoint(EntryBB);
+ Builder.CreateBr(HeaderBB);
+
+ Builder.SetInsertPoint(HeaderBB);
+ PHINode *Turn = Builder.CreatePHI(I32, 2, "omp.critical.turn");
+ Turn->addIncoming(ConstantInt::get(I32, 0), EntryBB);
+ Builder.CreateCondBr(Builder.CreateICmpSLT(Turn, NumThreads), TurnBB, ExitBB);
+
+ Builder.SetInsertPoint(TurnBB);
+ Builder.CreateCondBr(Builder.CreateICmpEQ(TId, Turn), RegionBB, SyncBB);
+
+ Builder.SetInsertPoint(RegionBB);
+ Builder.CreateBr(SyncBB);
+ Builder.SetInsertPoint(RegionBB->getTerminator());
+ // EmitOMPInlinedRegion only relocates the exit call; the entry call stays
+ // where it was built, which is before the loop.
+ EntryCall->moveBefore(*RegionBB, RegionBB->getTerminator()->getIterator());
+ InsertPointOrErrorTy AfterIP = EmitOMPInlinedRegion(
+ Directive::OMPD_critical, EntryCall, ExitCall, BodyGenCB, FiniCB,
+ /*Conditional*/ false, /*hasFinalize*/ true);
+ if (!AfterIP)
+ return AfterIP.takeError();
+
+ Builder.SetInsertPoint(SyncBB);
+ createRuntimeFunctionCall(
+ getOrCreateRuntimeFunctionPtr(OMPRTL___kmpc_syncwarp), {Mask});
+ Value *Next = Builder.CreateAdd(Turn, ConstantInt::get(I32, 1));
+ Turn->addIncoming(Next, SyncBB);
+ Builder.CreateBr(HeaderBB);
+
+ Builder.SetInsertPoint(ExitBB, ExitBB->getFirstInsertionPt());
+ return Builder.saveIP();
+}
+
OpenMPIRBuilder::InsertPointTy
OpenMPIRBuilder::createOrderedDepend(const LocationDescription &Loc,
InsertPointTy AllocaIP, unsigned NumLoops,
diff --git a/mlir/test/Target/LLVMIR/omptarget-critical-device.mlir b/mlir/test/Target/LLVMIR/omptarget-critical-device.mlir
new file mode 100644
index 0000000000000..c3466fd4d7c72
--- /dev/null
+++ b/mlir/test/Target/LLVMIR/omptarget-critical-device.mlir
@@ -0,0 +1,46 @@
+// RUN: mlir-translate -mlir-to-llvmir %s | FileCheck %s
+
+// On a GPU the runtime lock is taken by one lane of a wavefront, so the region
+// must additionally be given to one thread of the block at a time.
+
+module attributes {llvm.target_triple = "amdgcn-amd-amdhsa", omp.is_gpu = true, omp.is_target_device = true} {
+ llvm.func @critical_device(%x : !llvm.ptr, %xval : i32) attributes {omp.declare_target = #omp.declaretarget<device_type = (nohost), capture_clause = (to)>} {
+ omp.critical {
+ llvm.store %xval, %x : i32, !llvm.ptr
+ omp.terminator
+ }
+ llvm.return
+ }
+}
+
+// CHECK-LABEL: define hidden void @critical_device(
+// CHECK: %[[MASK:.*]] = call i64 @__kmpc_warp_active_thread_mask()
+// CHECK: %[[TID:.*]] = call i32 @__kmpc_get_hardware_thread_id_in_block()
+// CHECK: %[[NTHREADS:.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
+// CHECK: br label %omp.critical.serial.header
+
+// CHECK: omp.critical.serial.header:
+// CHECK: %[[TURN:.*]] = phi i32 [ 0, %{{.*}} ], [ %[[NEXT:.*]], %omp.critical.serial.sync ]
+// CHECK: %[[GO:.*]] = icmp slt i32 %[[TURN]], %[[NTHREADS]]
+// CHECK: br i1 %[[GO]], label %omp.critical.serial.turn, label %omp.critical.serial.exit
+
+// CHECK: omp.critical.serial.turn:
+// CHECK: %[[MINE:.*]] = icmp eq i32 %[[TID]], %[[TURN]]
+// CHECK: br i1 %[[MINE]], label %omp.critical.serial.region, label %omp.critical.serial.sync
+
+// The lock must be acquired and released inside the region a single thread
+// enters, not around the loop.
+// CHECK: omp.critical.serial.region:
+// CHECK: call void @__kmpc_critical(
+// CHECK: omp.critical.region:
+// CHECK: store i32 %{{.*}}, ptr %{{.*}}
+// CHECK: omp_region.finalize:
+// CHECK: call void @__kmpc_end_critical(
+// CHECK: br label %omp.critical.serial.sync
+
+// CHECK: omp.critical.serial.sync:
+// CHECK: call void @__kmpc_syncwarp(i64 %[[MASK]])
+// CHECK: %[[NEXT]] = add i32 %[[TURN]], 1
+// CHECK: br label %omp.critical.serial.header
+
+// CHECK: omp.critical.serial.exit:
``````````
</details>
https://github.com/llvm/llvm-project/pull/215009
More information about the llvm-commits
mailing list