[OpenMP] Add thread limit environment variable support to plugins

[jhuber6]

Summary

The OpenMP 5.1 standard defines the environment variable
OMP_TEAMS_THREAD_LIMIT to limit the number of threads that will be run in a
single block. This patch adds support for this into the AMDGPU and CUDA
plugins.

[grokos]

openmp/libomptarget/plugins/cuda/src/rtl.cpp

623

exceeds the env team thread limit

627–632

This else does not take care of this scenario correctly:

EnvTeamThreadLimit < HardThreadLimit < MaxBlockDimX

MaxBlockDimX will be set to HardThreadLimit whereas it should be set to EnvTeamThreadLimit.

[jhuber6]

Fixing. Don't know why I kept that mess.

[jhuber6]

Ping.

[jdoerfert]

LGTM

openmp/libomptarget/plugins/amdgpu/src/rtl.cpp

 #include <cstdio>
 #include <cstdlib>
 #include <cstring>
 #include <dlfcn.h>
 #include <elf.h>
 #include <ffi.h>
 #include <fstream>
 #include <iostream>
 #include <libelf.h>
 #include <list>
 #include <memory>
 #include <mutex>
 #include <shared_mutex>
 #include <thread>
 #include <unordered_map>
 #include <vector>
 
public:
 
   // OpenMP properties
   std::vector<int> NumTeams;
   std::vector<int> NumThreads;
 
   // OpenMP Environment properties
   int EnvNumTeams;
   int EnvTeamLimit;
+  int EnvTeamThreadLimit;
   int EnvMaxTeamsDefault;
 
   // OpenMP Requires Flags
   int64_t RequiresFlags;
 
   // Resource pools
   SignalPoolT FreeSignalPool;
 
