[llvm] 4905109 - [OpenMP] Analyze the loop-body callback of the static-loop runtime entries (#211287)
via llvm-commits
llvm-commits at lists.llvm.org
Mon Aug 3 07:22:03 PDT 2026
Author: Spencer Bryngelson
Date: 2026-08-03T10:21:58-04:00
New Revision: 4905109b00e6916a310cf7c521bd8df19c0d4a11
URL: https://github.com/llvm/llvm-project/commit/4905109b00e6916a310cf7c521bd8df19c0d4a11
DIFF: https://github.com/llvm/llvm-project/commit/4905109b00e6916a310cf7c521bd8df19c0d4a11.diff
LOG: [OpenMP] Analyze the loop-body callback of the static-loop runtime entries (#211287)
Fixes #211132. Also removes the trigger for #198621, see below.
`AAKernelInfo` treats the loop body passed to the
`__kmpc_*_static_loop_*` entries as opaque and records an unknown
parallel region for it, per the TODO at the site. Consequently
`NestedParallelism` is true for any kernel whose parallel region
contains a device workshare loop, and `MayUseNestedParallelism` is
written to the kernel environment as 1 where it should be 0.
The callback is a direct function operand at the callsite, so resolve it
and consult its `AAKernelInfo`, exactly as the `__kmpc_parallel_60`
handling already does for its parallel-region operand a few lines away.
Only record an unknown region when it does not resolve, or does reach
parallel regions. The SPMD-izability half of the TODO is left alone.
Only flang lowers device workshare loops through these entries; clang
emits `__kmpc_for_static_init_4` plus an explicit loop, which the
preceding case already handles. So flang kernels get 1 and clang kernels
get 0 on identical source. Controlled pair, flang, gfx90a:
| construct | runtime loop entry | before | after |
|---|---|---|---|
| `!$omp target parallel` | none | 0 | 0 |
| `!$omp target parallel do` | `__kmpc_for_static_loop_4u` | 1 | **0** |
What the 1 costs. It stops `config::mayUseNestedParallelism()` folding,
so the serialized branch in `__kmpc_parallel_60` stays live and carries
its own call to the microtask. After `__kmpc_parallel_60` is inlined
into the kernel there are then two calls to the outlined region rather
than one, `isSoleCallToLocalFunction` is false, and
`LastCallToStaticBonus` never applies. On AMDGPU that is 15000 x 11 =
165000; the analysis starts at -165045 with one callsite and -45 with
two. With two the region stays out of line and the kernel is no longer a
leaf.
`-Rpass-analysis=kernel-resource-usage`, VGPRs / scratch / occupancy:
| | before | after |
|---|---|---|
| reproducer from #211132, gfx90a | 212 / 48 B / 2 | **94 / 0 / 5** |
| WENO5 + HLLC NEQ=8, gfx942 | 196 / 64 B / 2 | **110 / 0 / 4** |
| NEQ=16 | 214 / 328 B / 2 | **138 / 0 / 3** |
| NEQ=24 | 196 / 456 B / 2 | **110 / 392 B / 4** |
End-to-end on gfx942 (MI325X), 1M cells, best of 50, Mcell/s, checksums
bit-identical: 1.35x, 1.47x, 1.28x at NEQ=8/16/24. Baseline run-to-run
spread across jobs is wider than the patched one, so treat the resource
numbers above, which are deterministic, as the primary evidence.
For #198621: step 2 of that root cause identifies the same
`MayUseNestedParallelism=1` as what prevents LTO folding
`omp_get_num_threads()` into a register read, which is what leaves
`DistributeFor` with `NumThreads=1` and skips a suffix of iterations.
Scalar kernels are immune there because they get 0. This refines the
field for the array-expression kernels too, so it addresses that cause
rather than the symptom.
Testing: added `spmdization_kernel_env_static_loop.ll`, which covers
both directions, a callback with no parallel region refining to 0 and a
callback that does contain one staying at 1. It fails without the patch.
`llvm/test/Transforms` (11676) and `llvm/test/CodeGen/AMDGPU` (4920),
16596 tests, 14668 passed with 39 expected failures and no regression.
The one failure, `Transforms/ThinLTOBitcodeWriter/no-type-md.ll`, is
pre-existing and fails identically with the patch reverted.
This affects performance-critical applications on large AMD GPU
supercomputers, including [MFC](https://github.com/MFlowCode/MFC).
All numbers above come from the validated reproducers attached to
#211132 and are independently reproducible; they stand on their own.
This was found and root-caused with the assistance of AI tools.
Added:
llvm/test/Transforms/OpenMP/spmdization_kernel_env_static_loop.ll
Modified:
llvm/lib/Transforms/IPO/OpenMPOpt.cpp
Removed:
################################################################################
diff --git a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
index 22a5efb402504..1165d8628634a 100644
--- a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
+++ b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
@@ -5103,17 +5103,8 @@ struct AAKernelInfoCallSite : AAKernelInfo {
case OMPRTL___kmpc_for_static_loop_4u:
case OMPRTL___kmpc_for_static_loop_8:
case OMPRTL___kmpc_for_static_loop_8u:
- // Parallel regions might be reached by these calls, as they take a
- // callback argument potentially containing arbitrary user-provided
- // code.
- ReachedUnknownParallelRegions.insert(&CB);
- // TODO: The presence of these calls on their own does not prevent a
- // kernel from being SPMD-izable. We mark it as such because we need
- // further changes in order to also consider the contents of the
- // callbacks passed to them.
- SPMDCompatibilityTracker.indicatePessimisticFixpoint();
- SPMDCompatibilityTracker.insert(&CB);
- break;
+ handleStaticLoop(A, CB);
+ return;
default:
// Unknown OpenMP runtime calls cannot be executed in SPMD-mode,
// generally. However, they do not hide parallel regions.
@@ -5169,6 +5160,12 @@ struct AAKernelInfoCallSite : AAKernelInfo {
return indicatePessimisticFixpoint();
CallBase &CB = cast<CallBase>(getAssociatedValue());
+ if (isStaticLoopRTL(It->getSecond())) {
+ handleStaticLoop(A, CB);
+ return StateBefore == getState() ? ChangeStatus::UNCHANGED
+ : ChangeStatus::CHANGED;
+ }
+
if (It->getSecond() == OMPRTL___kmpc_parallel_60) {
if (!handleParallel60(A, CB))
return indicatePessimisticFixpoint();
@@ -5232,6 +5229,54 @@ struct AAKernelInfoCallSite : AAKernelInfo {
/// Deal with a __kmpc_parallel_60 call (\p CB). Returns true if the call was
/// handled, if a problem occurred, false is returned.
+ /// The __kmpc_*_static_loop_* runtime entries take the loop body as a
+ /// callback. Analyse it when it resolves to a function instead of assuming
+ /// it hides arbitrary parallelism.
+ void handleStaticLoop(Attributor &A, CallBase &CB) {
+ const unsigned int LoopBodyArgNo = 1;
+
+ auto *LoopBody = dyn_cast<Function>(
+ CB.getArgOperand(LoopBodyArgNo)->stripPointerCasts());
+ auto *BodyAA =
+ LoopBody
+ ? A.getAAFor<AAKernelInfo>(*this, IRPosition::function(*LoopBody),
+ DepClassTy::OPTIONAL)
+ : nullptr;
+ if (!BodyAA || !BodyAA->getState().isValidState() ||
+ !BodyAA->ReachedKnownParallelRegions.isValidState() ||
+ !BodyAA->ReachedKnownParallelRegions.empty() ||
+ !BodyAA->ReachedUnknownParallelRegions.isValidState() ||
+ !BodyAA->ReachedUnknownParallelRegions.empty())
+ ReachedUnknownParallelRegions.insert(&CB);
+
+ // TODO: The presence of these calls on their own does not prevent a
+ // kernel from being SPMD-izable. We mark it as such because we need
+ // further changes in order to also consider the contents of the
+ // callbacks passed to them.
+ SPMDCompatibilityTracker.indicatePessimisticFixpoint();
+ SPMDCompatibilityTracker.insert(&CB);
+ }
+
+ static bool isStaticLoopRTL(RuntimeFunction RF) {
+ switch (RF) {
+ case OMPRTL___kmpc_distribute_static_loop_4:
+ case OMPRTL___kmpc_distribute_static_loop_4u:
+ case OMPRTL___kmpc_distribute_static_loop_8:
+ case OMPRTL___kmpc_distribute_static_loop_8u:
+ case OMPRTL___kmpc_distribute_for_static_loop_4:
+ case OMPRTL___kmpc_distribute_for_static_loop_4u:
+ case OMPRTL___kmpc_distribute_for_static_loop_8:
+ case OMPRTL___kmpc_distribute_for_static_loop_8u:
+ case OMPRTL___kmpc_for_static_loop_4:
+ case OMPRTL___kmpc_for_static_loop_4u:
+ case OMPRTL___kmpc_for_static_loop_8:
+ case OMPRTL___kmpc_for_static_loop_8u:
+ return true;
+ default:
+ return false;
+ }
+ }
+
bool handleParallel60(Attributor &A, CallBase &CB) {
const unsigned int NonWrapperFunctionArgNo = 5;
const unsigned int WrapperFunctionArgNo = 6;
diff --git a/llvm/test/Transforms/OpenMP/spmdization_kernel_env_static_loop.ll b/llvm/test/Transforms/OpenMP/spmdization_kernel_env_static_loop.ll
new file mode 100644
index 0000000000000..98bbfca760f88
--- /dev/null
+++ b/llvm/test/Transforms/OpenMP/spmdization_kernel_env_static_loop.ll
@@ -0,0 +1,94 @@
+; RUN: opt -S -passes=openmp-opt < %s | FileCheck %s
+
+; The __kmpc_*_static_loop_* entries take the loop body as a callback. Treating
+; it as opaque made every kernel whose parallel region contains a device
+; workshare loop look like it nested parallelism. Resolve the callback: here it
+; is a plain loop body, so MayUseNestedParallelism must be refined to 0.
+; Field order is UseGenericStateMachine, MayUseNestedParallelism, ExecMode.
+
+%struct.KernelEnvironmentTy = type { %struct.ConfigurationEnvironmentTy, ptr, ptr }
+%struct.ConfigurationEnvironmentTy = type { i8, i8, i8, i32, i32, i32, i32 }
+
+ at kernel_environment = local_unnamed_addr constant %struct.KernelEnvironmentTy {
+ %struct.ConfigurationEnvironmentTy { i8 0, i8 1, i8 2, i32 1, i32 256, i32 1, i32 1 },
+ ptr null, ptr null }
+
+; CHECK: @kernel_environment = local_unnamed_addr constant %struct.KernelEnvironmentTy { %struct.ConfigurationEnvironmentTy { i8 0, i8 0, i8 2,
+
+define weak amdgpu_kernel void @kernel(ptr %args) #0 {
+entry:
+ %0 = call i32 @__kmpc_target_init(ptr @kernel_environment, ptr null)
+ call void @__kmpc_parallel_60(ptr null, i32 0, i32 1, i32 -1, i32 -1,
+ ptr @omp_par, ptr null, ptr %args, i64 1, i32 0)
+ call void @__kmpc_target_deinit()
+ ret void
+}
+
+; The parallel region: its whole body is a device workshare loop.
+define internal void @omp_par(ptr %tid, ptr %zero, ptr %args) {
+entry:
+ call void @__kmpc_distribute_for_static_loop_4u(ptr null, ptr @loop_body,
+ ptr %args, i32 64, i32 0,
+ i32 0, i32 0, i8 0)
+ ret void
+}
+
+; The callback: plain work, no parallel region of its own.
+define internal void @loop_body(i32 %iv, ptr %args) {
+entry:
+ %p = load ptr, ptr %args, align 8
+ %idx = zext i32 %iv to i64
+ %gep = getelementptr inbounds i32, ptr %p, i64 %idx
+ store i32 %iv, ptr %gep, align 4
+ ret void
+}
+
+; Negative case: the callback does contain a parallel region, so the refinement
+; must not fire and MayUseNestedParallelism stays 1.
+
+ at kernel_environment_nested = local_unnamed_addr constant %struct.KernelEnvironmentTy {
+ %struct.ConfigurationEnvironmentTy { i8 0, i8 1, i8 2, i32 1, i32 256, i32 1, i32 1 },
+ ptr null, ptr null }
+
+; CHECK: @kernel_environment_nested = local_unnamed_addr constant %struct.KernelEnvironmentTy { %struct.ConfigurationEnvironmentTy { i8 0, i8 1, i8 2,
+
+define weak amdgpu_kernel void @kernel_nested(ptr %args) #0 {
+entry:
+ %0 = call i32 @__kmpc_target_init(ptr @kernel_environment_nested, ptr null)
+ call void @__kmpc_parallel_60(ptr null, i32 0, i32 1, i32 -1, i32 -1,
+ ptr @omp_par_nested, ptr null, ptr %args,
+ i64 1, i32 0)
+ call void @__kmpc_target_deinit()
+ ret void
+}
+
+define internal void @omp_par_nested(ptr %tid, ptr %zero, ptr %args) {
+entry:
+ call void @__kmpc_distribute_for_static_loop_4u(ptr null, ptr @loop_body_nested,
+ ptr %args, i32 64, i32 0,
+ i32 0, i32 0, i8 0)
+ ret void
+}
+
+define internal void @loop_body_nested(i32 %iv, ptr %args) {
+entry:
+ call void @__kmpc_parallel_60(ptr null, i32 0, i32 1, i32 -1, i32 -1,
+ ptr @inner_par, ptr null, ptr %args, i64 1, i32 0)
+ ret void
+}
+
+define internal void @inner_par(ptr %tid, ptr %zero, ptr %args) {
+entry:
+ ret void
+}
+
+declare i32 @__kmpc_target_init(ptr, ptr)
+declare void @__kmpc_target_deinit()
+declare void @__kmpc_parallel_60(ptr, i32, i32, i32, i32, ptr, ptr, ptr, i64, i32)
+declare void @__kmpc_distribute_for_static_loop_4u(ptr, ptr, ptr, i32, i32, i32, i32, i8)
+
+attributes #0 = { "kernel" "omp_target_thread_limit"="256" }
+
+!llvm.module.flags = !{!0, !1}
+!0 = !{i32 7, !"openmp", i32 51}
+!1 = !{i32 7, !"openmp-device", i32 51}
More information about the llvm-commits
mailing list