[llvm] [OpenMP] Analyze the loop-body callback of the static-loop runtime entries (PR #211287)
Spencer Bryngelson via llvm-commits
llvm-commits at lists.llvm.org
Thu Jul 30 15:54:07 PDT 2026
https://github.com/sbryngelson updated https://github.com/llvm/llvm-project/pull/211287
>From d689283d035a5e1c8836697ed0e3eaa519bbd440 Mon Sep 17 00:00:00 2001
From: Spencer Bryngelson <sbryngelson at gmail.com>
Date: Thu, 30 Jul 2026 17:53:24 -0500
Subject: [PATCH] [OpenMP] Analyse the loop-body callback of the static-loop
runtime entries
The __kmpc_*_static_loop_* entries take the loop body as a callback, and
AAKernelInfo treated it as opaque, recording an unknown parallel region.
That makes NestedParallelism true for any kernel whose parallel region
contains a device workshare loop, so MayUseNestedParallelism is written
as 1 where it should be 0.
The callback is a direct function operand at the callsite. Resolve it and
consult its AAKernelInfo, as the __kmpc_parallel_60 handling already does
for its parallel-region operand, and only record an unknown region when it
does not resolve or does reach parallel regions.
Only flang lowers device workshare loops through these entries; clang uses
__kmpc_for_static_init_4 plus an explicit loop, which is already handled.
The SPMD-izability half of the existing TODO is left as is.
---
llvm/lib/Transforms/IPO/OpenMPOpt.cpp | 67 ++++++++++---
.../spmdization_kernel_env_static_loop.ll | 94 +++++++++++++++++++
2 files changed, 150 insertions(+), 11 deletions(-)
create mode 100644 llvm/test/Transforms/OpenMP/spmdization_kernel_env_static_loop.ll
diff --git a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
index e36b05564ff1c..64f6ab690e22e 100644
--- a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
+++ b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
@@ -5040,17 +5040,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.
@@ -5106,6 +5097,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();
@@ -5169,6 +5166,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