[OpenMP] Use default grid value for static grid size

[jdoerfert]

Summary

If the user did not provide any static clause to override the grid size,
we assume the default grid size as upper bound and use it to improve
code generation through vendor specific attributes.

Fixes: https://github.com/llvm/llvm-project/issues/64816

[jhuber6]

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

4134

Can't we just get the module from the Function and check the triple?

[tianshilei1992]

Is amdgpu-flat-work-group-size a bound or exact value?

[jdoerfert]

In D158382#4603950, @tianshilei1992 wrote:

Is amdgpu-flat-work-group-size a bound or exact value?

It's a range. What really matters is the upper bound, that is the value you see in the test (max flat ...).

[jdoerfert]

Added checks.

[jhuber6]

LG, thanks

[dhruvachak]

This patch produces the following difference in IR out of CodeGen.

Without this patch:

%nvptx_num_threads.i = tail call i32 @__kmpc_get_hardware_num_threads_in_block() #2
call void @__kmpc_distribute_static_init_4(ptr addrspacecast (ptr addrspace(1) @2 to ptr), i32 %1, i32 91, ptr nonnull %.omp.is_last.ascast.i, ptr nonnull %.omp.comb.lb.ascast.i, ptr nonnull %.omp.comb.ub.ascast.i, ptr nonnull %.omp.stride.ascast.i, i32 1, i32 %nvptx_num_threads.i) #2

With this patch:

call void @__kmpc_distribute_static_init_4(ptr addrspacecast (ptr addrspace(1) @2 to ptr), i32 %1, i32 91, ptr nonnull %.omp.is_last.ascast.i, ptr nonnull %.omp.comb.lb.ascast.i, ptr nonnull %.omp.comb.ub.ascast.i, ptr nonnull %.omp.stride.ascast.i, i32 1, i32 256) #2

Setting the blocksize to a constant too early would be a problem if the runtime changes the blocksize, e.g. because of an environment variable or because of a low trip count (D152014). Comments? @jdoerfert

[jdoerfert]

In D158382#4621885, @dhruvachak wrote:

This patch produces the following difference in IR out of CodeGen.

Without this patch:

%nvptx_num_threads.i = tail call i32 @__kmpc_get_hardware_num_threads_in_block() #2
call void @__kmpc_distribute_static_init_4(ptr addrspacecast (ptr addrspace(1) @2 to ptr), i32 %1, i32 91, ptr nonnull %.omp.is_last.ascast.i, ptr nonnull %.omp.comb.lb.ascast.i, ptr nonnull %.omp.comb.ub.ascast.i, ptr nonnull %.omp.stride.ascast.i, i32 1, i32 %nvptx_num_threads.i) #2

With this patch:

call void @__kmpc_distribute_static_init_4(ptr addrspacecast (ptr addrspace(1) @2 to ptr), i32 %1, i32 91, ptr nonnull %.omp.is_last.ascast.i, ptr nonnull %.omp.comb.lb.ascast.i, ptr nonnull %.omp.comb.ub.ascast.i, ptr nonnull %.omp.stride.ascast.i, i32 1, i32 256) #2

Setting the blocksize to a constant too early would be a problem if the runtime changes the blocksize, e.g. because of an environment variable or because of a low trip count (D152014). Comments? @jdoerfert

From OpenMP-Opt:

case OMPRTL___kmpc_get_hardware_num_threads_in_block:
   Changed = Changed | foldKernelFnAttribute(A, "omp_target_thread_limit");
   break;

this is wrong. We should fold thread limit, not num_threads_in_block.
The latter can only be folded if we will not lower it, which we currently cannot guarantee.

llvm/lib/Frontend/OpenMP/OMPIRBuilder.cpp

 #include "llvm/ADT/StringRef.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/CodeMetrics.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/ScalarEvolution.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Bitcode/BitcodeReader.h"
+#include "llvm/Frontend/OpenMP/OMPGridValues.h"
 #include "llvm/IR/Attributes.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/CallingConv.h"
 #include "llvm/IR/Constant.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/GlobalVariable.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/PassManager.h"
 #include "llvm/IR/Value.h"
 #include "llvm/MC/TargetRegistry.h"
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/Target/TargetOptions.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/CodeExtractor.h"
 #include "llvm/Transforms/Utils/LoopPeel.h"
 #include "llvm/Transforms/Utils/UnrollLoop.h"
if (!updateToLocation(Loc))
     return;
 
   Function *Fn = getOrCreateRuntimeFunctionPtr(
       omp::RuntimeFunction::OMPRTL___kmpc_target_deinit);
 
   Builder.CreateCall(Fn, {});
 }
 
