[OpenMP] Folding threadLimit and numThreads when single value in kernels

[josemonsalve2]

Summary

The device runtime contains several calls to __kmpc_get_hardware_num_threads_in_block
and __kmpc_get_hardware_num_blocks. If the thread_limit and the num_teams are constant,
these calls can be folded to the constant value.

In this patch we use the already introduced AAFoldRuntimeCall and the NumTeams and
NumThreads kernel attributes (to be introduced in a different patch) to fold these functions.
The code checks all the kernels, and if their attributes match, the functions are folded.

In the future we will explore specializing for multiple values of NumThreads and NumTeams.

Depends on D106390

[josemonsalve2]

The test is breaking the compiler. I need to change it to see if it passes or not.

[josemonsalve2]

Fixing syntax, removing unnecessary code. Changing a break
to a indicatePessimisticFixpoint()

[tianshilei1992]

Just a comment before reviewing the patch: please don't rebase the patch as D105787 has been reverted.

[josemonsalve2]

In D106033#2880245, @tianshilei1992 wrote:

Just a comment before reviewing the patch: please don't rebase the patch as D105787 has been reverted.

ACK

[JonChesterfield]

Renaming the functions and adding to OMPKinds.def et al looks separate to folding them, should land that first. Makes the functional diff easier to read and reduces the amount of churn if the functional change needs to be backed out.

openmp/libomptarget/deviceRTLs/common/src/omptarget.cu

98

git-clang-format HEAD^ may be useful for things like this

openmp/libomptarget/deviceRTLs/common/support.h

53

Noinline seems bad. Why? Also looks like the functions tagged noinline here are different to the ones that are renamed.

openmp/libomptarget/deviceRTLs/nvptx/src/target_impl.cu

97

rename the functions in amdgcn/src/target_impl.hip too please

[JonChesterfield]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3614

I haven't read through this part, but if we can only fold it to a constant sometimes, we shouldn't mark the calls that survive noinline, as that'll be expensive for the cases that this pass misses.

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3614

We are about to remove noinline from known runtime functions such that we can keep them around until we get to OpenMP-Opt as calls. This will have the effect we want without any drawbacks. Thus, adding noinline in the runitme will be totally fine.

[tianshilei1992]

llvm/include/llvm/Frontend/OpenMP/OMPKinds.def

208

Probably worth a separate patch to add them.

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3632

No need to have else here because the code above already returns.

3635

no need for else as well

3636

Directly indicate pessimistic state because we don't know clearly what the number is.

3646

This has to be updated as SimplifiedValue is not part of BooleanState. Refer to foldIsSPMDExecMode.

3652

Update accordingly

3705

This part needs to be changed. Refer to the trunk for more details. Basically it should be IRPosition::callsite_returned(*CI).

openmp/libomptarget/deviceRTLs/common/src/reduction.cu

200

Since this part is moved to another patch, they should go in this patch.

[josemonsalve2]

Modifying to adapt to most of the comments already in here. Will provide more detail soon. Tests are still failing.

[josemonsalve2]

Ups, missed one comment from Shilei

[josemonsalve2]

llvm/include/llvm/Frontend/OpenMP/OMPKinds.def

208

Do you mean just to add these 4 lines? Or should I add these lines to the front end patch D106298?

[jdoerfert]

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

1766–1767

copy and paste.

3623

if (!CallerKernelInfoAA.ReachingKernelEntries.isValid())
  return indicatePessimisticFixpoint();

also below

3634

use early exit instead,
if (!...)

return ...

3643

no reaching kernels is fine, keep it at none (which is the default) and just return UNCHANGED.

3644

make this function take the string attribute such that we can have a single one and not two.

3702

leftover. (and you can also just pipe an IRPosition into errs())

[josemonsalve2]

• Created a single function receiving a string instead of one per attribute: foldKernelFnAttribute
• Removed pesimisticFixpoint if no kernel is found.
• Adding Check for the ReachingKernelEntries valid state
• Removed unnecesary comments

[josemonsalve2]

Format

[tianshilei1992]

LGTM with one nit. @jdoerfert WDYT?

llvm/include/llvm/Frontend/OpenMP/OMPKinds.def

208

It's better to add them to the patch where they are introduced.

[josemonsalve2]

My test is still failing, but it fails on an assertion on the changeToSPMD method:

