[llvm] [OpenMPOpt] Look inside the callbacks the loop runtime functions are handed (PR #221449)
Larry Meadows via llvm-commits
llvm-commits at lists.llvm.org
Tue Sep 8 02:41:09 PDT 2026
https://github.com/lfmeadow updated https://github.com/llvm/llvm-project/pull/221449
>From cc836f720adfd03657432ea10e9ca156fb5f55c1 Mon Sep 17 00:00:00 2001
From: Larry Meadows <Lawrence.Meadows at amd.com>
Date: Fri, 21 Aug 2026 22:22:51 -0500
Subject: [PATCH] [OpenMPOpt] Look inside the callbacks the loop runtime
functions are handed
The __kmpc_{distribute_,for_,distribute_for_}static_loop_* functions receive the
loop body as a callback, so a parallel region written inside that body is
reachable from the kernel through the runtime call. AAKernelInfo could not see
that, and recorded the call as reaching an unknown parallel region. A kernel
using these functions therefore always got a worker state machine whose only
option was to indirectly call whatever work function it was handed.
Describe the callback argument of each of these functions in OMPKinds.def and
attach the corresponding !callback metadata in OpenMPOpt, then fold the
callback's AAKernelInfo state into the caller's. The state machine can now
dispatch directly to the regions the loop body actually reaches. Relax the two
"more than one callee means give up" checks for functions carrying !callback,
since the callback edge is a second edge by construction and is analyzable.
The conservative unknown-region record is kept for the case that motivated it, a
callback we only see a declaration of.
These calls still block SPMD-ization. The existing TODO says they need not, once
the contents of the callbacks are considered, which is what this change does --
but acting on it faults, so it is replaced with the obstacles that actually
remain. Two different things are going on, so the cases are split apart.
A plain `distribute` spreads its iterations over the teams rather than the
threads of a team: the runtime runs __kmpc_distribute_static_loop_* with TId 0
and a team size of one and asserts parallel level 0, so exactly one thread per
block may call it. SPMD-izing such a kernel has every thread call it, each
running the whole of its block's share of the body and entering any nested
parallel region from level 0. No amount of analysis of the body changes that, so
this family is unconditionally incompatible.
The thread-indexed families are a different matter. They are meant to be called
by every thread of the block, and a kernel reaching one is not SPMD-incompatible
for that reason alone; what stops us is the transform. Guarding keeps
single-threaded code on one thread with a block-wide barrier, and a barrier
inside a loop body only some threads run is divergent. The values that code
produces are also handed to the other threads through shared memory, written
under a `tid == 0` guard, and which threads the runtime invokes the body on is
its own choice; when thread 0 is not among them the others read a shared slot
nobody wrote. That is a missing capability rather than a property of the
construct, so it is written down as one.
Ported from AMD's ROCm fork, commit 8176194f36b9, which also used this to
SPMD-ize such kernels. That part is left out: it is disabled by default
downstream after it regressed applications.
callback_parallel_regions.ll covers the analysis on all twelve functions.
target-teams-distribute-nested-parallel-do.f90 is the reproducer for the fault:
its loop body is a definition the compiler can see through, so the analysis
knows every region the kernel reaches, and SPMD-izing on the strength of that
writes through a null pointer read out of the unwritten shared slot.
No functional change to any kernel's execution mode: both groups block
SPMD-ization exactly as before. The exec mode of all 766 kernels in the AOMP
smoke and smoke-fort suites is unchanged.
Co-authored-by: Cursor <cursoragent at cursor.com>
---
.../include/llvm/Frontend/OpenMP/OMPKinds.def | 51 ++++
llvm/lib/Transforms/IPO/OpenMPOpt.cpp | 109 +++++++-
.../OpenMP/callback_parallel_regions.ll | 239 ++++++++++++++++++
...et-teams-distribute-nested-parallel-do.f90 | 48 ++++
4 files changed, 437 insertions(+), 10 deletions(-)
create mode 100644 llvm/test/Transforms/OpenMP/callback_parallel_regions.ll
create mode 100644 offload/test/offloading/fortran/target-teams-distribute-nested-parallel-do.f90
diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def b/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
index 8e2a386cf55e9..8d94d952fb39f 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
@@ -1215,6 +1215,57 @@ __OMP_PROC_BIND_KIND(unknown, 7)
///}
+/// Which argument of a runtime function is a callback, and what the callback
+/// receives, in the encoding !callback metadata uses: the argument number the
+/// callee is passed in, then one entry per callback parameter giving the
+/// argument of the runtime call it is forwarded from, with -1 for a parameter
+/// the runtime supplies itself.
+///
+///{
+
+#ifndef OMP_RTL_CB_INFO
+#define OMP_RTL_CB_INFO(Enum, Str, ArgNo, ArgIndices, IsVarArg)
+#endif
+
+#define __OMP_RTL_CB_INFO(Name, ArgNo, ArgIndices, IsVarArg) \
+ OMP_RTL_CB_INFO(OMPRTL_##Name, #Name, ArgNo, ArgIndices, IsVarArg)
+
+/// The loop runtime functions all take (loc, fn, arg, num_iters, ...) and hand
+/// the loop body (fn) an iteration the runtime picks and the arg they were
+/// given, so the body's second parameter is in fact argument 2. It is left
+/// unknown deliberately: the runtime invokes the body once per iteration while
+/// the IR shows a single call, and naming the argument it comes from would let
+/// a consumer reason about what the payload carries from one iteration to the
+/// next as if the body ran once.
+#define __OMP_LOOP_CB_INDICES ArrayRef<int>({-1, -1})
+
+__OMP_RTL_CB_INFO(__kmpc_distribute_static_loop_4, 1, __OMP_LOOP_CB_INDICES,
+ false)
+__OMP_RTL_CB_INFO(__kmpc_distribute_static_loop_4u, 1, __OMP_LOOP_CB_INDICES,
+ false)
+__OMP_RTL_CB_INFO(__kmpc_distribute_static_loop_8, 1, __OMP_LOOP_CB_INDICES,
+ false)
+__OMP_RTL_CB_INFO(__kmpc_distribute_static_loop_8u, 1, __OMP_LOOP_CB_INDICES,
+ false)
+__OMP_RTL_CB_INFO(__kmpc_distribute_for_static_loop_4, 1, __OMP_LOOP_CB_INDICES,
+ false)
+__OMP_RTL_CB_INFO(__kmpc_distribute_for_static_loop_4u, 1,
+ __OMP_LOOP_CB_INDICES, false)
+__OMP_RTL_CB_INFO(__kmpc_distribute_for_static_loop_8, 1, __OMP_LOOP_CB_INDICES,
+ false)
+__OMP_RTL_CB_INFO(__kmpc_distribute_for_static_loop_8u, 1,
+ __OMP_LOOP_CB_INDICES, false)
+__OMP_RTL_CB_INFO(__kmpc_for_static_loop_4, 1, __OMP_LOOP_CB_INDICES, false)
+__OMP_RTL_CB_INFO(__kmpc_for_static_loop_4u, 1, __OMP_LOOP_CB_INDICES, false)
+__OMP_RTL_CB_INFO(__kmpc_for_static_loop_8, 1, __OMP_LOOP_CB_INDICES, false)
+__OMP_RTL_CB_INFO(__kmpc_for_static_loop_8u, 1, __OMP_LOOP_CB_INDICES, false)
+
+#undef __OMP_LOOP_CB_INDICES
+#undef __OMP_RTL_CB_INFO
+#undef OMP_RTL_CB_INFO
+
+///}
+
/// OpenMP context related definitions:
/// - trait set selector
/// - trait selector
diff --git a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
index e1d5b99e867a6..e1dbbea8dbcdc 100644
--- a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
+++ b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
@@ -51,6 +51,7 @@
#include "llvm/IR/IntrinsicsAMDGPU.h"
#include "llvm/IR/IntrinsicsNVPTX.h"
#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/MDBuilder.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
@@ -551,6 +552,53 @@ struct OMPInformationCache : public InformationCache {
collectUses(RFI, /*CollectStats*/ false);
}
+ /// Attach !callback metadata to a runtime function that takes one, so that
+ /// the Attributor sees the edge from the runtime call to the callback and
+ /// AAKernelInfo can look inside it. The runtime declares these functions
+ /// without the metadata, so OpenMPOpt supplies it from the table in
+ /// OMPKinds.def.
+ void setCallbackMetadata(Function *F, unsigned ArgNo, ArrayRef<int> Indices,
+ bool IsVarArg) {
+ if (!F || F->hasMetadata(LLVMContext::MD_callback))
+ return;
+
+ LLVMContext &Ctx = F->getContext();
+ MDBuilder MDB(Ctx);
+ F->addMetadata(LLVMContext::MD_callback,
+ *MDNode::get(Ctx, {MDB.createCallbackEncoding(ArgNo, Indices,
+ IsVarArg)}));
+ }
+
+ /// The callback a runtime function was handed, if it is one we can analyze.
+ /// Returns null when the call takes no callback, or when the callback is not
+ /// a definition this module can see, in which case its contents are unknown
+ /// and callers have to stay conservative.
+ static Function *getAnalyzableCallback(const CallBase &CB) {
+ Function *Callee = CB.getCalledFunction();
+ if (!Callee)
+ return nullptr;
+ MDNode *CallbackMD = Callee->getMetadata(LLVMContext::MD_callback);
+ if (!CallbackMD || CallbackMD->getNumOperands() == 0)
+ return nullptr;
+ // TODO: A runtime function with more than one callback would need each of
+ // them checked; none of the ones in the table have more than one.
+ auto *Encoding = dyn_cast<MDNode>(CallbackMD->getOperand(0));
+ if (!Encoding || Encoding->getNumOperands() == 0)
+ return nullptr;
+ auto *ArgNoMD = dyn_cast<ConstantAsMetadata>(Encoding->getOperand(0));
+ if (!ArgNoMD)
+ return nullptr;
+ uint64_t ArgNo =
+ cast<ConstantInt>(ArgNoMD->getValue())->getLimitedValue(UINT64_MAX);
+ if (ArgNo >= CB.arg_size())
+ return nullptr;
+ auto *Callback =
+ dyn_cast<Function>(CB.getArgOperand(ArgNo)->stripPointerCasts());
+ if (!Callback || Callback->isDeclaration())
+ return nullptr;
+ return Callback;
+ }
+
// Helper function to recollect uses of all runtime functions.
void recollectUses() {
for (int Idx = 0; Idx < RFIs.size(); ++Idx)
@@ -634,6 +682,10 @@ struct OMPInformationCache : public InformationCache {
}); \
} \
}
+
+#define OMP_RTL_CB_INFO(_Enum, _Name, _ArgNo, _ArgIndices, _IsVarArg) \
+ setCallbackMetadata(M.getFunction(_Name), _ArgNo, _ArgIndices, _IsVarArg);
+
#include "llvm/Frontend/OpenMP/OMPKinds.def"
// Remove the `noinline` attribute from `__kmpc`, `ompx::` and `omp_`
@@ -4831,6 +4883,24 @@ struct AAKernelInfoFunction : AAKernelInfo {
bool AllSPMDStatesWereFixed = true;
auto CheckCallInst = [&](Instruction &I) {
auto &CB = cast<CallBase>(I);
+ // A runtime function that takes a callback runs the user's code inside
+ // it, so whatever the callback reaches this kernel reaches too. Fold the
+ // callback's state in; without this the call tells us nothing about the
+ // parallel regions on the other side of it.
+ if (Function *Callback = OMPInformationCache::getAnalyzableCallback(CB)) {
+ LLVM_DEBUG(dbgs() << TAG << "folding in callback "
+ << Callback->getName() << " of " << CB << "\n");
+ if (auto *CallbackAA = A.getAAFor<AAKernelInfo>(
+ *this, IRPosition::function(*Callback), DepClassTy::OPTIONAL)) {
+ getState() ^= CallbackAA->getState();
+ AllSPMDStatesWereFixed &=
+ CallbackAA->SPMDCompatibilityTracker.isAtFixpoint();
+ AllParallelRegionStatesWereFixed &=
+ CallbackAA->ReachedKnownParallelRegions.isAtFixpoint();
+ AllParallelRegionStatesWereFixed &=
+ CallbackAA->ReachedUnknownParallelRegions.isAtFixpoint();
+ }
+ }
auto *CBAA = A.getAAFor<AAKernelInfo>(
*this, IRPosition::callsite_function(CB), DepClassTy::OPTIONAL);
if (!CBAA)
@@ -5006,7 +5076,10 @@ struct AAKernelInfoCallSite : AAKernelInfo {
// state based on the callee state in updateImpl.
return;
}
- if (NumCallees > 1) {
+ // More than one callee normally means an indirect call we cannot resolve.
+ // A runtime function carrying !callback is the exception: the extra edge
+ // is the callback, which we analyze rather than give up on.
+ if (NumCallees > 1 && !Callee->hasMetadata(LLVMContext::MD_callback)) {
indicatePessimisticFixpoint();
return;
}
@@ -5106,6 +5179,18 @@ struct AAKernelInfoCallSite : AAKernelInfo {
case OMPRTL___kmpc_distribute_static_loop_4u:
case OMPRTL___kmpc_distribute_static_loop_8:
case OMPRTL___kmpc_distribute_static_loop_8u:
+ // A plain `distribute` spreads its iterations over the teams, not over
+ // the threads of a team: the runtime runs it with TId 0 and a team size
+ // of one, and asserts the kernel is at parallel level 0. One thread per
+ // block calls it, which is what generic mode gives it. In SPMD mode
+ // every thread would call it, each running the whole of its block's
+ // share of the loop body, so the kernel cannot be SPMD-ized however
+ // analyzable the body is.
+ if (!OMPInformationCache::getAnalyzableCallback(CB))
+ ReachedUnknownParallelRegions.insert(&CB);
+ SPMDCompatibilityTracker.indicatePessimisticFixpoint();
+ SPMDCompatibilityTracker.insert(&CB);
+ break;
case OMPRTL___kmpc_distribute_for_static_loop_4:
case OMPRTL___kmpc_distribute_for_static_loop_4u:
case OMPRTL___kmpc_distribute_for_static_loop_8:
@@ -5114,14 +5199,16 @@ 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.
+ // These index by the thread's own id, so unlike a plain distribute they
+ // are meant to be called by every thread of the block, and a kernel
+ // reaching one is not SPMD-incompatible for that reason alone. What
+ // stops us is the transform rather than the analysis: SPMD-ization
+ // guards whatever has to stay single-threaded with a block-wide
+ // barrier, and a barrier placed inside a loop body only some threads
+ // run is divergent. Until guarding can express "the thread that owns
+ // this iteration", stay conservative here too.
+ if (!OMPInformationCache::getAnalyzableCallback(CB))
+ ReachedUnknownParallelRegions.insert(&CB);
SPMDCompatibilityTracker.indicatePessimisticFixpoint();
SPMDCompatibilityTracker.insert(&CB);
break;
@@ -5176,7 +5263,9 @@ struct AAKernelInfoCallSite : AAKernelInfo {
getState() = FnAA->getState();
return ChangeStatus::CHANGED;
}
- if (NumCallees > 1)
+ // See the matching check in initialize: a !callback runtime function has
+ // a second call edge by construction, and it is one we can analyze.
+ if (NumCallees > 1 && !F->hasMetadata(LLVMContext::MD_callback))
return indicatePessimisticFixpoint();
CallBase &CB = cast<CallBase>(getAssociatedValue());
diff --git a/llvm/test/Transforms/OpenMP/callback_parallel_regions.ll b/llvm/test/Transforms/OpenMP/callback_parallel_regions.ll
new file mode 100644
index 0000000000000..9031c05350b5e
--- /dev/null
+++ b/llvm/test/Transforms/OpenMP/callback_parallel_regions.ll
@@ -0,0 +1,239 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --check-globals all --version 6
+; The state machine rewrite has to be on for this test to say anything; it is
+; the default upstream but not everywhere, so ask for it explicitly.
+; RUN: opt --mtriple=amdgpu-amd-amdhsa --data-layout=A5 -S -passes=openmp-opt \
+; RUN: -openmp-opt-disable-state-machine-rewrite=false < %s | FileCheck %s
+
+; The loop runtime functions take the loop body as a callback, so a parallel
+; region inside the body is reached by the kernel through the runtime call.
+; When the body is a definition we can see, the state machine should dispatch
+; to that region directly; when it is opaque, the kernel has to keep the
+; fallback that runs whatever work function it is handed.
+;
+; Both kernels stay in generic mode. Seeing the body does not make a kernel
+; that distributes over teams SPMD-izable: __kmpc_distribute_static_loop_* is
+; called by one thread per block by construction, so knowing what the body
+; reaches says nothing about running the call itself on every thread.
+
+%struct.ident_t = type { i32, i32, i32, i32, ptr }
+%struct.KernelEnvironmentTy = type { %struct.ConfigurationEnvironmentTy, ptr, ptr }
+%struct.ConfigurationEnvironmentTy = type { i8, i8, i8, i32, i32, i32, i32, i32, i32 }
+
+ at 0 = private unnamed_addr constant [23 x i8] c";unknown;unknown;0;0;;\00", align 1
+ at 1 = private unnamed_addr constant %struct.ident_t { i32 0, i32 2, i32 0, i32 0, ptr @0 }, align 8
+ at visible_body_kernel_environment = local_unnamed_addr constant %struct.KernelEnvironmentTy { %struct.ConfigurationEnvironmentTy { i8 1, i8 0, i8 1, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0 }, ptr @1, ptr null }
+ at opaque_body_kernel_environment = local_unnamed_addr constant %struct.KernelEnvironmentTy { %struct.ConfigurationEnvironmentTy { i8 1, i8 0, i8 1, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0 }, ptr @1, ptr null }
+
+;.
+; CHECK: @[[GLOB0:[0-9]+]] = private unnamed_addr constant [23 x i8] c"
+; CHECK: @[[GLOB1:[0-9]+]] = private unnamed_addr constant %struct.ident_t { i32 0, i32 2, i32 0, i32 0, ptr @[[GLOB0]] }, align 8
+; CHECK: @visible_body_kernel_environment = local_unnamed_addr constant %struct.KernelEnvironmentTy { %struct.ConfigurationEnvironmentTy { i8 0, i8 1, i8 1, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0 }, ptr @[[GLOB1]], ptr null }
+; CHECK: @opaque_body_kernel_environment = local_unnamed_addr constant %struct.KernelEnvironmentTy { %struct.ConfigurationEnvironmentTy { i8 0, i8 0, i8 1, i32 0, i32 0, i32 0, i32 0, i32 0, i32 0 }, ptr @[[GLOB1]], ptr null }
+;.
+define weak amdgpu_kernel void @visible_body(ptr %dyn, ptr %payload) #0 {
+; CHECK-LABEL: define weak amdgpu_kernel void @visible_body(
+; CHECK-SAME: ptr [[DYN:%.*]], ptr [[PAYLOAD:%.*]]) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[WORKER_WORK_FN_ADDR:%.*]] = alloca ptr, align 8, addrspace(5)
+; CHECK-NEXT: [[TMP0:%.*]] = call i32 @__kmpc_target_init(ptr @visible_body_kernel_environment, ptr [[DYN]])
+; CHECK-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
+; CHECK-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
+; CHECK: [[IS_WORKER_CHECK]]:
+; CHECK-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
+; CHECK-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
+; CHECK-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; CHECK-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
+; CHECK-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
+; CHECK: [[WORKER_STATE_MACHINE_BEGIN]]:
+; CHECK-NEXT: call void @__kmpc_barrier_simple_generic(ptr @[[GLOB1]], i32 [[TMP0]])
+; CHECK-NEXT: [[WORKER_WORK_FN_ADDR_GENERIC:%.*]] = addrspacecast ptr addrspace(5) [[WORKER_WORK_FN_ADDR]] to ptr
+; CHECK-NEXT: [[WORKER_IS_ACTIVE:%.*]] = call i1 @__kmpc_kernel_parallel(ptr [[WORKER_WORK_FN_ADDR_GENERIC]])
+; CHECK-NEXT: [[WORKER_WORK_FN:%.*]] = load ptr, ptr [[WORKER_WORK_FN_ADDR_GENERIC]], align 8
+; CHECK-NEXT: [[WORKER_IS_DONE:%.*]] = icmp eq ptr [[WORKER_WORK_FN]], null
+; CHECK-NEXT: br i1 [[WORKER_IS_DONE]], label %[[WORKER_STATE_MACHINE_FINISHED]], label %[[WORKER_STATE_MACHINE_IS_ACTIVE_CHECK:.*]]
+; CHECK: [[WORKER_STATE_MACHINE_FINISHED]]:
+; CHECK-NEXT: ret void
+; CHECK: [[WORKER_STATE_MACHINE_IS_ACTIVE_CHECK]]:
+; CHECK-NEXT: br i1 [[WORKER_IS_ACTIVE]], label %[[WORKER_STATE_MACHINE_PARALLEL_REGION_CHECK:.*]], label %[[WORKER_STATE_MACHINE_DONE_BARRIER:.*]]
+; CHECK: [[WORKER_STATE_MACHINE_PARALLEL_REGION_CHECK]]:
+; CHECK-NEXT: br i1 true, label %[[WORKER_STATE_MACHINE_PARALLEL_REGION_EXECUTE:.*]], label %[[WORKER_STATE_MACHINE_PARALLEL_REGION_CHECK1:.*]]
+; CHECK: [[WORKER_STATE_MACHINE_PARALLEL_REGION_EXECUTE]]:
+; CHECK-NEXT: call void @outlined_wrapper(i16 0, i32 [[TMP0]])
+; CHECK-NEXT: br label %[[WORKER_STATE_MACHINE_PARALLEL_REGION_END:.*]]
+; CHECK: [[WORKER_STATE_MACHINE_PARALLEL_REGION_CHECK1]]:
+; CHECK-NEXT: br label %[[WORKER_STATE_MACHINE_PARALLEL_REGION_END]]
+; CHECK: [[WORKER_STATE_MACHINE_PARALLEL_REGION_END]]:
+; CHECK-NEXT: call void @__kmpc_kernel_end_parallel()
+; CHECK-NEXT: br label %[[WORKER_STATE_MACHINE_DONE_BARRIER]]
+; CHECK: [[WORKER_STATE_MACHINE_DONE_BARRIER]]:
+; CHECK-NEXT: call void @__kmpc_barrier_simple_generic(ptr @[[GLOB1]], i32 [[TMP0]])
+; CHECK-NEXT: br label %[[WORKER_STATE_MACHINE_BEGIN]]
+; CHECK: [[THREAD_USER_CODE_CHECK]]:
+; CHECK-NEXT: [[EXEC_USER_CODE:%.*]] = icmp eq i32 [[TMP0]], -1
+; CHECK-NEXT: br i1 [[EXEC_USER_CODE]], label %[[USER_CODE_ENTRY:.*]], label %[[WORKER_EXIT:.*]]
+; CHECK: [[USER_CODE_ENTRY]]:
+; CHECK-NEXT: call void @__kmpc_distribute_static_loop_4u(ptr @[[GLOB1]], ptr @loop_body, ptr [[PAYLOAD]], i32 100, i32 0, i8 0)
+; CHECK-NEXT: call void @__kmpc_target_deinit()
+; CHECK-NEXT: ret void
+; CHECK: [[WORKER_EXIT]]:
+; CHECK-NEXT: ret void
+;
+entry:
+ %0 = call i32 @__kmpc_target_init(ptr @visible_body_kernel_environment, ptr %dyn)
+ %exec_user_code = icmp eq i32 %0, -1
+ br i1 %exec_user_code, label %user_code.entry, label %worker.exit
+
+user_code.entry:
+ call void @__kmpc_distribute_static_loop_4u(ptr @1, ptr @loop_body, ptr %payload, i32 100, i32 0, i8 0)
+ call void @__kmpc_target_deinit()
+ ret void
+
+worker.exit:
+ ret void
+}
+
+define weak amdgpu_kernel void @opaque_body(ptr %dyn, ptr %payload) #0 {
+; CHECK-LABEL: define weak amdgpu_kernel void @opaque_body(
+; CHECK-SAME: ptr [[DYN:%.*]], ptr [[PAYLOAD:%.*]]) #[[ATTR0]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[WORKER_WORK_FN_ADDR:%.*]] = alloca ptr, align 8, addrspace(5)
+; CHECK-NEXT: [[TMP0:%.*]] = call i32 @__kmpc_target_init(ptr @opaque_body_kernel_environment, ptr [[DYN]])
+; CHECK-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
+; CHECK-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
+; CHECK: [[IS_WORKER_CHECK]]:
+; CHECK-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
+; CHECK-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
+; CHECK-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; CHECK-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
+; CHECK-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
+; CHECK: [[WORKER_STATE_MACHINE_BEGIN]]:
+; CHECK-NEXT: call void @__kmpc_barrier_simple_generic(ptr @[[GLOB1]], i32 [[TMP0]])
+; CHECK-NEXT: [[WORKER_WORK_FN_ADDR_GENERIC:%.*]] = addrspacecast ptr addrspace(5) [[WORKER_WORK_FN_ADDR]] to ptr
+; CHECK-NEXT: [[WORKER_IS_ACTIVE:%.*]] = call i1 @__kmpc_kernel_parallel(ptr [[WORKER_WORK_FN_ADDR_GENERIC]])
+; CHECK-NEXT: [[WORKER_WORK_FN:%.*]] = load ptr, ptr [[WORKER_WORK_FN_ADDR_GENERIC]], align 8
+; CHECK-NEXT: [[WORKER_IS_DONE:%.*]] = icmp eq ptr [[WORKER_WORK_FN]], null
+; CHECK-NEXT: br i1 [[WORKER_IS_DONE]], label %[[WORKER_STATE_MACHINE_FINISHED]], label %[[WORKER_STATE_MACHINE_IS_ACTIVE_CHECK:.*]]
+; CHECK: [[WORKER_STATE_MACHINE_FINISHED]]:
+; CHECK-NEXT: ret void
+; CHECK: [[WORKER_STATE_MACHINE_IS_ACTIVE_CHECK]]:
+; CHECK-NEXT: br i1 [[WORKER_IS_ACTIVE]], label %[[WORKER_STATE_MACHINE_PARALLEL_REGION_FALLBACK_EXECUTE:.*]], label %[[WORKER_STATE_MACHINE_DONE_BARRIER:.*]]
+; CHECK: [[WORKER_STATE_MACHINE_PARALLEL_REGION_FALLBACK_EXECUTE]]:
+; CHECK-NEXT: call void [[WORKER_WORK_FN]](i16 0, i32 [[TMP0]])
+; CHECK-NEXT: br label %[[WORKER_STATE_MACHINE_PARALLEL_REGION_END:.*]]
+; CHECK: [[WORKER_STATE_MACHINE_PARALLEL_REGION_END]]:
+; CHECK-NEXT: call void @__kmpc_kernel_end_parallel()
+; CHECK-NEXT: br label %[[WORKER_STATE_MACHINE_DONE_BARRIER]]
+; CHECK: [[WORKER_STATE_MACHINE_DONE_BARRIER]]:
+; CHECK-NEXT: call void @__kmpc_barrier_simple_generic(ptr @[[GLOB1]], i32 [[TMP0]])
+; CHECK-NEXT: br label %[[WORKER_STATE_MACHINE_BEGIN]]
+; CHECK: [[THREAD_USER_CODE_CHECK]]:
+; CHECK-NEXT: [[EXEC_USER_CODE:%.*]] = icmp eq i32 [[TMP0]], -1
+; CHECK-NEXT: br i1 [[EXEC_USER_CODE]], label %[[USER_CODE_ENTRY:.*]], label %[[WORKER_EXIT:.*]]
+; CHECK: [[USER_CODE_ENTRY]]:
+; CHECK-NEXT: call void @__kmpc_distribute_static_loop_4u(ptr @[[GLOB1]], ptr @opaque_loop_body, ptr [[PAYLOAD]], i32 100, i32 0, i8 0)
+; CHECK-NEXT: call void @__kmpc_target_deinit()
+; CHECK-NEXT: ret void
+; CHECK: [[WORKER_EXIT]]:
+; CHECK-NEXT: ret void
+;
+entry:
+ %0 = call i32 @__kmpc_target_init(ptr @opaque_body_kernel_environment, ptr %dyn)
+ %exec_user_code = icmp eq i32 %0, -1
+ br i1 %exec_user_code, label %user_code.entry, label %worker.exit
+
+user_code.entry:
+ call void @__kmpc_distribute_static_loop_4u(ptr @1, ptr @opaque_loop_body, ptr %payload, i32 100, i32 0, i8 0)
+ call void @__kmpc_target_deinit()
+ ret void
+
+worker.exit:
+ ret void
+}
+
+define internal void @loop_body(i32 %iv, ptr %payload) #1 {
+; CHECK-LABEL: define internal void @loop_body(
+; CHECK-SAME: i32 [[IV:%.*]], ptr [[PAYLOAD:%.*]]) #[[ATTR1:[0-9]+]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TID:%.*]] = call i32 @__kmpc_global_thread_num(ptr @[[GLOB1]]) #[[ATTR1]]
+; CHECK-NEXT: call void @__kmpc_parallel_60(ptr @[[GLOB1]], i32 [[TID]], i32 1, i32 -1, i32 -1, ptr @outlined, ptr @outlined_wrapper, ptr [[PAYLOAD]], i64 1, i32 0)
+; CHECK-NEXT: ret void
+;
+entry:
+ %tid = call i32 @__kmpc_global_thread_num(ptr @1)
+ call void @__kmpc_parallel_60(ptr @1, i32 %tid, i32 1, i32 -1, i32 -1, ptr @outlined, ptr @outlined_wrapper, ptr %payload, i64 1, i32 0)
+ ret void
+}
+
+define internal void @outlined(ptr noalias %tid.addr, ptr noalias %zero.addr) #1 {
+; CHECK-LABEL: define internal void @outlined(
+; CHECK-SAME: ptr noalias [[TID_ADDR:%.*]], ptr noalias [[ZERO_ADDR:%.*]]) #[[ATTR2:[0-9]+]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TID:%.*]] = load i32, ptr [[TID_ADDR]], align 4
+; CHECK-NEXT: [[ZERO:%.*]] = load i32, ptr [[ZERO_ADDR]], align 4
+; CHECK-NEXT: call void @p0(i32 [[TID]], i32 [[ZERO]]) #[[ATTR1]]
+; CHECK-NEXT: ret void
+;
+entry:
+ %tid = load i32, ptr %tid.addr, align 4
+ %zero = load i32, ptr %zero.addr, align 4
+ call void @p0(i32 %tid, i32 %zero)
+ ret void
+}
+
+define internal void @outlined_wrapper(i16 zeroext %0, i32 %1) #1 {
+; CHECK-LABEL: define internal void @outlined_wrapper(
+; CHECK-SAME: i16 zeroext [[TMP0:%.*]], i32 [[TMP1:%.*]]) #[[ATTR2]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TID_ADDR:%.*]] = alloca i32, align 4, addrspace(5)
+; CHECK-NEXT: [[TID_ADDR_CAST:%.*]] = addrspacecast ptr addrspace(5) [[TID_ADDR]] to ptr
+; CHECK-NEXT: [[ZERO_ADDR:%.*]] = alloca i32, align 4, addrspace(5)
+; CHECK-NEXT: [[ZERO_ADDR_CAST:%.*]] = addrspacecast ptr addrspace(5) [[ZERO_ADDR]] to ptr
+; CHECK-NEXT: store i32 [[TMP1]], ptr addrspace(5) [[TID_ADDR]], align 4
+; CHECK-NEXT: store i32 0, ptr addrspace(5) [[ZERO_ADDR]], align 4
+; CHECK-NEXT: call void @outlined(ptr [[TID_ADDR_CAST]], ptr [[ZERO_ADDR_CAST]]) #[[ATTR1]]
+; CHECK-NEXT: ret void
+;
+entry:
+ %tid.addr = alloca i32, align 4, addrspace(5)
+ %tid.addr.cast = addrspacecast ptr addrspace(5) %tid.addr to ptr
+ %zero.addr = alloca i32, align 4, addrspace(5)
+ %zero.addr.cast = addrspacecast ptr addrspace(5) %zero.addr to ptr
+ store i32 %1, ptr %tid.addr.cast, align 4
+ store i32 0, ptr %zero.addr.cast, align 4
+ call void @outlined(ptr %tid.addr.cast, ptr %zero.addr.cast)
+ ret void
+}
+
+define weak i32 @__kmpc_target_init(ptr %0, ptr %1) {
+; CHECK-LABEL: define weak i32 @__kmpc_target_init(
+; CHECK-SAME: ptr [[TMP0:%.*]], ptr [[TMP1:%.*]]) {
+; CHECK-NEXT: ret i32 0
+;
+ ret i32 0
+}
+
+declare void @opaque_loop_body(i32, ptr)
+declare void @p0(i32, i32) #1
+declare void @__kmpc_target_deinit()
+declare i32 @__kmpc_global_thread_num(ptr)
+declare void @__kmpc_parallel_60(ptr, i32, i32, i32, i32, ptr, ptr, ptr, i64, i32)
+declare void @__kmpc_distribute_static_loop_4u(ptr, ptr, ptr, i32, i32, i8)
+
+attributes #0 = { "kernel" }
+attributes #1 = { convergent nounwind }
+
+!llvm.module.flags = !{!0, !1}
+
+!0 = !{i32 7, !"openmp", i32 50}
+!1 = !{i32 7, !"openmp-device", i32 50}
+;.
+; CHECK: attributes #[[ATTR0]] = { "kernel" }
+; CHECK: attributes #[[ATTR1]] = { nounwind }
+; CHECK: attributes #[[ATTR2]] = { convergent nounwind }
+; CHECK: attributes #[[ATTR3:[0-9]+]] = { alwaysinline }
+;.
+; CHECK: [[META0:![0-9]+]] = !{i32 7, !"openmp", i32 50}
+; CHECK: [[META1:![0-9]+]] = !{i32 7, !"openmp-device", i32 50}
+; CHECK: [[META2:![0-9]+]] = !{[[META3:![0-9]+]]}
+; CHECK: [[META3]] = !{i64 1, i64 -1, i64 -1, i1 false}
+;.
diff --git a/offload/test/offloading/fortran/target-teams-distribute-nested-parallel-do.f90 b/offload/test/offloading/fortran/target-teams-distribute-nested-parallel-do.f90
new file mode 100644
index 0000000000000..82b449094b7e3
--- /dev/null
+++ b/offload/test/offloading/fortran/target-teams-distribute-nested-parallel-do.f90
@@ -0,0 +1,48 @@
+! A `distribute` loop whose body contains a parallel region must keep its kernel
+! in generic mode. __kmpc_distribute_static_loop_* spreads its iterations over
+! the teams, not over the threads of a team: the runtime runs it with a team
+! size of one and asserts the kernel is at parallel level 0, so exactly one
+! thread per block may call it. If the kernel were SPMD-ized every thread would
+! call it, each running the whole of its block's share of the body and entering
+! the nested parallel region from level 0.
+!
+! Seeing the loop body does not change that, which is why this is worth a test:
+! the body here is a definition the compiler can look inside, so the analysis
+! knows every parallel region the kernel reaches. Acting on that by SPMD-izing
+! the kernel faults with a memory access at a null address rather than giving a
+! wrong answer quietly.
+!
+! Run at -O2 as well, since the mode is only decided once the optimizer runs.
+
+! REQUIRES: flang, amdgpu
+! RUN: %libomptarget-compile-fortran-generic
+! RUN: %libomptarget-run-generic | %fcheck-generic
+! RUN: %libomptarget-compile-fortran-generic -O2
+! RUN: %libomptarget-run-generic | %fcheck-generic
+
+program main
+ implicit none
+ integer, parameter :: n = 10
+ integer :: array(n, n), i, j, wrong
+
+ array = 0
+
+ !$omp target teams distribute map(tofrom: array)
+ do i = 1, n
+ !$omp parallel do
+ do j = 1, n
+ array(j, i) = i + j
+ end do
+ end do
+
+ wrong = 0
+ do i = 1, n
+ do j = 1, n
+ if (array(j, i) /= i + j) wrong = wrong + 1
+ end do
+ end do
+
+ print *, "wrong elements:", wrong
+end program main
+
+! CHECK: wrong elements: 0
More information about the llvm-commits
mailing list