[Openmp-commits] [llvm] [openmp] [OpenMPOpt] Analyse loop-runtime callbacks, bound specialization, ask the runtime for the worker count (PR #219323)
Larry Meadows via Openmp-commits
openmp-commits at lists.llvm.org
Thu Aug 27 22:16:11 PDT 2026
https://github.com/lfmeadow updated https://github.com/llvm/llvm-project/pull/219323
>From d651ec29d14448b5e1a0811efb87ada7b5c01d29 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 1/5] [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 | 44 ++++
llvm/lib/Transforms/IPO/OpenMPOpt.cpp | 109 ++++++++-
.../OpenMP/callback_parallel_regions.ll | 227 ++++++++++++++++++
...et-teams-distribute-nested-parallel-do.f90 | 48 ++++
4 files changed, 418 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..507f296279a77 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
@@ -1215,6 +1215,50 @@ __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)
+
+#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 22a5efb402504..27726efc1d185 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"
@@ -545,6 +546,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)
@@ -628,6 +676,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_`
@@ -4820,6 +4872,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)
@@ -4995,7 +5065,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;
}
@@ -5095,6 +5168,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:
@@ -5103,14 +5188,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;
@@ -5165,7 +5252,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..e6d57681faee8
--- /dev/null
+++ b/llvm/test/Transforms/OpenMP/callback_parallel_regions.ll
@@ -0,0 +1,227 @@
+; 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: call void @p0() #[[ATTR1]]
+; CHECK-NEXT: ret void
+;
+entry:
+ call void @p0()
+ 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: call void @outlined(ptr undef, ptr undef) #[[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
+ 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() #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
>From 1dea6273298483f6ef725658c220cb29c8c9c1d7 Mon Sep 17 00:00:00 2001
From: Larry Meadows <Lawrence.Meadows at amd.com>
Date: Fri, 21 Aug 2026 22:21:29 -0500
Subject: [PATCH 2/5] [OpenMPOpt] Bound indirect-call specialization instead of
paying for all of it
AAIndirectCallInfo replaces an indirect call with an if-cascade of direct
calls, one arm per assumed callee. Device code routes every parallel region
through a single runtime dispatcher, so one call site there can see every
outlined region in the module, and the cascade it builds is quadratic in a way
nothing later undoes. On a module with 64 target regions the device LTO
pipeline takes 30.6s, of which 24s is this; capping the cascade at three
callees brings it to 6.6s, and the call sites that keep their indirect call
drop from 392 to 203.
The cap is on the number of possible callees, so it is a property of the call
site rather than the module: a dispatcher that could go anywhere keeps its
indirect call and a call site with one or two known callees still gets direct
ones. A single known callee is promoted either way, which is the case that
pays for itself.
ROCm carries this region under `#if 0` for the same reason. A bound keeps the
specialization that is worth having and lets that go.
Co-authored-by: Cursor <cursoragent at cursor.com>
---
llvm/lib/Transforms/IPO/OpenMPOpt.cpp | 16 +++
.../indirect_call_specialization_cap.ll | 123 ++++++++++++++++++
2 files changed, 139 insertions(+)
create mode 100644 llvm/test/Transforms/OpenMP/indirect_call_specialization_cap.ll
diff --git a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
index 27726efc1d185..e12636111f0a0 100644
--- a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
+++ b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
@@ -5919,6 +5919,20 @@ AAFoldRuntimeCall &AAFoldRuntimeCall::createForPosition(const IRPosition &IRP,
return *AA;
}
+/// Bound the if-cascade AAIndirectCallInfo builds for an indirect call. Device
+/// code routes many parallel regions through a single runtime dispatcher, so a
+/// call site there can see every outlined region in the module; specializing
+/// all of them costs more in code size and compile time than the direct calls
+/// are worth.
+static constexpr unsigned MaxIndirectCallSpecializations = 3;
+
+static bool shouldSpecializeIndirectCallee(Attributor &,
+ const AbstractAttribute &,
+ CallBase &, Function &,
+ unsigned NumAssumedCallees) {
+ return NumAssumedCallees <= MaxIndirectCallSpecializations;
+}
+
PreservedAnalyses OpenMPOptPass::run(Module &M, ModuleAnalysisManager &AM) {
if (!containsOpenMP(M))
return PreservedAnalyses::all();
@@ -6005,6 +6019,7 @@ PreservedAnalyses OpenMPOptPass::run(Module &M, ModuleAnalysisManager &AM) {
AC.OREGetter = OREGetter;
AC.PassName = DEBUG_TYPE;
AC.InitializationCallback = OpenMPOpt::registerAAsForFunction;
+ AC.IndirectCalleeSpecializationCallback = shouldSpecializeIndirectCallee;
AC.IPOAmendableCB = [](const Function &F) {
return F.hasFnAttribute("kernel");
};
@@ -6085,6 +6100,7 @@ PreservedAnalyses OpenMPOptCGSCCPass::run(LazyCallGraph::SCC &C,
AC.OREGetter = OREGetter;
AC.PassName = DEBUG_TYPE;
AC.InitializationCallback = OpenMPOpt::registerAAsForFunction;
+ AC.IndirectCalleeSpecializationCallback = shouldSpecializeIndirectCallee;
Attributor A(Functions, InfoCache, AC);
diff --git a/llvm/test/Transforms/OpenMP/indirect_call_specialization_cap.ll b/llvm/test/Transforms/OpenMP/indirect_call_specialization_cap.ll
new file mode 100644
index 0000000000000..862c4a815a57c
--- /dev/null
+++ b/llvm/test/Transforms/OpenMP/indirect_call_specialization_cap.ll
@@ -0,0 +1,123 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt -passes=openmp-opt -S < %s | FileCheck %s
+; The bound is on the number of possible callees, so it is a property of the call
+; site and not of the module: @three gets its if-cascade of direct calls and
+; @four keeps its indirect call.
+
+target triple = "amdgcn-amd-amdhsa"
+
+ at g = external global i32
+
+define internal void @c1() {
+;
+;
+;
+; CHECK-LABEL: define internal void @c1(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT: store volatile i32 1, ptr @g, align 4
+; CHECK-NEXT: ret void
+;
+ store volatile i32 1, ptr @g
+ ret void
+}
+
+define internal void @c2() {
+;
+;
+;
+; CHECK-LABEL: define internal void @c2(
+; CHECK-SAME: ) #[[ATTR0]] {
+; CHECK-NEXT: store volatile i32 2, ptr @g, align 4
+; CHECK-NEXT: ret void
+;
+ store volatile i32 2, ptr @g
+ ret void
+}
+
+define internal void @c3() {
+;
+;
+;
+; CHECK-LABEL: define internal void @c3(
+; CHECK-SAME: ) #[[ATTR0]] {
+; CHECK-NEXT: store volatile i32 3, ptr @g, align 4
+; CHECK-NEXT: ret void
+;
+ store volatile i32 3, ptr @g
+ ret void
+}
+
+define internal void @c4() {
+;
+;
+;
+; CHECK-LABEL: define internal void @c4(
+; CHECK-SAME: ) #[[ATTR0]] {
+; CHECK-NEXT: store volatile i32 4, ptr @g, align 4
+; CHECK-NEXT: ret void
+;
+ store volatile i32 4, ptr @g
+ ret void
+}
+
+define void @three(i1 %a, i1 %b) {
+;
+;
+;
+; CHECK-LABEL: define void @three(
+; CHECK-SAME: i1 [[A:%.*]], i1 [[B:%.*]]) {
+; CHECK-NEXT: [[S0:%.*]] = select i1 [[A]], ptr @c1, ptr @c2
+; CHECK-NEXT: [[S1:%.*]] = select i1 [[B]], ptr [[S0]], ptr @c3
+; CHECK-NEXT: [[TMP1:%.*]] = icmp eq ptr [[S1]], @c3
+; CHECK-NEXT: br i1 [[TMP1]], label %[[BB2:.*]], label %[[BB3:.*]]
+; CHECK: [[BB2]]:
+; CHECK-NEXT: call void @c3()
+; CHECK-NEXT: br label %[[BB9:.*]]
+; CHECK: [[BB3]]:
+; CHECK-NEXT: [[TMP4:%.*]] = icmp eq ptr [[S1]], @c2
+; CHECK-NEXT: br i1 [[TMP4]], label %[[BB5:.*]], label %[[BB6:.*]]
+; CHECK: [[BB5]]:
+; CHECK-NEXT: call void @c2()
+; CHECK-NEXT: br label %[[BB9]]
+; CHECK: [[BB6]]:
+; CHECK-NEXT: br i1 true, label %[[BB7:.*]], label %[[BB8:.*]]
+; CHECK: [[BB7]]:
+; CHECK-NEXT: call void @c1()
+; CHECK-NEXT: br label %[[BB9]]
+; CHECK: [[BB8]]:
+; CHECK-NEXT: unreachable
+; CHECK: [[BB9]]:
+; CHECK-NEXT: ret void
+;
+ %s0 = select i1 %a, ptr @c1, ptr @c2
+ %s1 = select i1 %b, ptr %s0, ptr @c3
+ call void %s1()
+ ret void
+}
+
+define void @four(i1 %a, i1 %b, i1 %c) {
+;
+;
+;
+; CHECK-LABEL: define void @four(
+; CHECK-SAME: i1 [[A:%.*]], i1 [[B:%.*]], i1 [[C:%.*]]) {
+; CHECK-NEXT: [[S0:%.*]] = select i1 [[A]], ptr @c1, ptr @c2
+; CHECK-NEXT: [[S1:%.*]] = select i1 [[B]], ptr [[S0]], ptr @c3
+; CHECK-NEXT: [[S2:%.*]] = select i1 [[C]], ptr [[S1]], ptr @c4
+; CHECK-NEXT: call void [[S2]]() #[[ATTR0]], !callees [[META2:![0-9]+]]
+; CHECK-NEXT: ret void
+;
+ %s0 = select i1 %a, ptr @c1, ptr @c2
+ %s1 = select i1 %b, ptr %s0, ptr @c3
+ %s2 = select i1 %c, ptr %s1, ptr @c4
+ call void %s2()
+ ret void
+}
+
+!llvm.module.flags = !{!0, !1}
+
+!0 = !{i32 7, !"openmp", i32 50}
+!1 = !{i32 7, !"openmp-device", i32 50}
+;.
+; CHECK: [[META2]] = !{ptr @c4, ptr @c3, ptr @c2, ptr @c1}
+;.
>From 4590bc2635dd4b9cc25a07f876ff0ef8d79d1460 Mon Sep 17 00:00:00 2001
From: Larry Meadows <Lawrence.Meadows at amd.com>
Date: Tue, 18 Aug 2026 11:56:16 -0500
Subject: [PATCH 3/5] [OpenMPOpt] Ask the runtime how many of a block's threads
can be workers
The custom state machine gates a thread on InitCB < BlockHwSize - WarpSize,
reconstructing the number of worker threads from the block size on the
assumption that the main thread occupies a whole warp above them. The DeviceRTL
already computes that number, in mapping::getMaxTeamThreads(), and its own
generic state machine gates on it in shouldEnterStateMachine(). Export it as
__kmpc_get_max_team_threads() and call that instead, so the compiler's state
machine and the runtime's agree by construction rather than by arithmetic that
has to be kept in step with the launch geometry.
This is NFC here: getMaxTeamThreads() in generic mode is BlockSize - WarpSize,
the same three instructions folded into one call. It is not NFC for a toolchain
whose launch geometry differs. In ROCm, CGOpenMPRuntimeGPU starts a single extra
thread rather than a warp -- "Only one additional thread is started, not an
entire warp" -- so thread_limit(1024) on a 64-lane target launches 961 threads
and the runtime reports 960 workers, while the state machine's own arithmetic
says 961 - 64 = 897. The threads in between are in neither group: the state
machine returns immediately for them, and the parallel region still hands them
iterations. Asking the runtime removes the assumption rather than adding a
second one.
A device runtime that predates the new entry point is handled by the existing
runtimeFnsAvailable() check, which skips the rewrite rather than emitting a call
to a function that is not there.
The two calls this replaces, __kmpc_get_hardware_num_threads_in_block and
__kmpc_get_warp_size, keep their other uses in the pass, including the constant
folding of the former.
Co-authored-by: Cursor <cursoragent at cursor.com>
---
.../include/llvm/Frontend/OpenMP/OMPKinds.def | 2 +
llvm/lib/Transforms/IPO/OpenMPOpt.cpp | 52 ++++++--------
.../OpenMP/callback_parallel_regions.ll | 8 +--
.../OpenMP/custom_state_machines.ll | 72 +++++++------------
llvm/test/Transforms/OpenMP/spmdization.ll | 64 +++++------------
.../Transforms/OpenMP/spmdization_guarding.ll | 6 +-
.../Transforms/OpenMP/spmdization_indirect.ll | 16 ++---
...zation_no_guarding_two_reaching_kernels.ll | 18 ++---
openmp/device/include/Interface.h | 3 +
openmp/device/src/Mapping.cpp | 4 ++
10 files changed, 85 insertions(+), 160 deletions(-)
diff --git a/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def b/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
index 507f296279a77..cd14a4e7bc676 100644
--- a/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
+++ b/llvm/include/llvm/Frontend/OpenMP/OMPKinds.def
@@ -228,6 +228,7 @@ __OMP_RTL(__kmpc_omp_reg_task_with_affinity, false, Int32, IdentPtr, Int32,
__OMP_RTL(__kmpc_get_hardware_num_blocks, false, Int32, )
__OMP_RTL(__kmpc_get_hardware_num_threads_in_block, false, Int32, )
__OMP_RTL(__kmpc_get_warp_size, false, Int32, )
+__OMP_RTL(__kmpc_get_max_team_threads, false, Int32, )
__OMP_RTL(omp_get_thread_num, false, Int32, )
__OMP_RTL(omp_get_num_threads, false, Int32, )
@@ -721,6 +722,7 @@ __OMP_RTL_ATTRS(__kmpc_omp_reg_task_with_affinity, DefaultAttrs, SExt,
__OMP_RTL_ATTRS(__kmpc_get_hardware_num_blocks, GetterAttrs, ZExt, ParamAttrs())
__OMP_RTL_ATTRS(__kmpc_get_hardware_num_threads_in_block, GetterAttrs, ZExt, ParamAttrs())
__OMP_RTL_ATTRS(__kmpc_get_warp_size, GetterAttrs, ZExt, ParamAttrs())
+__OMP_RTL_ATTRS(__kmpc_get_max_team_threads, GetterAttrs, ZExt, ParamAttrs())
__OMP_RTL_ATTRS(omp_get_thread_num, GetterAttrs, SExt, ParamAttrs())
__OMP_RTL_ATTRS(omp_get_num_threads, GetterAttrs, SExt, ParamAttrs())
diff --git a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
index e12636111f0a0..1865fc0f07c2c 100644
--- a/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
+++ b/llvm/lib/Transforms/IPO/OpenMPOpt.cpp
@@ -3956,8 +3956,7 @@ struct AAKernelInfoFunction : AAKernelInfo {
Attributor::VirtualUseCallbackTy CustomStateMachineUseCB =
[&](Attributor &A, const AbstractAttribute *QueryingAA) {
// Whenever we create a custom state machine we will insert calls to
- // __kmpc_get_hardware_num_threads_in_block,
- // __kmpc_get_warp_size,
+ // __kmpc_get_max_team_threads,
// __kmpc_barrier_simple_generic,
// __kmpc_kernel_parallel, and
// __kmpc_kernel_end_parallel.
@@ -3972,9 +3971,8 @@ struct AAKernelInfoFunction : AAKernelInfo {
// Not needed if we are pre-runtime merge.
if (!KernelInitCB->getCalledFunction()->isDeclaration()) {
- RegisterVirtualUse(OMPRTL___kmpc_get_hardware_num_threads_in_block,
+ RegisterVirtualUse(OMPRTL___kmpc_get_max_team_threads,
CustomStateMachineUseCB);
- RegisterVirtualUse(OMPRTL___kmpc_get_warp_size, CustomStateMachineUseCB);
RegisterVirtualUse(OMPRTL___kmpc_barrier_simple_generic,
CustomStateMachineUseCB);
RegisterVirtualUse(OMPRTL___kmpc_kernel_parallel,
@@ -4433,10 +4431,10 @@ struct AAKernelInfoFunction : AAKernelInfo {
return false;
auto &OMPInfoCache = static_cast<OMPInformationCache &>(A.getInfoCache());
- if (!OMPInfoCache.runtimeFnsAvailable(
- {OMPRTL___kmpc_get_hardware_num_threads_in_block,
- OMPRTL___kmpc_get_warp_size, OMPRTL___kmpc_barrier_simple_generic,
- OMPRTL___kmpc_kernel_parallel, OMPRTL___kmpc_kernel_end_parallel}))
+ if (!OMPInfoCache.runtimeFnsAvailable({OMPRTL___kmpc_get_max_team_threads,
+ OMPRTL___kmpc_barrier_simple_generic,
+ OMPRTL___kmpc_kernel_parallel,
+ OMPRTL___kmpc_kernel_end_parallel}))
return false;
ConstantStruct *ExistingKernelEnvC =
@@ -4515,13 +4513,10 @@ struct AAKernelInfoFunction : AAKernelInfo {
// Create all the blocks:
//
// InitCB = __kmpc_target_init(...)
- // BlockHwSize =
- // __kmpc_get_hardware_num_threads_in_block();
- // WarpSize = __kmpc_get_warp_size();
- // BlockSize = BlockHwSize - WarpSize;
+ // MaxTeamThreads = __kmpc_get_max_team_threads();
// IsWorkerCheckBB: bool IsWorker = InitCB != -1;
// if (IsWorker) {
- // if (InitCB >= BlockSize) return;
+ // if (InitCB >= MaxTeamThreads) return;
// SMBeginBB: __kmpc_barrier_simple_generic(...);
// void *WorkFn;
// bool Active = __kmpc_kernel_parallel(&WorkFn);
@@ -4586,26 +4581,23 @@ struct AAKernelInfoFunction : AAKernelInfo {
IsWorker->setDebugLoc(DLoc);
CondBrInst::Create(IsWorker, IsWorkerCheckBB, UserCodeEntryBB, InitBB);
+ // How many of the block's threads can be worker threads is a property of
+ // the launch geometry, which the runtime knows and the block size alone
+ // does not determine: a target may host the main thread in a whole warp
+ // above the workers, or in a single thread above them. Ask the runtime
+ // rather than subtracting a warp here, or the workers in between are left
+ // in neither group, waiting for no parallel region while it is handed
+ // iterations.
Module &M = *Kernel->getParent();
- FunctionCallee BlockHwSizeFn =
+ FunctionCallee MaxTeamThreadsFn =
OMPInfoCache.OMPBuilder.getOrCreateRuntimeFunction(
- M, OMPRTL___kmpc_get_hardware_num_threads_in_block);
- FunctionCallee WarpSizeFn =
- OMPInfoCache.OMPBuilder.getOrCreateRuntimeFunction(
- M, OMPRTL___kmpc_get_warp_size);
- CallInst *BlockHwSize =
- CallInst::Create(BlockHwSizeFn, "block.hw_size", IsWorkerCheckBB);
- OMPInfoCache.setCallingConvention(BlockHwSizeFn, BlockHwSize);
- BlockHwSize->setDebugLoc(DLoc);
- CallInst *WarpSize =
- CallInst::Create(WarpSizeFn, "warp.size", IsWorkerCheckBB);
- OMPInfoCache.setCallingConvention(WarpSizeFn, WarpSize);
- WarpSize->setDebugLoc(DLoc);
- Instruction *BlockSize = BinaryOperator::CreateSub(
- BlockHwSize, WarpSize, "block.size", IsWorkerCheckBB);
- BlockSize->setDebugLoc(DLoc);
+ M, OMPRTL___kmpc_get_max_team_threads);
+ CallInst *MaxTeamThreads =
+ CallInst::Create(MaxTeamThreadsFn, "max_team_threads", IsWorkerCheckBB);
+ OMPInfoCache.setCallingConvention(MaxTeamThreadsFn, MaxTeamThreads);
+ MaxTeamThreads->setDebugLoc(DLoc);
Instruction *IsMainOrWorker = ICmpInst::Create(
- ICmpInst::ICmp, llvm::CmpInst::ICMP_SLT, KernelInitCB, BlockSize,
+ ICmpInst::ICmp, llvm::CmpInst::ICMP_SLT, KernelInitCB, MaxTeamThreads,
"thread.is_main_or_worker", IsWorkerCheckBB);
IsMainOrWorker->setDebugLoc(DLoc);
CondBrInst::Create(IsMainOrWorker, StateMachineBeginBB,
diff --git a/llvm/test/Transforms/OpenMP/callback_parallel_regions.ll b/llvm/test/Transforms/OpenMP/callback_parallel_regions.ll
index e6d57681faee8..84c6341d49c06 100644
--- a/llvm/test/Transforms/OpenMP/callback_parallel_regions.ll
+++ b/llvm/test/Transforms/OpenMP/callback_parallel_regions.ll
@@ -39,9 +39,7 @@ define weak amdgpu_kernel void @visible_body(ptr %dyn, ptr %payload) #0 {
; 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: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; 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]]:
@@ -101,9 +99,7 @@ define weak amdgpu_kernel void @opaque_body(ptr %dyn, ptr %payload) #0 {
; 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: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; 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]]:
diff --git a/llvm/test/Transforms/OpenMP/custom_state_machines.ll b/llvm/test/Transforms/OpenMP/custom_state_machines.ll
index a03a4b7f3fa67..ac35d7a4a9765 100644
--- a/llvm/test/Transforms/OpenMP/custom_state_machines.ll
+++ b/llvm/test/Transforms/OpenMP/custom_state_machines.ll
@@ -1152,10 +1152,8 @@ attributes #9 = { convergent nounwind willreturn memory(read) }
; AMDGPU-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; AMDGPU-NEXT: br i1 [[THREAD_IS_WORKER]], label [[IS_WORKER_CHECK:%.*]], label [[THREAD_USER_CODE_CHECK:%.*]]
; AMDGPU: is_worker_check:
-; AMDGPU-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; AMDGPU-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; AMDGPU-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
-; AMDGPU-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
+; AMDGPU-NEXT: [[MAX_TEAM_THREADS:%.*]] = call i32 @__kmpc_get_max_team_threads()
+; AMDGPU-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[MAX_TEAM_THREADS]]
; AMDGPU-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label [[WORKER_STATE_MACHINE_BEGIN:%.*]], label [[WORKER_STATE_MACHINE_FINISHED:%.*]]
; AMDGPU: worker_state_machine.begin:
; AMDGPU-NEXT: call void @__kmpc_barrier_simple_generic(ptr @[[GLOB1]], i32 [[TMP0]])
@@ -1273,10 +1271,8 @@ attributes #9 = { convergent nounwind willreturn memory(read) }
; AMDGPU-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; AMDGPU-NEXT: br i1 [[THREAD_IS_WORKER]], label [[IS_WORKER_CHECK:%.*]], label [[THREAD_USER_CODE_CHECK:%.*]]
; AMDGPU: is_worker_check:
-; AMDGPU-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; AMDGPU-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; AMDGPU-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
-; AMDGPU-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
+; AMDGPU-NEXT: [[MAX_TEAM_THREADS:%.*]] = call i32 @__kmpc_get_max_team_threads()
+; AMDGPU-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[MAX_TEAM_THREADS]]
; AMDGPU-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label [[WORKER_STATE_MACHINE_BEGIN:%.*]], label [[WORKER_STATE_MACHINE_FINISHED:%.*]]
; AMDGPU: worker_state_machine.begin:
; AMDGPU-NEXT: call void @__kmpc_barrier_simple_generic(ptr @[[GLOB1]], i32 [[TMP0]])
@@ -1420,10 +1416,8 @@ attributes #9 = { convergent nounwind willreturn memory(read) }
; AMDGPU-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; AMDGPU-NEXT: br i1 [[THREAD_IS_WORKER]], label [[IS_WORKER_CHECK:%.*]], label [[THREAD_USER_CODE_CHECK:%.*]]
; AMDGPU: is_worker_check:
-; AMDGPU-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; AMDGPU-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; AMDGPU-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
-; AMDGPU-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
+; AMDGPU-NEXT: [[MAX_TEAM_THREADS:%.*]] = call i32 @__kmpc_get_max_team_threads()
+; AMDGPU-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[MAX_TEAM_THREADS]]
; AMDGPU-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label [[WORKER_STATE_MACHINE_BEGIN:%.*]], label [[WORKER_STATE_MACHINE_FINISHED:%.*]]
; AMDGPU: worker_state_machine.begin:
; AMDGPU-NEXT: call void @__kmpc_barrier_simple_generic(ptr @[[GLOB1]], i32 [[TMP0]])
@@ -1542,10 +1536,8 @@ attributes #9 = { convergent nounwind willreturn memory(read) }
; AMDGPU-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; AMDGPU-NEXT: br i1 [[THREAD_IS_WORKER]], label [[IS_WORKER_CHECK:%.*]], label [[THREAD_USER_CODE_CHECK:%.*]]
; AMDGPU: is_worker_check:
-; AMDGPU-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; AMDGPU-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; AMDGPU-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
-; AMDGPU-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
+; AMDGPU-NEXT: [[MAX_TEAM_THREADS:%.*]] = call i32 @__kmpc_get_max_team_threads()
+; AMDGPU-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[MAX_TEAM_THREADS]]
; AMDGPU-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label [[WORKER_STATE_MACHINE_BEGIN:%.*]], label [[WORKER_STATE_MACHINE_FINISHED:%.*]]
; AMDGPU: worker_state_machine.begin:
; AMDGPU-NEXT: call void @__kmpc_barrier_simple_generic(ptr @[[GLOB1]], i32 [[TMP0]])
@@ -1662,10 +1654,8 @@ attributes #9 = { convergent nounwind willreturn memory(read) }
; AMDGPU-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; AMDGPU-NEXT: br i1 [[THREAD_IS_WORKER]], label [[IS_WORKER_CHECK:%.*]], label [[THREAD_USER_CODE_CHECK:%.*]]
; AMDGPU: is_worker_check:
-; AMDGPU-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; AMDGPU-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; AMDGPU-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
-; AMDGPU-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
+; AMDGPU-NEXT: [[MAX_TEAM_THREADS:%.*]] = call i32 @__kmpc_get_max_team_threads()
+; AMDGPU-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[MAX_TEAM_THREADS]]
; AMDGPU-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label [[WORKER_STATE_MACHINE_BEGIN:%.*]], label [[WORKER_STATE_MACHINE_FINISHED:%.*]]
; AMDGPU: worker_state_machine.begin:
; AMDGPU-NEXT: call void @__kmpc_barrier_simple_generic(ptr @[[GLOB1]], i32 [[TMP0]])
@@ -1854,10 +1844,8 @@ attributes #9 = { convergent nounwind willreturn memory(read) }
; AMDGPU-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; AMDGPU-NEXT: br i1 [[THREAD_IS_WORKER]], label [[IS_WORKER_CHECK:%.*]], label [[THREAD_USER_CODE_CHECK:%.*]]
; AMDGPU: is_worker_check:
-; AMDGPU-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; AMDGPU-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; AMDGPU-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
-; AMDGPU-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
+; AMDGPU-NEXT: [[MAX_TEAM_THREADS:%.*]] = call i32 @__kmpc_get_max_team_threads()
+; AMDGPU-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[MAX_TEAM_THREADS]]
; AMDGPU-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label [[WORKER_STATE_MACHINE_BEGIN:%.*]], label [[WORKER_STATE_MACHINE_FINISHED:%.*]]
; AMDGPU: worker_state_machine.begin:
; AMDGPU-NEXT: call void @__kmpc_barrier_simple_generic(ptr @[[GLOB1]], i32 [[TMP0]])
@@ -2078,10 +2066,8 @@ attributes #9 = { convergent nounwind willreturn memory(read) }
; NVPTX-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; NVPTX-NEXT: br i1 [[THREAD_IS_WORKER]], label [[IS_WORKER_CHECK:%.*]], label [[THREAD_USER_CODE_CHECK:%.*]]
; NVPTX: is_worker_check:
-; NVPTX-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; NVPTX-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; NVPTX-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
-; NVPTX-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
+; NVPTX-NEXT: [[MAX_TEAM_THREADS:%.*]] = call i32 @__kmpc_get_max_team_threads()
+; NVPTX-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[MAX_TEAM_THREADS]]
; NVPTX-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label [[WORKER_STATE_MACHINE_BEGIN:%.*]], label [[WORKER_STATE_MACHINE_FINISHED:%.*]]
; NVPTX: worker_state_machine.begin:
; NVPTX-NEXT: call void @__kmpc_barrier_simple_generic(ptr @[[GLOB1]], i32 [[TMP0]])
@@ -2198,10 +2184,8 @@ attributes #9 = { convergent nounwind willreturn memory(read) }
; NVPTX-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; NVPTX-NEXT: br i1 [[THREAD_IS_WORKER]], label [[IS_WORKER_CHECK:%.*]], label [[THREAD_USER_CODE_CHECK:%.*]]
; NVPTX: is_worker_check:
-; NVPTX-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; NVPTX-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; NVPTX-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
-; NVPTX-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
+; NVPTX-NEXT: [[MAX_TEAM_THREADS:%.*]] = call i32 @__kmpc_get_max_team_threads()
+; NVPTX-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[MAX_TEAM_THREADS]]
; NVPTX-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label [[WORKER_STATE_MACHINE_BEGIN:%.*]], label [[WORKER_STATE_MACHINE_FINISHED:%.*]]
; NVPTX: worker_state_machine.begin:
; NVPTX-NEXT: call void @__kmpc_barrier_simple_generic(ptr @[[GLOB1]], i32 [[TMP0]])
@@ -2344,10 +2328,8 @@ attributes #9 = { convergent nounwind willreturn memory(read) }
; NVPTX-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; NVPTX-NEXT: br i1 [[THREAD_IS_WORKER]], label [[IS_WORKER_CHECK:%.*]], label [[THREAD_USER_CODE_CHECK:%.*]]
; NVPTX: is_worker_check:
-; NVPTX-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; NVPTX-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; NVPTX-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
-; NVPTX-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
+; NVPTX-NEXT: [[MAX_TEAM_THREADS:%.*]] = call i32 @__kmpc_get_max_team_threads()
+; NVPTX-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[MAX_TEAM_THREADS]]
; NVPTX-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label [[WORKER_STATE_MACHINE_BEGIN:%.*]], label [[WORKER_STATE_MACHINE_FINISHED:%.*]]
; NVPTX: worker_state_machine.begin:
; NVPTX-NEXT: call void @__kmpc_barrier_simple_generic(ptr @[[GLOB1]], i32 [[TMP0]])
@@ -2465,10 +2447,8 @@ attributes #9 = { convergent nounwind willreturn memory(read) }
; NVPTX-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; NVPTX-NEXT: br i1 [[THREAD_IS_WORKER]], label [[IS_WORKER_CHECK:%.*]], label [[THREAD_USER_CODE_CHECK:%.*]]
; NVPTX: is_worker_check:
-; NVPTX-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; NVPTX-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; NVPTX-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
-; NVPTX-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
+; NVPTX-NEXT: [[MAX_TEAM_THREADS:%.*]] = call i32 @__kmpc_get_max_team_threads()
+; NVPTX-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[MAX_TEAM_THREADS]]
; NVPTX-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label [[WORKER_STATE_MACHINE_BEGIN:%.*]], label [[WORKER_STATE_MACHINE_FINISHED:%.*]]
; NVPTX: worker_state_machine.begin:
; NVPTX-NEXT: call void @__kmpc_barrier_simple_generic(ptr @[[GLOB1]], i32 [[TMP0]])
@@ -2584,10 +2564,8 @@ attributes #9 = { convergent nounwind willreturn memory(read) }
; NVPTX-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; NVPTX-NEXT: br i1 [[THREAD_IS_WORKER]], label [[IS_WORKER_CHECK:%.*]], label [[THREAD_USER_CODE_CHECK:%.*]]
; NVPTX: is_worker_check:
-; NVPTX-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; NVPTX-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; NVPTX-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
-; NVPTX-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
+; NVPTX-NEXT: [[MAX_TEAM_THREADS:%.*]] = call i32 @__kmpc_get_max_team_threads()
+; NVPTX-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[MAX_TEAM_THREADS]]
; NVPTX-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label [[WORKER_STATE_MACHINE_BEGIN:%.*]], label [[WORKER_STATE_MACHINE_FINISHED:%.*]]
; NVPTX: worker_state_machine.begin:
; NVPTX-NEXT: call void @__kmpc_barrier_simple_generic(ptr @[[GLOB1]], i32 [[TMP0]])
@@ -2775,10 +2753,8 @@ attributes #9 = { convergent nounwind willreturn memory(read) }
; NVPTX-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; NVPTX-NEXT: br i1 [[THREAD_IS_WORKER]], label [[IS_WORKER_CHECK:%.*]], label [[THREAD_USER_CODE_CHECK:%.*]]
; NVPTX: is_worker_check:
-; NVPTX-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; NVPTX-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; NVPTX-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
-; NVPTX-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
+; NVPTX-NEXT: [[MAX_TEAM_THREADS:%.*]] = call i32 @__kmpc_get_max_team_threads()
+; NVPTX-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[MAX_TEAM_THREADS]]
; NVPTX-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label [[WORKER_STATE_MACHINE_BEGIN:%.*]], label [[WORKER_STATE_MACHINE_FINISHED:%.*]]
; NVPTX: worker_state_machine.begin:
; NVPTX-NEXT: call void @__kmpc_barrier_simple_generic(ptr @[[GLOB1]], i32 [[TMP0]])
diff --git a/llvm/test/Transforms/OpenMP/spmdization.ll b/llvm/test/Transforms/OpenMP/spmdization.ll
index 7d53a808ef4d3..88a4d53ae9c82 100644
--- a/llvm/test/Transforms/OpenMP/spmdization.ll
+++ b/llvm/test/Transforms/OpenMP/spmdization.ll
@@ -271,9 +271,7 @@ define internal void @__omp_offloading_fd02_2044372e_sequential_loop_l5__debug()
; AMDGPU-DISABLED1-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; AMDGPU-DISABLED1-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
; AMDGPU-DISABLED1: [[IS_WORKER_CHECK]]:
-; AMDGPU-DISABLED1-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; AMDGPU-DISABLED1-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; AMDGPU-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; AMDGPU-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; AMDGPU-DISABLED1-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
; AMDGPU-DISABLED1-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
; AMDGPU-DISABLED1: [[WORKER_STATE_MACHINE_BEGIN]]:
@@ -343,9 +341,7 @@ define internal void @__omp_offloading_fd02_2044372e_sequential_loop_l5__debug()
; NVPTX-DISABLED1-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; NVPTX-DISABLED1-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
; NVPTX-DISABLED1: [[IS_WORKER_CHECK]]:
-; NVPTX-DISABLED1-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; NVPTX-DISABLED1-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; NVPTX-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; NVPTX-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; NVPTX-DISABLED1-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
; NVPTX-DISABLED1-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
; NVPTX-DISABLED1: [[WORKER_STATE_MACHINE_BEGIN]]:
@@ -753,9 +749,7 @@ define weak ptx_kernel void @__omp_offloading_fd02_2044372e_sequential_loop_to_s
; AMDGPU-DISABLED1-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; AMDGPU-DISABLED1-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
; AMDGPU-DISABLED1: [[IS_WORKER_CHECK]]:
-; AMDGPU-DISABLED1-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; AMDGPU-DISABLED1-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; AMDGPU-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; AMDGPU-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; AMDGPU-DISABLED1-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
; AMDGPU-DISABLED1-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
; AMDGPU-DISABLED1: [[WORKER_STATE_MACHINE_BEGIN]]:
@@ -825,9 +819,7 @@ define weak ptx_kernel void @__omp_offloading_fd02_2044372e_sequential_loop_to_s
; NVPTX-DISABLED1-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; NVPTX-DISABLED1-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
; NVPTX-DISABLED1: [[IS_WORKER_CHECK]]:
-; NVPTX-DISABLED1-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; NVPTX-DISABLED1-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; NVPTX-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; NVPTX-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; NVPTX-DISABLED1-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
; NVPTX-DISABLED1-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
; NVPTX-DISABLED1: [[WORKER_STATE_MACHINE_BEGIN]]:
@@ -1253,9 +1245,7 @@ define weak ptx_kernel void @__omp_offloading_fd02_2044372e_sequential_loop_to_s
; AMDGPU-DISABLED1-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; AMDGPU-DISABLED1-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
; AMDGPU-DISABLED1: [[IS_WORKER_CHECK]]:
-; AMDGPU-DISABLED1-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; AMDGPU-DISABLED1-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; AMDGPU-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; AMDGPU-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; AMDGPU-DISABLED1-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
; AMDGPU-DISABLED1-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
; AMDGPU-DISABLED1: [[WORKER_STATE_MACHINE_BEGIN]]:
@@ -1325,9 +1315,7 @@ define weak ptx_kernel void @__omp_offloading_fd02_2044372e_sequential_loop_to_s
; NVPTX-DISABLED1-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; NVPTX-DISABLED1-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
; NVPTX-DISABLED1: [[IS_WORKER_CHECK]]:
-; NVPTX-DISABLED1-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; NVPTX-DISABLED1-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; NVPTX-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; NVPTX-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; NVPTX-DISABLED1-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
; NVPTX-DISABLED1-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
; NVPTX-DISABLED1: [[WORKER_STATE_MACHINE_BEGIN]]:
@@ -1779,9 +1767,7 @@ define weak ptx_kernel void @__omp_offloading_fd02_2044372e_sequential_loop_to_s
; AMDGPU-DISABLED1-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; AMDGPU-DISABLED1-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
; AMDGPU-DISABLED1: [[IS_WORKER_CHECK]]:
-; AMDGPU-DISABLED1-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; AMDGPU-DISABLED1-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; AMDGPU-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; AMDGPU-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; AMDGPU-DISABLED1-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
; AMDGPU-DISABLED1-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
; AMDGPU-DISABLED1: [[WORKER_STATE_MACHINE_BEGIN]]:
@@ -1851,9 +1837,7 @@ define weak ptx_kernel void @__omp_offloading_fd02_2044372e_sequential_loop_to_s
; NVPTX-DISABLED1-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; NVPTX-DISABLED1-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
; NVPTX-DISABLED1: [[IS_WORKER_CHECK]]:
-; NVPTX-DISABLED1-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; NVPTX-DISABLED1-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; NVPTX-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; NVPTX-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; NVPTX-DISABLED1-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
; NVPTX-DISABLED1-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
; NVPTX-DISABLED1: [[WORKER_STATE_MACHINE_BEGIN]]:
@@ -2300,9 +2284,7 @@ define weak ptx_kernel void @__omp_offloading_fd02_2044372e_do_not_spmdize_targe
; AMDGPU-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; AMDGPU-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
; AMDGPU: [[IS_WORKER_CHECK]]:
-; AMDGPU-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; AMDGPU-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; AMDGPU-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; AMDGPU-NEXT: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; AMDGPU-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
; AMDGPU-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
; AMDGPU: [[WORKER_STATE_MACHINE_BEGIN]]:
@@ -2348,9 +2330,7 @@ define weak ptx_kernel void @__omp_offloading_fd02_2044372e_do_not_spmdize_targe
; NVPTX-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; NVPTX-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
; NVPTX: [[IS_WORKER_CHECK]]:
-; NVPTX-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; NVPTX-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; NVPTX-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; NVPTX-NEXT: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; NVPTX-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
; NVPTX-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
; NVPTX: [[WORKER_STATE_MACHINE_BEGIN]]:
@@ -2395,9 +2375,7 @@ define weak ptx_kernel void @__omp_offloading_fd02_2044372e_do_not_spmdize_targe
; AMDGPU-DISABLED1-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; AMDGPU-DISABLED1-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
; AMDGPU-DISABLED1: [[IS_WORKER_CHECK]]:
-; AMDGPU-DISABLED1-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; AMDGPU-DISABLED1-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; AMDGPU-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; AMDGPU-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; AMDGPU-DISABLED1-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
; AMDGPU-DISABLED1-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
; AMDGPU-DISABLED1: [[WORKER_STATE_MACHINE_BEGIN]]:
@@ -2461,9 +2439,7 @@ define weak ptx_kernel void @__omp_offloading_fd02_2044372e_do_not_spmdize_targe
; NVPTX-DISABLED1-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; NVPTX-DISABLED1-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
; NVPTX-DISABLED1: [[IS_WORKER_CHECK]]:
-; NVPTX-DISABLED1-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; NVPTX-DISABLED1-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; NVPTX-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; NVPTX-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; NVPTX-DISABLED1-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
; NVPTX-DISABLED1-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
; NVPTX-DISABLED1: [[WORKER_STATE_MACHINE_BEGIN]]:
@@ -2589,9 +2565,7 @@ define weak ptx_kernel void @__omp_offloading_fd02_2044372e_do_not_spmdize_task_
; AMDGPU-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; AMDGPU-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
; AMDGPU: [[IS_WORKER_CHECK]]:
-; AMDGPU-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; AMDGPU-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; AMDGPU-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; AMDGPU-NEXT: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; AMDGPU-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
; AMDGPU-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
; AMDGPU: [[WORKER_STATE_MACHINE_BEGIN]]:
@@ -2643,9 +2617,7 @@ define weak ptx_kernel void @__omp_offloading_fd02_2044372e_do_not_spmdize_task_
; NVPTX-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; NVPTX-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
; NVPTX: [[IS_WORKER_CHECK]]:
-; NVPTX-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; NVPTX-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; NVPTX-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; NVPTX-NEXT: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; NVPTX-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
; NVPTX-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
; NVPTX: [[WORKER_STATE_MACHINE_BEGIN]]:
@@ -2696,9 +2668,7 @@ define weak ptx_kernel void @__omp_offloading_fd02_2044372e_do_not_spmdize_task_
; AMDGPU-DISABLED1-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; AMDGPU-DISABLED1-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
; AMDGPU-DISABLED1: [[IS_WORKER_CHECK]]:
-; AMDGPU-DISABLED1-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; AMDGPU-DISABLED1-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; AMDGPU-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; AMDGPU-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; AMDGPU-DISABLED1-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
; AMDGPU-DISABLED1-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
; AMDGPU-DISABLED1: [[WORKER_STATE_MACHINE_BEGIN]]:
@@ -2768,9 +2738,7 @@ define weak ptx_kernel void @__omp_offloading_fd02_2044372e_do_not_spmdize_task_
; NVPTX-DISABLED1-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; NVPTX-DISABLED1-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
; NVPTX-DISABLED1: [[IS_WORKER_CHECK]]:
-; NVPTX-DISABLED1-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; NVPTX-DISABLED1-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; NVPTX-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; NVPTX-DISABLED1-NEXT: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; NVPTX-DISABLED1-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
; NVPTX-DISABLED1-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
; NVPTX-DISABLED1: [[WORKER_STATE_MACHINE_BEGIN]]:
diff --git a/llvm/test/Transforms/OpenMP/spmdization_guarding.ll b/llvm/test/Transforms/OpenMP/spmdization_guarding.ll
index 2fb1f878cce7b..4dae6889beb9d 100644
--- a/llvm/test/Transforms/OpenMP/spmdization_guarding.ll
+++ b/llvm/test/Transforms/OpenMP/spmdization_guarding.ll
@@ -201,10 +201,8 @@ define weak ptx_kernel void @__omp_offloading_2a_fbfa7a_sequential_loop_l6(ptr %
; CHECK-DISABLED-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; CHECK-DISABLED-NEXT: br i1 [[THREAD_IS_WORKER]], label [[IS_WORKER_CHECK:%.*]], label [[THREAD_USER_CODE_CHECK:%.*]]
; CHECK-DISABLED: is_worker_check:
-; CHECK-DISABLED-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; CHECK-DISABLED-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; CHECK-DISABLED-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
-; CHECK-DISABLED-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
+; CHECK-DISABLED-NEXT: [[MAX_TEAM_THREADS:%.*]] = call i32 @__kmpc_get_max_team_threads()
+; CHECK-DISABLED-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[MAX_TEAM_THREADS]]
; CHECK-DISABLED-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label [[WORKER_STATE_MACHINE_BEGIN:%.*]], label [[WORKER_STATE_MACHINE_FINISHED:%.*]]
; CHECK-DISABLED: worker_state_machine.begin:
; CHECK-DISABLED-NEXT: call void @__kmpc_barrier_simple_generic(ptr @[[GLOB1]], i32 [[TMP0]])
diff --git a/llvm/test/Transforms/OpenMP/spmdization_indirect.ll b/llvm/test/Transforms/OpenMP/spmdization_indirect.ll
index 1de719e1b85a8..71847b72ae24f 100644
--- a/llvm/test/Transforms/OpenMP/spmdization_indirect.ll
+++ b/llvm/test/Transforms/OpenMP/spmdization_indirect.ll
@@ -395,9 +395,7 @@ define weak ptx_kernel void @spmd_and_non_spmd_callee(i1 %c) #0 {
; AMDGPU-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; AMDGPU-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
; AMDGPU: [[IS_WORKER_CHECK]]:
-; AMDGPU-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; AMDGPU-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; AMDGPU-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; AMDGPU-NEXT: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; AMDGPU-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
; AMDGPU-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
; AMDGPU: [[WORKER_STATE_MACHINE_BEGIN]]:
@@ -458,9 +456,7 @@ define weak ptx_kernel void @spmd_and_non_spmd_callee(i1 %c) #0 {
; NVPTX-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; NVPTX-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
; NVPTX: [[IS_WORKER_CHECK]]:
-; NVPTX-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; NVPTX-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; NVPTX-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; NVPTX-NEXT: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; NVPTX-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
; NVPTX-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
; NVPTX: [[WORKER_STATE_MACHINE_BEGIN]]:
@@ -751,9 +747,7 @@ define weak ptx_kernel void @spmd_and_non_spmd_callees_metadata(ptr %fp) #0 {
; AMDGPU-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; AMDGPU-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
; AMDGPU: [[IS_WORKER_CHECK]]:
-; AMDGPU-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; AMDGPU-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; AMDGPU-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; AMDGPU-NEXT: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; AMDGPU-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
; AMDGPU-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
; AMDGPU: [[WORKER_STATE_MACHINE_BEGIN]]:
@@ -813,9 +807,7 @@ define weak ptx_kernel void @spmd_and_non_spmd_callees_metadata(ptr %fp) #0 {
; NVPTX-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; NVPTX-NEXT: br i1 [[THREAD_IS_WORKER]], label %[[IS_WORKER_CHECK:.*]], label %[[THREAD_USER_CODE_CHECK:.*]]
; NVPTX: [[IS_WORKER_CHECK]]:
-; NVPTX-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; NVPTX-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; NVPTX-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
+; NVPTX-NEXT: [[BLOCK_SIZE:%.*]] = call i32 @__kmpc_get_max_team_threads()
; NVPTX-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
; NVPTX-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label %[[WORKER_STATE_MACHINE_BEGIN:.*]], label %[[WORKER_STATE_MACHINE_FINISHED:.*]]
; NVPTX: [[WORKER_STATE_MACHINE_BEGIN]]:
diff --git a/llvm/test/Transforms/OpenMP/spmdization_no_guarding_two_reaching_kernels.ll b/llvm/test/Transforms/OpenMP/spmdization_no_guarding_two_reaching_kernels.ll
index 4f64839deaf53..dc08df96dfd64 100644
--- a/llvm/test/Transforms/OpenMP/spmdization_no_guarding_two_reaching_kernels.ll
+++ b/llvm/test/Transforms/OpenMP/spmdization_no_guarding_two_reaching_kernels.ll
@@ -78,10 +78,8 @@ define weak ptx_kernel void @__omp_offloading_2b_10393b5_spmd_l12(ptr %dyn) #0 {
; CHECK-DISABLE-SPMDIZATION-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; CHECK-DISABLE-SPMDIZATION-NEXT: br i1 [[THREAD_IS_WORKER]], label [[IS_WORKER_CHECK:%.*]], label [[THREAD_USER_CODE_CHECK:%.*]]
; CHECK-DISABLE-SPMDIZATION: is_worker_check:
-; CHECK-DISABLE-SPMDIZATION-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; CHECK-DISABLE-SPMDIZATION-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; CHECK-DISABLE-SPMDIZATION-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
-; CHECK-DISABLE-SPMDIZATION-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
+; CHECK-DISABLE-SPMDIZATION-NEXT: [[MAX_TEAM_THREADS:%.*]] = call i32 @__kmpc_get_max_team_threads()
+; CHECK-DISABLE-SPMDIZATION-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[MAX_TEAM_THREADS]]
; CHECK-DISABLE-SPMDIZATION-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label [[WORKER_STATE_MACHINE_BEGIN:%.*]], label [[WORKER_STATE_MACHINE_FINISHED:%.*]]
; CHECK-DISABLE-SPMDIZATION: worker_state_machine.begin:
; CHECK-DISABLE-SPMDIZATION-NEXT: call void @__kmpc_barrier_simple_generic(ptr @[[GLOB1]], i32 [[TMP0]])
@@ -155,10 +153,8 @@ define weak ptx_kernel void @__omp_offloading_2b_10393b5_generic_l20(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: [[MAX_TEAM_THREADS:%.*]] = call i32 @__kmpc_get_max_team_threads()
+; CHECK-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[MAX_TEAM_THREADS]]
; 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]])
@@ -197,10 +193,8 @@ define weak ptx_kernel void @__omp_offloading_2b_10393b5_generic_l20(ptr %dyn) #
; CHECK-DISABLE-SPMDIZATION-NEXT: [[THREAD_IS_WORKER:%.*]] = icmp ne i32 [[TMP0]], -1
; CHECK-DISABLE-SPMDIZATION-NEXT: br i1 [[THREAD_IS_WORKER]], label [[IS_WORKER_CHECK:%.*]], label [[THREAD_USER_CODE_CHECK:%.*]]
; CHECK-DISABLE-SPMDIZATION: is_worker_check:
-; CHECK-DISABLE-SPMDIZATION-NEXT: [[BLOCK_HW_SIZE:%.*]] = call i32 @__kmpc_get_hardware_num_threads_in_block()
-; CHECK-DISABLE-SPMDIZATION-NEXT: [[WARP_SIZE:%.*]] = call i32 @__kmpc_get_warp_size()
-; CHECK-DISABLE-SPMDIZATION-NEXT: [[BLOCK_SIZE:%.*]] = sub i32 [[BLOCK_HW_SIZE]], [[WARP_SIZE]]
-; CHECK-DISABLE-SPMDIZATION-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[BLOCK_SIZE]]
+; CHECK-DISABLE-SPMDIZATION-NEXT: [[MAX_TEAM_THREADS:%.*]] = call i32 @__kmpc_get_max_team_threads()
+; CHECK-DISABLE-SPMDIZATION-NEXT: [[THREAD_IS_MAIN_OR_WORKER:%.*]] = icmp slt i32 [[TMP0]], [[MAX_TEAM_THREADS]]
; CHECK-DISABLE-SPMDIZATION-NEXT: br i1 [[THREAD_IS_MAIN_OR_WORKER]], label [[WORKER_STATE_MACHINE_BEGIN:%.*]], label [[WORKER_STATE_MACHINE_FINISHED:%.*]]
; CHECK-DISABLE-SPMDIZATION: worker_state_machine.begin:
; CHECK-DISABLE-SPMDIZATION-NEXT: call void @__kmpc_barrier_simple_generic(ptr @[[GLOB1]], i32 [[TMP0]])
diff --git a/openmp/device/include/Interface.h b/openmp/device/include/Interface.h
index 601694871597b..07ad56e8dbef0 100644
--- a/openmp/device/include/Interface.h
+++ b/openmp/device/include/Interface.h
@@ -217,6 +217,9 @@ uint32_t __kmpc_get_hardware_num_threads_in_block();
/// External interface to get the warp size.
uint32_t __kmpc_get_warp_size();
+/// External interface to get the maximum number of threads in a team.
+uint32_t __kmpc_get_max_team_threads();
+
/// Kernel
///
///{
diff --git a/openmp/device/src/Mapping.cpp b/openmp/device/src/Mapping.cpp
index 1096e5f496590..f65ccb5203056 100644
--- a/openmp/device/src/Mapping.cpp
+++ b/openmp/device/src/Mapping.cpp
@@ -175,6 +175,10 @@ extern "C" {
[[gnu::noinline]] uint32_t __kmpc_get_warp_size() {
return mapping::getWarpSize();
}
+
+[[gnu::noinline]] uint32_t __kmpc_get_max_team_threads() {
+ return mapping::getMaxTeamThreads();
+}
}
#define _TGT_KERNEL_LANGUAGE(NAME, MAPPER_NAME) \
>From 2872cc1096b278ea774473b7e881fececf8d3387 Mon Sep 17 00:00:00 2001
From: Larry Meadows <Lawrence.Meadows at amd.com>
Date: Wed, 19 Aug 2026 05:10:16 -0500
Subject: [PATCH 4/5] Say what __kmpc_get_max_team_threads returns
"The maximum number of threads in a team" reads as the block size. It is not:
in generic mode the threads hosting the main thread are not available to the
team, and how many those are is the runtime's to decide, which is the reason
the compiler asks instead of subtracting a warp itself.
Co-authored-by: Cursor <cursoragent at cursor.com>
---
openmp/device/include/Interface.h | 6 +++++-
1 file changed, 5 insertions(+), 1 deletion(-)
diff --git a/openmp/device/include/Interface.h b/openmp/device/include/Interface.h
index 07ad56e8dbef0..dabec09f94524 100644
--- a/openmp/device/include/Interface.h
+++ b/openmp/device/include/Interface.h
@@ -217,7 +217,11 @@ uint32_t __kmpc_get_hardware_num_threads_in_block();
/// External interface to get the warp size.
uint32_t __kmpc_get_warp_size();
-/// External interface to get the maximum number of threads in a team.
+/// External interface to get the number of a block's threads that can execute a
+/// team's parallel region. This is not the block size: in generic mode the
+/// threads hosting the main thread are not available to the team. How many
+/// those are is the runtime's to decide, so callers must ask rather than
+/// subtract for themselves.
uint32_t __kmpc_get_max_team_threads();
/// Kernel
>From 1ffe44784c4d5cf384275e5627327f30818619e8 Mon Sep 17 00:00:00 2001
From: Larry Meadows <Lawrence.Meadows at amd.com>
Date: Fri, 28 Aug 2026 00:15:58 -0500
Subject: [PATCH 5/5] Use live arguments in the new callback test
Co-authored-by: Cursor <cursoragent at cursor.com>
---
.../OpenMP/callback_parallel_regions.ll | 20 +++++++++++++++----
1 file changed, 16 insertions(+), 4 deletions(-)
diff --git a/llvm/test/Transforms/OpenMP/callback_parallel_regions.ll b/llvm/test/Transforms/OpenMP/callback_parallel_regions.ll
index 84c6341d49c06..704b9f00584a6 100644
--- a/llvm/test/Transforms/OpenMP/callback_parallel_regions.ll
+++ b/llvm/test/Transforms/OpenMP/callback_parallel_regions.ll
@@ -164,11 +164,15 @@ 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: call void @p0() #[[ATTR1]]
+; 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:
- call void @p0()
+ %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
}
@@ -176,7 +180,13 @@ 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: call void @outlined(ptr undef, ptr undef) #[[ATTR1]]
+; 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:
@@ -184,6 +194,8 @@ entry:
%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
}
@@ -197,7 +209,7 @@ define weak i32 @__kmpc_target_init(ptr %0, ptr %1) {
}
declare void @opaque_loop_body(i32, ptr)
-declare void @p0() #1
+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)
More information about the Openmp-commits
mailing list