opt: /home/josem/postdoc/llvm_stuff/llvm-project/llvm/lib/Transforms/IPO/OpenMPOpt.cpp:2898: bool (anonymous namespace)::AAKernelInfoFunction::changeToSPMDMode(llvm::Attributor &): Assertion `ExecMode && "Kernel without exec mode?"' failed.

I'm trying to debug it, but if it is obvious for any of you let me know if it is something I am doing or I am just triggering another bug

[jdoerfert]

Some minor notes, assuming the tests passes properly after, LGTM

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

3615

3637

You should set CurrentAttrValue = NextAttrValue at the end of the loop.

3642

llvm/test/Transforms/OpenMP/get_hardware_num_threads_in_block_fold.ll

7

you need a variable KERNEL_NAME_exec_mode per kernel here. Look at other tests or what clang generates.

[josemonsalve2]

Fixing final comments from @jdoerfert

[JonChesterfield]

Can we add the pass that removes NOINLINE from these functions before landing this one?

[jhuber6]

In D106033#2905091, @JonChesterfield wrote:

Can we add the pass that removes NOINLINE from these functions before landing this one?

I landed D106482 which does this. One problem is that this was getting run on the .bc libraries and removing it prematurely. @tianshilei1992 was going to make a patch to stop running OpenMPOpt when we build the libraries for this.

[tianshilei1992]

In D106033#2905101, @jhuber6 wrote:

@tianshilei1992 was going to make a patch to stop running OpenMPOpt when we build the libraries for this.

It's already done. ;-)

[JonChesterfield]

Ah, good. I wasn't tagged on D106482 and missed the emails about it.

[josemonsalve2]

Rebasing to main

[josemonsalve2]

Sync to main

[josemonsalve2]

Resync again

[josemonsalve2]

Rebasing to main this time for real

llvm/include/llvm/Frontend/OpenMP/OMPKinds.def

 __OMP_RTL(__kmpc_omp_taskyield, false, Int32, IdentPtr, Int32, /* Int */ Int32)
 __OMP_RTL(__kmpc_push_num_threads, false, Void, IdentPtr, Int32,
           /* Int */ Int32)
 __OMP_RTL(__kmpc_push_proc_bind, false, Void, IdentPtr, Int32, /* Int */ Int32)
 __OMP_RTL(__kmpc_omp_reg_task_with_affinity, false, Int32, IdentPtr, Int32,
           /* kmp_task_t */ VoidPtr, Int32,
           /* kmp_task_affinity_info_t */ VoidPtr)
 
+__OMP_RTL(__kmpc_get_hardware_num_blocks, false, Int32, )
+__OMP_RTL(__kmpc_get_hardware_num_threads_in_block, false, Int32, )
+
 __OMP_RTL(omp_get_thread_num, false, Int32, )
 __OMP_RTL(omp_get_num_threads, false, Int32, )
 __OMP_RTL(omp_get_max_threads, false, Int32, )
 __OMP_RTL(omp_in_parallel, false, Int32, )
 __OMP_RTL(omp_get_dynamic, false, Int32, )
 __OMP_RTL(omp_get_cancellation, false, Int32, )
 __OMP_RTL(omp_get_nested, false, Int32, )
 __OMP_RTL(omp_get_schedule, false, Void, Int32Ptr, Int32Ptr)

 __OMP_RTL_ATTRS(__kmpc_serialized_parallel, InaccessibleArgOnlyAttrs,
                 AttributeSet(), ParamAttrs(ReadOnlyPtrAttrs))
 __OMP_RTL_ATTRS(__kmpc_end_serialized_parallel, InaccessibleArgOnlyAttrs,
                 AttributeSet(), ParamAttrs(ReadOnlyPtrAttrs))
 __OMP_RTL_ATTRS(__kmpc_omp_reg_task_with_affinity, DefaultAttrs, AttributeSet(),
                 ParamAttrs(ReadOnlyPtrAttrs, AttributeSet(), ReadOnlyPtrAttrs,
                            AttributeSet(), ReadOnlyPtrAttrs))
 
+__OMP_RTL_ATTRS(__kmpc_get_hardware_num_blocks, GetterAttrs, AttributeSet(), ParamAttrs())
+__OMP_RTL_ATTRS(__kmpc_get_hardware_num_threads_in_block, GetterAttrs, AttributeSet(), ParamAttrs())
+
 __OMP_RTL_ATTRS(omp_get_thread_num, GetterAttrs, AttributeSet(), ParamAttrs())
 __OMP_RTL_ATTRS(omp_get_num_threads, GetterAttrs, AttributeSet(), ParamAttrs())
 __OMP_RTL_ATTRS(omp_get_max_threads, GetterAttrs, AttributeSet(), ParamAttrs())
 __OMP_RTL_ATTRS(omp_in_parallel, GetterAttrs, AttributeSet(), ParamAttrs())
 __OMP_RTL_ATTRS(omp_get_dynamic, GetterAttrs, AttributeSet(), ParamAttrs())
 __OMP_RTL_ATTRS(omp_get_cancellation, GetterAttrs, AttributeSet(), ParamAttrs())
 __OMP_RTL_ATTRS(omp_get_nested, GetterAttrs, AttributeSet(), ParamAttrs())
 __OMP_RTL_ATTRS(

llvm/lib/Transforms/IPO/OpenMPOpt.cpp

bool runAttributor(bool IsModulePass) {
     ChangeStatus Changed = A.run();
 
     LLVM_DEBUG(dbgs() << "[Attributor] Done with " << SCC.size()
                       << " functions, result: " << Changed << ".\n");
 
     return Changed == ChangeStatus::CHANGED;
   }
 
   /// Populate the Attributor with abstract attribute opportunities in the
   /// function.
+  void registerFoldRuntimeCall(RuntimeFunction RF);
+
+  /// Populate the Attributor with abstract attribute opportunities in the
+  /// function.
   void registerAAs(bool IsModulePass);
 };
 
 Kernel OpenMPOpt::getUniqueKernelFor(Function &F) {
   if (!OMPInfoCache.ModuleSlice.count(&F))
     return nullptr;
 
   // Use a scope to keep the lifetime of the CachedKernel short.
void initialize(Attributor &A) override {
 
     const unsigned int WrapperFunctionArgNo = 6;
     RuntimeFunction RF = It->getSecond();
     switch (RF) {
     // All the functions we know are compatible with SPMD mode.
     case OMPRTL___kmpc_is_spmd_exec_mode:
     case OMPRTL___kmpc_for_static_fini:
     case OMPRTL___kmpc_global_thread_num:
+    case OMPRTL___kmpc_get_hardware_num_threads_in_block:
+    case OMPRTL___kmpc_get_hardware_num_blocks:
     case OMPRTL___kmpc_single:
     case OMPRTL___kmpc_end_single:
     case OMPRTL___kmpc_master:
     case OMPRTL___kmpc_end_master:
     case OMPRTL___kmpc_barrier:
       break;
     case OMPRTL___kmpc_for_static_init_4:
     case OMPRTL___kmpc_for_static_init_4u:
A.registerSimplificationCallback(
               A.recordDependence(*this, *AA, DepClassTy::OPTIONAL);
           }
           return SimplifiedValue;
         });
   }
 
   ChangeStatus updateImpl(Attributor &A) override {
     ChangeStatus Changed = ChangeStatus::UNCHANGED;
-
     switch (RFKind) {
     case OMPRTL___kmpc_is_spmd_exec_mode:
       Changed = Changed | foldIsSPMDExecMode(A);
       break;
+    case OMPRTL___kmpc_get_hardware_num_threads_in_block:
+      Changed = Changed | foldHardwareNumThreads(A);
+      break;
+    case OMPRTL___kmpc_get_hardware_num_blocks:
+      Changed = Changed | foldHardwareNumTeams(A);
+      break;
     default:
       llvm_unreachable("Unhandled OpenMP runtime function!");
     }
 
     return Changed;
   }
 
   ChangeStatus manifest(Attributor &A) override {
private:
   ChangeStatus foldIsSPMDExecMode(Attributor &A) {
     BooleanState StateBefore = getState();
 
     unsigned AssumedSPMDCount = 0, KnownSPMDCount = 0;
     unsigned AssumedNonSPMDCount = 0, KnownNonSPMDCount = 0;
     auto &CallerKernelInfoAA = A.getAAFor<AAKernelInfo>(
         *this, IRPosition::function(*getAnchorScope()), DepClassTy::REQUIRED);
 
+    // What actual kernels reach this function
     for (Kernel K : CallerKernelInfoAA.ReachingKernelEntries) {
       auto &AA = A.getAAFor<AAKernelInfo>(*this, IRPosition::function(*K),
                                           DepClassTy::REQUIRED);
 
       if (!AA.isValidState()) {
         SimplifiedValue = nullptr;
         return indicatePessimisticFixpoint();
       }
if (KnownSPMDCount || AssumedSPMDCount) {
       // calls to __kmpc_is_spmd_exec_mode to 0.
       SimplifiedValue = ConstantInt::get(Type::getInt8Ty(Ctx), false);
     }
 
     return getState() == StateBefore ? ChangeStatus::UNCHANGED
                                      : ChangeStatus::CHANGED;
   }
 
+  /// Fold __kmpc_get_hardware_num_blocks into a constant if possible.
+  /// The value is an attribute in the kernel
+  ChangeStatus foldHardwareNumTeams(Attributor &A) {
+    // Specialize only if all the calls agree with the number of teams
+    int32_t CurrentNumTeams = -1;
+    BooleanState StateBefore = getState();
+
+    auto &CallerKernelInfoAA = A.getAAFor<AAKernelInfo>(
+        *this, IRPosition::function(*getAnchorScope()), DepClassTy::REQUIRED);
+
+    // What actual kernels reach this function
+    for (Kernel K : CallerKernelInfoAA.ReachingKernelEntries) {
+      if (K->hasFnAttribute("NumTeams")) {
+        int32_t NumT = 
+            std::stoi(K->getFnAttribute("NumTeams").getValueAsString().str());
+        if (CurrentNumTeams != -1 && CurrentNumTeams != NumT) {
+          SimplifiedValue = nullptr;
+          return indicatePessimisticFixpoint();
+        } else {
+          CurrentNumTeams = NumT;
+        }
+      } else {
+        // TODO: No attribute, then default?
+      }
+    }
+
+    if (CurrentNumTeams != -1) {
+      auto &Ctx = getAnchorValue().getContext();
+      SimplifiedValue = 
+          ConstantInt::get(Type::getInt8Ty(Ctx), CurrentNumTeams);
+    }
+
+    return getState() == StateBefore ? ChangeStatus::UNCHANGED
+                                     : ChangeStatus::CHANGED;
+  }
+
+  /// Fold __kmpc_get_hardware_num_threads_in_block into a constant if possible.
+  /// The value is an attribute in the kernel
+  ChangeStatus foldHardwareNumThreads(Attributor &A) {
+    // Specialize only if all the calls agree with the number of threads
+    int32_t CurrentNumThreads = -1;
+    BooleanState StateBefore = getState();
+
+    auto &CallerKernelInfoAA = A.getAAFor<AAKernelInfo>(
+        *this, IRPosition::function(*getAnchorScope()), DepClassTy::REQUIRED);
+
+    // What actual kernels reach this function
+    for (Kernel K : CallerKernelInfoAA.ReachingKernelEntries) {
+      if (K->hasFnAttribute("NumThreads")) {
+        int32_t NumT = 
+            std::stoi(K->getFnAttribute("NumThreads").getValueAsString().str());
+        if (CurrentNumThreads != -1 && CurrentNumThreads != NumT) {
+          SimplifiedValue = nullptr;
+          return indicatePessimisticFixpoint();
+        } else {
+          CurrentNumThreads = NumT;
+        }
+      } else {
+        // TODO: No attribute, then default?
+      }
+    }
+
+    if (CurrentNumThreads != -1) {
+      auto &Ctx = getAnchorValue().getContext();
+      SimplifiedValue = 
+          ConstantInt::get(Type::getInt8Ty(Ctx), CurrentNumThreads);
+    }
+    
+    return getState() == StateBefore ? ChangeStatus::UNCHANGED
+                                     : ChangeStatus::CHANGED;
+  }
+
   /// An optional value the associated value is assumed to fold to. That is, we
   /// assume the associated value (which is a call) can be replaced by this
   /// simplified value.
   Optional<Value *> SimplifiedValue;
 
   /// The runtime function kind of the callee of the associated call site.
   RuntimeFunction RFKind;
 };
 
 } // namespace
 
+/// Register folding callsite
+void OpenMPOpt::registerFoldRuntimeCall(RuntimeFunction RF) {
+  auto &RFI = OMPInfoCache.RFIs[RF];
+  RFI.foreachUse(SCC, [&](Use &U, Function &) {
+    CallInst *CI = OpenMPOpt::getCallIfRegularCall(U, &RFI);
+    if (!CI)
+      return false;
+    A.getOrCreateAAFor<AAFoldRuntimeCall>(
+        IRPosition::callsite_function(*CI), /* QueryingAA */ nullptr,
+        DepClassTy::NONE, /* ForceUpdate */ false,
+        /* UpdateAfterInit */ false);
+    return false;
+  });
+}
+
 void OpenMPOpt::registerAAs(bool IsModulePass) {
   if (SCC.empty())
 
     return;
   if (IsModulePass) {
     // Ensure we create the AAKernelInfo AAs first and without triggering an
     // update. This will make sure we register all value simplification
     // callbacks before any other AA has the chance to create an AAValueSimplify
     // or similar.
-    for (Function *Kernel : OMPInfoCache.Kernels)
+    for (Function *Kernel : OMPInfoCache.Kernels) {
       A.getOrCreateAAFor<AAKernelInfo>(
           IRPosition::function(*Kernel), /* QueryingAA */ nullptr,
           DepClassTy::NONE, /* ForceUpdate */ false,
           /* UpdateAfterInit */ false);
+    }
 
-    auto &IsSPMDRFI = OMPInfoCache.RFIs[OMPRTL___kmpc_is_spmd_exec_mode];
-    IsSPMDRFI.foreachUse(SCC, [&](Use &U, Function &) {
-      CallInst *CI = OpenMPOpt::getCallIfRegularCall(U, &IsSPMDRFI);
-      if (!CI)
-        return false;
-      A.getOrCreateAAFor<AAFoldRuntimeCall>(
-          IRPosition::callsite_function(*CI), /* QueryingAA */ nullptr,
-          DepClassTy::NONE, /* ForceUpdate */ false,
-          /* UpdateAfterInit */ false);
-      return false;
-    });
+    registerFoldRuntimeCall(OMPRTL___kmpc_is_spmd_exec_mode);
+    registerFoldRuntimeCall(OMPRTL___kmpc_get_hardware_num_threads_in_block);
+    registerFoldRuntimeCall(OMPRTL___kmpc_get_hardware_num_blocks);
   }
 
   // Create CallSite AA for all Getters.
   for (int Idx = 0; Idx < OMPInfoCache.ICVs.size() - 1; ++Idx) {
     auto ICVInfo = OMPInfoCache.ICVs[static_cast<InternalControlVar>(Idx)];
 
     auto &GetterRFI = OMPInfoCache.RFIs[ICVInfo.Getter];
 

llvm/test/Transforms/OpenMP/get_hardware_num_threads_in_block_fold.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature --check-globals
+; RUN: opt -S -passes=openmp-opt < %s | FileCheck %s
+target triple = "nvptx64"
+
+%struct.ident_t = type { i32, i32, i32, i32, i8* }
+
+@G = external global i32
+;.
+; CHECK: @[[G:[a-zA-Z0-9_$"\\.-]+]] = external global i8
+;.
+define weak void @kernel0() #0 {
+; CHECK-LABEL: define {{[^@]+}}@kernel0() #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:    [[I:%.*]] = call i32 @__kmpc_target_init(%struct.ident_t* null, i1 true, i1 false, i1 false)
+; CHECK-NEXT:    call void @helper0()
+; CHECK-NEXT:    call void @helper1()
+; CHECK-NEXT:    call void @helper2()
+; CHECK-NEXT:    call void @__kmpc_target_deinit(%struct.ident_t* null, i1 true, i1 false)
+; CHECK-NEXT:    ret void
+;
+  %i = call i32 @__kmpc_target_init(%struct.ident_t* null, i1 true, i1 false, i1 false)
+  call void @helper0()
+  call void @helper1()
+  call void @helper2()
+  call void @__kmpc_target_deinit(%struct.ident_t* null, i1 true, i1 false)
+  ret void
+}
+
+define weak void @kernel1() #0 {
+; CHECK-LABEL: define {{[^@]+}}@kernel1() #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:    [[I:%.*]] = call i32 @__kmpc_target_init(%struct.ident_t* null, i1 true, i1 false, i1 false)
+; CHECK-NEXT:    call void @helper1()
+; CHECK-NEXT:    call void @__kmpc_target_deinit(%struct.ident_t* null, i1 true, i1 false)
+; CHECK-NEXT:    ret void
+;
+  %i = call i32 @__kmpc_target_init(%struct.ident_t* null, i1 false, i1 false, i1 false)
+  call void @helper1()
+  call void @__kmpc_target_deinit(%struct.ident_t* null, i1 false, i1 false)
+  ret void
+}
+
+define weak void @kernel2() #0 {
+; CHECK-LABEL: define {{[^@]+}}@kernel2() #[[ATTR0:[0-9]+]] {
+; CHECK-NEXT:    [[I:%.*]] = call i32 @__kmpc_target_init(%struct.ident_t* null, i1 false, i1 false, i1 false)
+; CHECK-NEXT:    call void @helper0()
+; CHECK-NEXT:    call void @helper1()
+; CHECK-NEXT:    call void @helper2()
+; CHECK-NEXT:    call void @__kmpc_target_deinit(%struct.ident_t* null, i1 false, i1 false)
+; CHECK-NEXT:    ret void
+;
+  %i = call i32 @__kmpc_target_init(%struct.ident_t* null, i1 false, i1 false, i1 false)
+  call void @helper0()
+  call void @helper1()
+  call void @helper2()
+  call void @__kmpc_target_deinit(%struct.ident_t* null, i1 false, i1 false)
+  ret void
+}
+
+define internal void @helper0() {
+; CHECK-LABEL: define {{[^@]+}}@helper0() {
+; CHECK-NEXT:    store i32 666, i32* @G, align 1
+; CHECK-NEXT:    ret void
+;
+  %threadLimit = call i32 @__kmpc_get_hardware_num_threads_in_block()
+  store i32 %threadLimit, i32* @G
+  ret void
+}
+
+define internal void @helper1() {
+; CHECK-LABEL: define {{[^@]+}}@helper1() {
+; CHECK-NEXT:    br label [[F:%.*]]
+; CHECK:       t:
+; CHECK-NEXT:    unreachable
+; CHECK:       f:
+; CHECK-NEXT:    ret void
+;
+  %threadLimit = call i32 @__kmpc_get_hardware_num_threads_in_block()
+  %c = icmp eq i32 %threadLimit, 666
+  br i1 %c, label %t, label %f
+t:
+  call void @helper0()
+  ret void
+f:
+  ret void
+}
+
+define internal void @helper2() {
+; CHECK-LABEL: define {{[^@]+}}@helper2() {
+; CHECK-NEXT:    store i32 666, i8* @G
+; CHECK-NEXT:    ret void
+;
+  %threadLimit = call i32 @__kmpc_get_hardware_num_threads_in_block()
+  store i32 %threadLimit, i32* @G
+  ret void
+}
+
+declare i32 @__kmpc_get_hardware_num_threads_in_block()
+declare i32 @__kmpc_target_init(%struct.ident_t*, i1 zeroext, i1 zeroext, i1 zeroext) #1
+declare void @__kmpc_target_deinit(%struct.ident_t* nocapture readnone, i1 zeroext, i1 zeroext) #1
+
+
+!llvm.module.flags = !{!0, !1}
+!nvvm.annotations = !{!2, !3, !4}
+
+attributes #0 = { "NumThreads"="666" "NumTeams"="777"}
+
+!0 = !{i32 7, !"openmp", i32 50}
+!1 = !{i32 7, !"openmp-device", i32 50}
+!2 = !{void ()* @kernel0, !"kernel", i32 1}
+!3 = !{void ()* @kernel1, !"kernel", i32 1}
+!4 = !{void ()* @kernel2, !"kernel", i32 1}
+;.
+; CHECK: attributes #[[ATTR0:[0-9]+]] = { "llvm.assume"="ompx_spmd_amenable" }
+;.
+; CHECK: [[META0:![0-9]+]] = !{i32 7, !"openmp", i32 50}
+; CHECK: [[META1:![0-9]+]] = !{i32 7, !"openmp-device", i32 50}
+; CHECK: [[META2:![0-9]+]] = !{void ()* @kernel0, !"kernel", i32 1}
+; CHECK: [[META3:![0-9]+]] = !{void ()* @kernel1, !"kernel", i32 1}
+; CHECK: [[META4:![0-9]+]] = !{void ()* @kernel2, !"kernel", i32 1}
+;.

openmp/libomptarget/deviceRTLs/amdgcn/src/target_impl.hip

 }
 uint32_t get_workgroup_dim(uint32_t group_id, uint32_t grid_size,
                            uint16_t group_size) {
   uint32_t r = grid_size - group_id * group_size;
   return (r < group_size) ? r : group_size;
 }
 } // namespace
 
-EXTERN int GetNumberOfBlocksInKernel() {
+EXTERN int __kmpc_get_hardware_num_blocks() {
   return get_grid_dim(__builtin_amdgcn_grid_size_x(),
                       __builtin_amdgcn_workgroup_size_x());
 }
 
-EXTERN int GetNumberOfThreadsInBlock() {
+EXTERN int __kmpc_get_hardware_num_threads_in_block() {
   return get_workgroup_dim(__builtin_amdgcn_workgroup_id_x(),
                            __builtin_amdgcn_grid_size_x(),
                            __builtin_amdgcn_workgroup_size_x());
 }
 
 EXTERN unsigned GetWarpId() { return GetThreadIdInBlock() / WARPSIZE; }
 EXTERN unsigned GetWarpSize() { return WARPSIZE; }
 EXTERN unsigned GetLaneId() {
   return __builtin_amdgcn_mbcnt_hi(~0u, __builtin_amdgcn_mbcnt_lo(~0u, 0u));
 }
 
 EXTERN uint32_t __kmpc_amdgcn_gpu_num_threads() {
-  return GetNumberOfThreadsInBlock();
+  return __kmpc_get_hardware_num_threads_in_block();
 }
 
 // Atomics
 uint32_t __kmpc_atomic_add(uint32_t *Address, uint32_t Val) {
   return __atomic_fetch_add(Address, Val, __ATOMIC_SEQ_CST);
 }
 uint32_t __kmpc_atomic_inc(uint32_t *Address, uint32_t Val) {
   return __builtin_amdgcn_atomic_inc32(Address, Val, __ATOMIC_SEQ_CST, "");

openmp/libomptarget/deviceRTLs/common/src/libcall.cu

if (parallelLevel[GetWarpId()] == 0)
     rc = nThreads;
   ASSERT0(LT_FUSSY, rc >= 0, "bad number of threads");
   PRINT(LD_IO, "call omp_get_max_threads() return %d\n", rc);
   return rc;
 }
 
 EXTERN int omp_get_thread_limit(void) {
   if (__kmpc_is_spmd_exec_mode())
-    return GetNumberOfThreadsInBlock();
+    return __kmpc_get_hardware_num_threads_in_block();
   int rc = threadLimit;
   PRINT(LD_IO, "call omp_get_thread_limit() return %d\n", rc);
   return rc;
 }
 
 EXTERN int omp_get_thread_num() {
   int rc = GetOmpThreadId();
   PRINT(LD_IO, "call omp_get_thread_num() returns %d\n", rc);
EXTERN int omp_get_ancestor_thread_num(int level) {
   }
   PRINT(LD_IO, "call omp_get_ancestor_thread_num(level %d) returns %d\n", level,
         rc)
   return rc;
 }
 
 EXTERN int omp_get_team_size(int level) {
   if (__kmpc_is_spmd_exec_mode())
-    return level == 1 ? GetNumberOfThreadsInBlock() : 1;
+    return level == 1 ? __kmpc_get_hardware_num_threads_in_block() : 1;
   int rc = -1;
   unsigned parLevel = parallelLevel[GetWarpId()];
   // If level is 0 or all parallel regions are not active - return 1.
   if (level == 1 && parLevel > OMP_ACTIVE_PARALLEL_LEVEL) {
     rc = threadsInTeam;
   } else if (level == 0 ||
              (level > 0 && parLevel < OMP_ACTIVE_PARALLEL_LEVEL &&
               level <= parLevel) ||

openmp/libomptarget/deviceRTLs/common/src/omptarget.cu

setExecutionParameters(Spmd, RequiresFullRuntime ? RuntimeInitialized
                          : RuntimeUninitialized);
   int threadId = GetThreadIdInBlock();
   if (threadId == 0) {
     usedSlotIdx = __kmpc_impl_smid() % MAX_SM;
   }
 
   if (GetLaneId() == 0) {
     parallelLevel[GetWarpId()] =
-        1 + (GetNumberOfThreadsInBlock() > 1 ? OMP_ACTIVE_PARALLEL_LEVEL : 0);
+        1 + (__kmpc_get_hardware_num_threads_in_block() > 1 ? OMP_ACTIVE_PARALLEL_LEVEL : 0);
   }
 
   __kmpc_data_sharing_init_stack();
   if (!RequiresFullRuntime)
     return;
 
   //
   // Team Context Initialization.
   //

openmp/libomptarget/deviceRTLs/common/src/reduction.cu

EXTERN int32_t __kmpc_nvptx_teams_reduce_nowait_v2(
 
   // In non-generic mode all workers participate in the teams reduction.
   // In generic mode only the team master participates in the teams
   // reduction because the workers are waiting for parallel work.
   uint32_t NumThreads =
       __kmpc_is_spmd_exec_mode() ? GetNumberOfOmpThreads(/*isSPMDExecutionMode=*/true)
                          : /*Master thread only*/ 1;
   uint32_t TeamId = GetBlockIdInKernel();
-  uint32_t NumTeams = GetNumberOfBlocksInKernel();
+  uint32_t NumTeams = __kmpc_get_hardware_num_blocks();
   static unsigned SHARED(Bound);
   static unsigned SHARED(ChunkTeamCount);
 
   // Block progress for teams greater than the current upper
   // limit. We always only allow a number of teams less or equal
   // to the number of slots in the buffer.
   bool IsMaster = isMaster(loc, ThreadId);
   while (IsMaster) {

openmp/libomptarget/deviceRTLs/common/src/support.cu

 // Thread id is 0 indexed.
 // E.g: If NumThreads is 33, master id is 32.
 //      If NumThreads is 64, master id is 32.
 //      If NumThreads is 97, master id is 96.
 //      If NumThreads is 1024, master id is 992.
 //
 // Called in Generic Execution Mode only.
 int GetMasterThreadID() {
-  return (GetNumberOfThreadsInBlock() - 1) & ~(WARPSIZE - 1);
+  return (__kmpc_get_hardware_num_threads_in_block() - 1) & ~(WARPSIZE - 1);
 }
 
 // The last warp is reserved for the master; other warps are workers.
 // Called in Generic Execution Mode only.
 int GetNumberOfWorkersInTeam() { return GetMasterThreadID(); }
 
 ////////////////////////////////////////////////////////////////////////////////
 // get thread id in team

 
 int GetNumberOfOmpThreads(bool isSPMDExecutionMode) {
   // omp_num_threads
   int rc;
   int Level = parallelLevel[GetWarpId()];
   if (Level != OMP_ACTIVE_PARALLEL_LEVEL + 1) {
     rc = 1;
   } else if (isSPMDExecutionMode) {
-    rc = GetNumberOfThreadsInBlock();
+    rc = __kmpc_get_hardware_num_threads_in_block();
   } else {
     rc = threadsInTeam;
   }
 
   return rc;
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 // Team id linked to OpenMP
 
 int GetOmpTeamId() {
   // omp_team_num
   return GetBlockIdInKernel(); // assume 1 block per team
 }
 
 int GetNumberOfOmpTeams() {
   // omp_num_teams
-  return GetNumberOfBlocksInKernel(); // assume 1 block per team
+  return __kmpc_get_hardware_num_blocks(); // assume 1 block per team
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 // Masters
 
 int IsTeamMaster(int ompThreadId) { return (ompThreadId == 0); }
 
 ////////////////////////////////////////////////////////////////////////////////

 
 ////////////////////////////////////////////////////////////////////////////////
 // get OpenMP number of procs
 
 // Get the number of processors in the device.
 int GetNumberOfProcsInDevice(bool isSPMDExecutionMode) {
   if (!isSPMDExecutionMode)
     return GetNumberOfWorkersInTeam();
-  return GetNumberOfThreadsInBlock();
+  return __kmpc_get_hardware_num_threads_in_block();
 }
 
 int GetNumberOfProcsInTeam(bool isSPMDExecutionMode) {
   return GetNumberOfProcsInDevice(isSPMDExecutionMode);
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 // Memory

openmp/libomptarget/deviceRTLs/common/support.h

 //===--------- support.h - OpenMP GPU support functions ---------- CUDA -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // Wrapper to some functions natively supported by the GPU.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef OMPTARGET_SUPPORT_H
 #define OMPTARGET_SUPPORT_H
 
 #include "interface.h"
 #include "target_impl.h"
 
 ////////////////////////////////////////////////////////////////////////////////
 // Execution Parameters
 ////////////////////////////////////////////////////////////////////////////////
 enum ExecutionMode {
   Spmd = 0x00u,
   Generic = 0x01u,

 int GetMasterThreadID();
 int GetNumberOfWorkersInTeam();
 
 // get OpenMP thread and team ids
 int GetOmpThreadId();                         // omp_thread_num
 int GetOmpTeamId();                           // omp_team_num
 
 // get OpenMP number of threads and team
-int GetNumberOfOmpThreads(bool isSPMDExecutionMode); // omp_num_threads
-int GetNumberOfOmpTeams();                           // omp_num_teams
+NOINLINE int GetNumberOfOmpThreads(bool isSPMDExecutionMode); // omp_num_threads
+NOINLINE int GetNumberOfOmpTeams();                           // omp_num_teams
 
 // get OpenMP number of procs
 int GetNumberOfProcsInTeam(bool isSPMDExecutionMode);
 int GetNumberOfProcsInDevice(bool isSPMDExecutionMode);
 
 // masters
 int IsTeamMaster(int ompThreadId);
 

openmp/libomptarget/deviceRTLs/nvptx/src/target_impl.cu

 
 EXTERN void __kmpc_impl_threadfence() { __nvvm_membar_gl(); }
 EXTERN void __kmpc_impl_threadfence_block() { __nvvm_membar_cta(); }
 EXTERN void __kmpc_impl_threadfence_system() { __nvvm_membar_sys(); }
 
 // Calls to the NVPTX layer (assuming 1D layout)
 EXTERN int GetThreadIdInBlock() { return __nvvm_read_ptx_sreg_tid_x(); }
 EXTERN int GetBlockIdInKernel() { return __nvvm_read_ptx_sreg_ctaid_x(); }
-EXTERN int GetNumberOfBlocksInKernel() {
+EXTERN int __kmpc_get_hardware_num_blocks() {
   return __nvvm_read_ptx_sreg_nctaid_x();
 }
-EXTERN int GetNumberOfThreadsInBlock() { return __nvvm_read_ptx_sreg_ntid_x(); }
+EXTERN int __kmpc_get_hardware_num_threads_in_block() { return __nvvm_read_ptx_sreg_ntid_x(); }
 EXTERN unsigned GetWarpId() { return GetThreadIdInBlock() / WARPSIZE; }
 EXTERN unsigned GetWarpSize() { return WARPSIZE; }
 EXTERN unsigned GetLaneId() { return GetThreadIdInBlock() & (WARPSIZE - 1); }
 
 // Atomics
 uint32_t __kmpc_atomic_add(uint32_t *Address, uint32_t Val) {
   return __atomic_fetch_add(Address, Val, __ATOMIC_SEQ_CST);
 }

openmp/libomptarget/deviceRTLs/target_interface.h

 //===------------- target_interface.h - Target interfaces --------- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file contains interfaces that must be implemented by each target.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef _OMPTARGET_TARGET_INTERFACE_H_
 #define _OMPTARGET_TARGET_INTERFACE_H_
 
 #include "target_impl.h"
 
 // Calls to the NVPTX layer (assuming 1D layout)
 EXTERN int GetThreadIdInBlock();
 EXTERN int GetBlockIdInKernel();
-EXTERN int GetNumberOfBlocksInKernel();
-EXTERN int GetNumberOfThreadsInBlock();
+EXTERN NOINLINE int __kmpc_get_hardware_num_blocks();
+EXTERN NOINLINE int __kmpc_get_hardware_num_threads_in_block();
 EXTERN unsigned GetWarpId();
 EXTERN unsigned GetWarpSize();
 EXTERN unsigned GetLaneId();
 
 // Atomics
 uint32_t __kmpc_atomic_add(uint32_t *, uint32_t);
 uint32_t __kmpc_atomic_inc(uint32_t *, uint32_t);
 uint32_t __kmpc_atomic_max(uint32_t *, uint32_t);