[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

625

exceeds the env team thread limit

629–634

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

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]);

 
 namespace core {
 hsa_status_t allow_access_to_all_gpu_agents(void *ptr) {
   return hsa_amd_agents_allow_access(DeviceInfo.HSAAgents.size(),
                                      &DeviceInfo.HSAAgents[0], NULL, ptr);
 }
 
 } // namespace core
 No newline at end of file

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

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 {
       DP("Using %d CUDA threads per block\n", MaxBlockDimX);
       DeviceData[DeviceId].ThreadsPerBlock = MaxBlockDimX;
-    } else {
+
+      if (EnvTeamThreadLimit > 0 &&
+          DeviceData[DeviceId].ThreadsPerBlock > EnvTeamThreadLimit) {
+        DP("Max CUDA threads per block %d exceeds the thread limit %d set by "
+           "OMP_TEAMS_THREAD_LIMIT, capping at the limit\n",
+           DeviceData[DeviceId].ThreadsPerBlock, EnvTeamThreadLimit);
+        DeviceData[DeviceId].ThreadsPerBlock = EnvTeamThreadLimit;
+      }
+      if (DeviceData[DeviceId].ThreadsPerBlock > DeviceRTLTy::HardThreadLimit) {
         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);
         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;
     } else {
       DP("Using warp size %d\n", WarpSize);
       DeviceData[DeviceId].WarpSize = WarpSize;
     }
 
     // Adjust teams to the env variables
     if (EnvTeamLimit > 0 && DeviceData[DeviceId].BlocksPerGrid > EnvTeamLimit) {
       DP("Capping max CUDA blocks per grid to OMP_TEAM_LIMIT=%d\n",
          EnvTeamLimit);
       DeviceData[DeviceId].BlocksPerGrid = EnvTeamLimit;
     }