+static const omp::GV &getGridValue(Function *Kernel) {
+  if (Kernel->getCallingConv() == CallingConv::AMDGPU_KERNEL) {
+    StringRef Features =
+        Kernel->getFnAttribute("target-features").getValueAsString();
+    if (Features.count("+wavefrontsize64"))
+      return omp::getAMDGPUGridValues<64>();
+    return omp::getAMDGPUGridValues<32>();
+  }
+  if (Triple(Kernel->getParent()->getTargetTriple()).isNVPTX())
+
+    return omp::NVPTXGridValues;
+  llvm_unreachable("No grid value available for this architecture!");
+}
+
 void OpenMPIRBuilder::setOutlinedTargetRegionFunctionAttributes(
     Function *OutlinedFn, int32_t NumTeams, int32_t NumThreads) {
   if (Config.isTargetDevice()) {
     OutlinedFn->setLinkage(GlobalValue::WeakODRLinkage);
     // TODO: Determine if DSO local can be set to true.
     OutlinedFn->setDSOLocal(false);
     OutlinedFn->setVisibility(GlobalValue::ProtectedVisibility);
     if (Triple(M.getTargetTriple()).isAMDGCN())
       OutlinedFn->setCallingConv(CallingConv::AMDGPU_KERNEL);
   }
 
   if (NumTeams > 0)
     OutlinedFn->addFnAttr("omp_target_num_teams", std::to_string(NumTeams));
 
+  if (NumThreads == -1 && Config.isGPU())
+    NumThreads = getGridValue(OutlinedFn).GV_Default_WG_Size;
+
   if (NumThreads > 0) {
     if (OutlinedFn->getCallingConv() == CallingConv::AMDGPU_KERNEL) {
       OutlinedFn->addFnAttr("amdgpu-flat-work-group-size",
                             "1," + llvm::utostr(NumThreads));
     } else {
       // TODO: Modify or create "maxntidx" module metadata.
     }
     OutlinedFn->addFnAttr("omp_target_thread_limit",

openmp/libomptarget/test/offloading/default_thread_limit.c

+// clang-format off
+// RUN: %libomptarget-compile-generic
+// RUN: env LIBOMPTARGET_INFO=16 \
+// RUN:   %libomptarget-run-generic 2>&1 | %fcheck-generic --check-prefix=DEFAULT
+
+// UNSUPPORTED: nvptx64-nvidia-cuda
+// UNSUPPORTED: nvptx64-nvidia-cuda-LTO
+// UNSUPPORTED: aarch64-unknown-linux-gnu
+// UNSUPPORTED: aarch64-unknown-linux-gnu-LTO
+// UNSUPPORTED: x86_64-pc-linux-gnu
+// UNSUPPORTED: x86_64-pc-linux-gnu-LTO
+
+__attribute__((optnone)) int optnone() { return 1; }
+
+int main() {
+  int N = optnone() * 4098 * 32;
+
+// DEFAULT: [[NT:(128|256)]] (MaxFlatWorkGroupSize: [[NT]]
+#pragma omp target teams distribute parallel for simd
+  for (int i = 0; i < N; ++i) {
+    optnone();
+  }
+// DEFAULT: [[NT:(128|256)]] (MaxFlatWorkGroupSize: [[NT]]
+#pragma omp target teams distribute parallel for simd
+  for (int i = 0; i < N; ++i) {
+    optnone();
+  }
+// DEFAULT: [[NT:(128|256)]] (MaxFlatWorkGroupSize: [[NT]]
+#pragma omp target teams distribute parallel for simd
+  for (int i = 0; i < N; ++i) {
+    optnone();
+  }
+// DEFAULT: [[NT:(128|256)]] (MaxFlatWorkGroupSize: [[NT]]
+#pragma omp target 
+#pragma omp teams distribute parallel for
+  for (int i = 0; i < N; ++i) {
+    optnone();
+  }
+// DEFAULT: 42 (MaxFlatWorkGroupSize: 1024
+#pragma omp target thread_limit(optnone() * 42)
+#pragma omp teams distribute parallel for
+  for (int i = 0; i < N; ++i) {
+    optnone();
+  }
+// DEFAULT: 42 (MaxFlatWorkGroupSize: 42
+#pragma omp target thread_limit(optnone() * 42) ompx_attribute(__attribute__((amdgpu_flat_work_group_size(42, 42))))
+#pragma omp teams distribute parallel for
+  for (int i = 0; i < N; ++i) {
+    optnone();
+  }
+// FIXME: Use the attribute value to imply a thread_limit
+// DEFAULT: {{(128|256)}} (MaxFlatWorkGroupSize: 42
+#pragma omp target ompx_attribute(__attribute__((amdgpu_flat_work_group_size(42, 42))))
+#pragma omp teams distribute parallel for
+  for (int i = 0; i < N; ++i) {
+    optnone();
+  }
+// DEFAULT: MaxFlatWorkGroupSize: 1024
+#pragma omp target
+#pragma omp teams distribute parallel for num_threads(optnone() * 42)
+  for (int i = 0; i < N; ++i) {
+    optnone();
+  }
+// DEFAULT: MaxFlatWorkGroupSize: 1024
+#pragma omp target teams distribute parallel for thread_limit(optnone() * 42)
+  for (int i = 0; i < N; ++i) {
+    optnone();
+  }
+// DEFAULT: MaxFlatWorkGroupSize: 1024
+#pragma omp target teams distribute parallel for num_threads(optnone() * 42)
+  for (int i = 0; i < N; ++i) {
+    optnone();
+  }
+// DEFAULT: 9 (MaxFlatWorkGroupSize: 9
+#pragma omp target
+#pragma omp teams distribute parallel for num_threads(9)
+  for (int i = 0; i < N; ++i) {
+    optnone();
+  }
+// DEFAULT: 4 (MaxFlatWorkGroupSize: 4
+#pragma omp target thread_limit(4)
+#pragma omp teams distribute parallel for
+  for (int i = 0; i < N; ++i) {
+    optnone();
+  }
+// DEFAULT: 4 (MaxFlatWorkGroupSize: 4
+#pragma omp target
+#pragma omp teams distribute parallel for thread_limit(4)
+  for (int i = 0; i < N; ++i) {
+    optnone();
+  }
+// DEFAULT: 9 (MaxFlatWorkGroupSize: 9
+#pragma omp target teams distribute parallel for num_threads(9)
+  for (int i = 0; i < N; ++i) {
+    optnone();
+  }
+// DEFAULT: 4 (MaxFlatWorkGroupSize: 4
+#pragma omp target teams distribute parallel for simd thread_limit(4)
+  for (int i = 0; i < N; ++i) {
+    optnone();
+  }
+}
+