RTLDeviceInfoTy() {
     // Get environment variables regarding expMaxTeams
     envStr = getenv("OMP_MAX_TEAMS_DEFAULT");
     if (envStr) {
       EnvMaxTeamsDefault = std::stoi(envStr);
       DP("Parsed OMP_MAX_TEAMS_DEFAULT=%d\n", EnvMaxTeamsDefault);
     } else {
       EnvMaxTeamsDefault = -1;
     }
+    envStr = getenv("OMP_TEAMS_THREAD_LIMIT");
+    if (envStr) {
+      EnvTeamThreadLimit = std::stoi(envStr);
+      DP("Parsed OMP_TEAMS_THREAD_LIMIT=%d\n", EnvTeamThreadLimit);
+    } else {
+      EnvTeamThreadLimit = -1;
+    }
 
     // Default state.
     RequiresFlags = OMP_REQ_UNDEFINED;
   }
 
   ~RTLDeviceInfoTy() {
     DP("Finalizing the HSA-ATMI DeviceInfo.\n");
     // Run destructors on types that use HSA before
int32_t __tgt_rtl_init_device(int device_id) {
   }
 
   if (enforce_upper_bound(&DeviceInfo.NumTeams[device_id],
                           DeviceInfo.GroupsPerDevice[device_id])) {
     DP("Default number of teams exceeds device limit, capping at %d\n",
        DeviceInfo.GroupsPerDevice[device_id]);
   }
 
+  // Adjust threads to the env variables
+  if (DeviceInfo.EnvTeamThreadLimit > 0 &&
+      (enforce_upper_bound(&DeviceInfo.NumThreads[device_id],
+                           DeviceInfo.EnvTeamThreadLimit))) {
+    DP("Capping max number of threads to OMP_TEAMS_THREAD_LIMIT=%d\n",
+       DeviceInfo.EnvTeamThreadLimit);
+  }
+
   // Set default number of threads
   DeviceInfo.NumThreads[device_id] = RTLDeviceInfoTy::Default_WG_Size;
   DP("Default number of threads set according to library's default %d\n",
      RTLDeviceInfoTy::Default_WG_Size);
   if (enforce_upper_bound(&DeviceInfo.NumThreads[device_id],
                           DeviceInfo.ThreadsPerGroup[device_id])) {
     DP("Default number of threads exceeds device limit, capping at %d\n",
        DeviceInfo.ThreadsPerGroup[device_id]);

openmp/libomptarget/plugins/cuda/src/rtl.cpp

 //===----RTLs/cuda/src/rtl.cpp - Target RTLs Implementation ------- 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
 //
 //===----------------------------------------------------------------------===//
 //
 // RTL for CUDA machine
 //
 //===----------------------------------------------------------------------===//
 
 #include <cassert>
 #include <cstddef>
 #include <cuda.h>
 #include <list>
 #include <memory>
 #include <mutex>
 #include <string>
 #include <unordered_map>
 #include <vector>
 
 #include "Debug.h"
public:
   }
 };
 
 class DeviceRTLTy {
   int NumberOfDevices;
   // OpenMP environment properties
   int EnvNumTeams;
   int EnvTeamLimit;
+  int EnvTeamThreadLimit;
   // OpenMP requires flags
   int64_t RequiresFlags;
 
   static constexpr const int HardTeamLimit = 1U << 16U; // 64k
   static constexpr const int HardThreadLimit = 1024;
   static constexpr const int DefaultNumTeams = 128;
   static constexpr const int DefaultNumThreads = 128;
 

 
 public:
   // This class should not be copied
   DeviceRTLTy(const DeviceRTLTy &) = delete;
   DeviceRTLTy(DeviceRTLTy &&) = delete;
 
   DeviceRTLTy()
       : NumberOfDevices(0), EnvNumTeams(-1), EnvTeamLimit(-1),
-        RequiresFlags(OMP_REQ_UNDEFINED) {
+        EnvTeamThreadLimit(-1), RequiresFlags(OMP_REQ_UNDEFINED) {
 
     DP("Start initializing CUDA\n");
 
     CUresult Err = cuInit(0);
     if (Err == CUDA_ERROR_INVALID_HANDLE) {
       // Can't call cuGetErrorString if dlsym failed
       DP("Failed to load CUDA shared library\n");
       return;
DeviceRTLTy()
     DeviceData.resize(NumberOfDevices);
 
     // Get environment variables regarding teams
     if (const char *EnvStr = getenv("OMP_TEAM_LIMIT")) {
       // OMP_TEAM_LIMIT has been set
       EnvTeamLimit = std::stoi(EnvStr);
       DP("Parsed OMP_TEAM_LIMIT=%d\n", EnvTeamLimit);
     }
+    if (const char *EnvStr = getenv("OMP_TEAMS_THREAD_LIMIT")) {
+      // OMP_TEAMS_THREAD_LIMIT has been set
+      EnvTeamThreadLimit = std::stoi(EnvStr);
+      DP("Parsed OMP_TEAMS_THREAD_LIMIT=%d\n", EnvTeamThreadLimit);
+    }
     if (const char *EnvStr = getenv("OMP_NUM_TEAMS")) {
       // OMP_NUM_TEAMS has been set
       EnvNumTeams = std::stoi(EnvStr);
       DP("Parsed OMP_NUM_TEAMS=%d\n", EnvNumTeams);
     }
 
     StreamManager =
         std::make_unique<StreamManagerTy>(NumberOfDevices, DeviceData);
int initDevice(const int DeviceId) {
     // We are only exploiting threads along the x axis.
     int MaxBlockDimX;
     Err = cuDeviceGetAttribute(&MaxBlockDimX,
                                CU_DEVICE_ATTRIBUTE_MAX_BLOCK_DIM_X, Device);
     if (Err != CUDA_SUCCESS) {
       DP("Error getting max block dimension, use default value %d\n",
          DeviceRTLTy::DefaultNumThreads);
       DeviceData[DeviceId].ThreadsPerBlock = DeviceRTLTy::DefaultNumThreads;
-    } else if (MaxBlockDimX <= DeviceRTLTy::HardThreadLimit) {
+    } else {
+      if (EnvTeamThreadLimit < 0) {
+        if (MaxBlockDimX <= DeviceRTLTy::HardThreadLimit) {
           DP("Using %d CUDA threads per block\n", MaxBlockDimX);
           DeviceData[DeviceId].ThreadsPerBlock = MaxBlockDimX;
         } else {
           DP("Max CUDA threads per block %d exceeds the hard thread limit %d, "
              "capping at the hard limit\n",
              MaxBlockDimX, DeviceRTLTy::HardThreadLimit);
           DeviceData[DeviceId].ThreadsPerBlock = DeviceRTLTy::HardThreadLimit;
         }
+      } else {
+        if (MaxBlockDimX <= DeviceRTLTy::HardThreadLimit &&
+            MaxBlockDimX <= EnvTeamThreadLimit) {
+          DP("Using %d CUDA threads per block\n", MaxBlockDimX);
+          DeviceData[DeviceId].ThreadsPerBlock = MaxBlockDimX;
+        } else if (MaxBlockDimX <= DeviceRTLTy::HardThreadLimit &&
+                   MaxBlockDimX > EnvTeamThreadLimit) {
+          DP("Max CUDA threads per block %d exceeds the hard thread limit %d "
+             "set by OMP_TEAMS_THREAD_LIMIT, capping at the limit\n",
+             MaxBlockDimX, EnvTeamThreadLimit);
+          DeviceData[DeviceId].ThreadsPerBlock = EnvTeamThreadLimit;
+        } else {
+          DP("Max CUDA threads per block %d exceeds the hard thread limit %d, "
+             "capping at the hard limit\n",
+             MaxBlockDimX, DeviceRTLTy::HardThreadLimit);
+          DeviceData[DeviceId].ThreadsPerBlock = DeviceRTLTy::HardThreadLimit;
+        }
+      }
+    }
 
     // Get and set warp size
     int WarpSize;
     Err =
         cuDeviceGetAttribute(&WarpSize, CU_DEVICE_ATTRIBUTE_WARP_SIZE, Device);
     if (Err != CUDA_SUCCESS) {
       DP("Error getting warp size, assume default value 32\n");
       DeviceData[DeviceId].WarpSize = 32;