[Libomptarget] Fix block and thread limit environment variables not being respected

[jhuber6]

Summary

The next-gen plugins did not properly set the values from
OMP_NUM_TEAMS and OMP_TEAMS_THREAD_LIMIT. This is because these
maximum values are set by each plugin to its hardware maximum. This
happens *after* the previous initialization. Move it to the correct
place and then add a test.

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

openmp/libomptarget/plugins-nextgen/common/PluginInterface/PluginInterface.cpp

: MemoryManager(nullptr), OMP_TeamLimit("OMP_TEAM_LIMIT"),
       // device initialization. These cannot be consulted until the device is
       // initialized correctly. We intialize them in GenericDeviceTy::init().
       OMPX_TargetStackSize(), OMPX_TargetHeapSize(),
       // By default, the initial number of streams and events are 32.
       OMPX_InitialNumStreams("LIBOMPTARGET_NUM_INITIAL_STREAMS", 32),
       OMPX_InitialNumEvents("LIBOMPTARGET_NUM_INITIAL_EVENTS", 32),
       DeviceId(DeviceId), GridValues(OMPGridValues),
       PeerAccesses(NumDevices, PeerAccessState::PENDING), PeerAccessesLock(),
-      PinnedAllocs(*this) {
-  if (OMP_NumTeams > 0)
-    GridValues.GV_Max_Teams =
-        std::min(GridValues.GV_Max_Teams, uint32_t(OMP_NumTeams));
-
-  if (OMP_TeamsThreadLimit > 0)
-    GridValues.GV_Max_WG_Size =
-        std::min(GridValues.GV_Max_WG_Size, uint32_t(OMP_TeamsThreadLimit));
-}
+      PinnedAllocs(*this) {}
 
 Error GenericDeviceTy::init(GenericPluginTy &Plugin) {
   if (auto Err = initImpl(Plugin))
     return Err;
 
   // Read and reinitialize the envars that depend on the device initialization.
   // Notice these two envars may change the stack size and heap size of the
   // device, so they need the device properly initialized.
   auto StackSizeEnvarOrErr = UInt64Envar::create(
       "LIBOMPTARGET_STACK_SIZE",
       [this](uint64_t &V) -> Error { return getDeviceStackSize(V); },
       [this](uint64_t V) -> Error { return setDeviceStackSize(V); });
   if (!StackSizeEnvarOrErr)
     return StackSizeEnvarOrErr.takeError();
   OMPX_TargetStackSize = std::move(*StackSizeEnvarOrErr);
 
   auto HeapSizeEnvarOrErr = UInt64Envar::create(
       "LIBOMPTARGET_HEAP_SIZE",
       [this](uint64_t &V) -> Error { return getDeviceHeapSize(V); },
       [this](uint64_t V) -> Error { return setDeviceHeapSize(V); });
   if (!HeapSizeEnvarOrErr)
     return HeapSizeEnvarOrErr.takeError();
   OMPX_TargetHeapSize = std::move(*HeapSizeEnvarOrErr);
 
+  // Update the maximum number of teams and threads after the device
+  // initialization sets the corresponding hardware limit.
+  if (OMP_NumTeams > 0)
+    GridValues.GV_Max_Teams =
+        std::min(GridValues.GV_Max_Teams, uint32_t(OMP_NumTeams));
+
+  if (OMP_TeamsThreadLimit > 0)
+    GridValues.GV_Max_WG_Size =
+        std::min(GridValues.GV_Max_WG_Size, uint32_t(OMP_TeamsThreadLimit));
+
   // Enable the memory manager if required.
   auto [ThresholdMM, EnableMM] = MemoryManagerTy::getSizeThresholdFromEnv();
   if (EnableMM)
     MemoryManager = new MemoryManagerTy(*this, ThresholdMM);
 
   if (RecordReplay.isRecordingOrReplaying())
     if (auto Err = RecordReplay.init(this))
       return Err;
int32_t __tgt_rtl_is_data_exchangable(int32_t SrcDeviceId,
   return Plugin::get().isDataExchangable(SrcDeviceId, DstDeviceId);
 }
 
 __tgt_target_table *__tgt_rtl_load_binary(int32_t DeviceId,
                                           __tgt_device_image *TgtImage) {
   GenericPluginTy &Plugin = Plugin::get();
   GenericDeviceTy &Device = Plugin.getDevice(DeviceId);
 
-
   auto TableOrErr = Device.loadBinary(Plugin, TgtImage);
   if (!TableOrErr) {
     auto Err = TableOrErr.takeError();
     REPORT("Failure to load binary image %p on device %d: %s\n", TgtImage,
            DeviceId, toString(std::move(Err)).data());
     return nullptr;
   }
 

openmp/libomptarget/test/api/omp_env_vars.c

+// RUN: %libomptarget-compile-generic
+// RUN: env OMP_NUM_TEAMS=1 OMP_TEAMS_THREAD_LIMIT=1 LIBOMPTARGET_INFO=16 \
+// RUN:   %libomptarget-run-generic 2>&1 | %fcheck-generic
+
+#define N 256
+
+int main() {
+  // CHECK: Launching kernel [[KERNEL:.+_main_.+]] with 1 blocks and 1 threads
+#pragma omp target teams
+#pragma omp parallel
+  {}
+}