[OpenMP][CUDA] Clean up the logic to initialize the limit and default number for teams and threads.

[tianshilei1992]

Summary

This patch makes some cleanup to the logic to determine the limit and
number of teams and threads to make it more clear.

[tianshilei1992]

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

1151

If both num_teams and OMP_NUM_TEAMS exist, which one to honor?

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

if (!StreamPool[DeviceId])
       StreamPool[DeviceId] = std::make_unique<StreamPoolTy>(
           StreamAllocatorTy(DeviceData[DeviceId].Context), NumInitialStreams);
 
     // Initialize the event pool.
     if (!EventPool[DeviceId])
       EventPool[DeviceId] = std::make_unique<EventPoolTy>(
           EventAllocatorTy(DeviceData[DeviceId].Context), NumInitialEvents);
 
-    // Query attributes to determine number of threads/block and blocks/grid.
+    // Set the limit for the number of teams.
     int MaxGridDimX;
     Err = cuDeviceGetAttribute(&MaxGridDimX, CU_DEVICE_ATTRIBUTE_MAX_GRID_DIM_X,
                                Device);
     if (Err != CUDA_SUCCESS) {
       DP("Error getting max grid dimension, use default value %d\n",
          DeviceRTLTy::DefaultNumTeams);
       DeviceData[DeviceId].BlocksPerGrid = DeviceRTLTy::DefaultNumTeams;
     } else {
       DP("Using %d CUDA blocks per grid\n", MaxGridDimX);
       DeviceData[DeviceId].BlocksPerGrid = MaxGridDimX;
     }
+    // Adjust the number to the env OMP_TEAM_LIMIT.
+    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;
+    }
 
-    // We are only exploiting threads along the x axis.
+    // Set the limit for the number of threads.
     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 {
       DP("Using %d CUDA threads per block\n", MaxBlockDimX);
       DeviceData[DeviceId].ThreadsPerBlock = MaxBlockDimX;
-
-      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",
-           DeviceData[DeviceId].ThreadsPerBlock, DeviceRTLTy::HardThreadLimit);
-        DeviceData[DeviceId].ThreadsPerBlock = DeviceRTLTy::HardThreadLimit;
-      }
+    }
+    // Adjust the number to the env OMP_TEAMS_THREAD_LIMIT.
+
+    // Set default number of teams used when user doesn't set in construct.
+    if (EnvNumTeams > 0) {
+      DP("Default number of teams set according to environment %d\n",
+         EnvNumTeams);
+      DeviceData[DeviceId].NumTeams = EnvNumTeams;
+    } else {
+      DP("Default number of teams set according to library's default %d\n",
+         DeviceRTLTy::DefaultNumTeams);
+      DeviceData[DeviceId].NumTeams = DeviceRTLTy::DefaultNumTeams;
+    if (DeviceData[DeviceId].NumTeams > DeviceData[DeviceId].BlocksPerGrid) {
+      DP("Default number of teams exceeds device limit, capping at %d\n",
+         DeviceData[DeviceId].BlocksPerGrid);
+      DeviceData[DeviceId].NumTeams = DeviceData[DeviceId].BlocksPerGrid;
+    }
+
+    // Set default number of threads used when user doesn't set in construct.
+    DeviceData[DeviceId].NumThreads = DeviceRTLTy::DefaultNumThreads;
+    DP("Default number of threads set according to library's default %d\n",
+       DeviceRTLTy::DefaultNumThreads);
+    if (DeviceData[DeviceId].NumThreads >
+        DeviceData[DeviceId].ThreadsPerBlock) {
+      DP("Default number of threads exceeds device limit, capping at %d\n",
+         DeviceData[DeviceId].ThreadsPerBlock);
+      DeviceData[DeviceId].NumThreads = DeviceData[DeviceId].ThreadsPerBlock;
     }
 
     // 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;
-    }
-
     size_t StackLimit;
     size_t HeapLimit;
     if (const char *EnvStr = getenv("LIBOMPTARGET_STACK_SIZE")) {
       StackLimit = std::stol(EnvStr);
       if (cuCtxSetLimit(CU_LIMIT_STACK_SIZE, StackLimit) != CUDA_SUCCESS)
         return OFFLOAD_FAIL;
     } else {
       if (cuCtxGetLimit(&StackLimit, CU_LIMIT_STACK_SIZE) != CUDA_SUCCESS)
INFO(OMP_INFOTYPE_PLUGIN_KERNEL, DeviceId,
          "warp size of %d\n",
          DeviceData[DeviceId].BlocksPerGrid,
          DeviceData[DeviceId].ThreadsPerBlock, DeviceData[DeviceId].WarpSize);
     INFO(OMP_INFOTYPE_PLUGIN_KERNEL, DeviceId,
          "Device heap size is %d Bytes, device stack size is %d Bytes per "
          "thread\n",
          (int)HeapLimit, (int)StackLimit);
 
-    // Set default number of teams
-    if (EnvNumTeams > 0) {
-      DP("Default number of teams set according to environment %d\n",
-         EnvNumTeams);
-      DeviceData[DeviceId].NumTeams = EnvNumTeams;
-    } else {
-      DeviceData[DeviceId].NumTeams = DeviceRTLTy::DefaultNumTeams;
-      DP("Default number of teams set according to library's default %d\n",
-         DeviceRTLTy::DefaultNumTeams);
-    }
-
-    if (DeviceData[DeviceId].NumTeams > DeviceData[DeviceId].BlocksPerGrid) {
-      DP("Default number of teams exceeds device limit, capping at %d\n",
-         DeviceData[DeviceId].BlocksPerGrid);
-      DeviceData[DeviceId].NumTeams = DeviceData[DeviceId].BlocksPerGrid;
-    }
-
-    // Set default number of threads
-    DeviceData[DeviceId].NumThreads = DeviceRTLTy::DefaultNumThreads;
-    DP("Default number of threads set according to library's default %d\n",
-       DeviceRTLTy::DefaultNumThreads);
-    if (DeviceData[DeviceId].NumThreads >
-        DeviceData[DeviceId].ThreadsPerBlock) {
-      DP("Default number of threads exceeds device limit, capping at %d\n",
-         DeviceData[DeviceId].ThreadsPerBlock);
-      DeviceData[DeviceId].NumThreads = DeviceData[DeviceId].ThreadsPerBlock;
-    }
-
     return OFFLOAD_SUCCESS;
   }
 
   int deinitDevice(const int DeviceId) {
     auto IsInitialized = InitializedFlags[DeviceId];
     if (!IsInitialized)
       return OFFLOAD_SUCCESS;
     InitializedFlags[DeviceId] = false;
if (ThreadLimit > 0) {
         DP("Adding master warp: +%d threads\n", DeviceData[DeviceId].WarpSize);
         CudaThreadsPerBlock += DeviceData[DeviceId].WarpSize;
       }
     } else {
       DP("Setting CUDA threads per block to default %d\n",
          DeviceData[DeviceId].NumThreads);
       CudaThreadsPerBlock = DeviceData[DeviceId].NumThreads;
     }
-
     if (CudaThreadsPerBlock > DeviceData[DeviceId].ThreadsPerBlock) {
       DP("Threads per block capped at device limit %d\n",
          DeviceData[DeviceId].ThreadsPerBlock);
       CudaThreadsPerBlock = DeviceData[DeviceId].ThreadsPerBlock;
     }
-
     if (!KernelInfo->MaxThreadsPerBlock) {
       Err = cuFuncGetAttribute(&KernelInfo->MaxThreadsPerBlock,
                                CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK,
                                KernelInfo->Func);
       if (!checkResult(Err, "Error returned from cuFuncGetAttribute\n"))
         return OFFLOAD_FAIL;
     }
-
     if (KernelInfo->MaxThreadsPerBlock < CudaThreadsPerBlock) {
       DP("Threads per block capped at kernel limit %d\n",
          KernelInfo->MaxThreadsPerBlock);
       CudaThreadsPerBlock = KernelInfo->MaxThreadsPerBlock;
     }
 
     unsigned int CudaBlocksPerGrid;
     if (TeamNum <= 0) {
       if (LoopTripCount > 0 && EnvNumTeams < 0) {
         if (IsSPMDGenericMode) {
           // If we reach this point, then we are executing a kernel that was
           // transformed from Generic-mode to SPMD-mode. This kernel has
           // SPMD-mode execution, but needs its blocks to be scheduled
           // differently because the current loop trip count only applies to the
           // `teams distribute` region and will create var too few blocks using
           // the regular SPMD-mode method.
if (TeamNum <= 0) {
       } else {
         DP("Using default number of teams %d\n", DeviceData[DeviceId].NumTeams);
         CudaBlocksPerGrid = DeviceData[DeviceId].NumTeams;
       }
     } else {
       DP("Using requested number of teams %d\n", TeamNum);
       CudaBlocksPerGrid = TeamNum;
     }
-
     if (CudaBlocksPerGrid > DeviceData[DeviceId].BlocksPerGrid) {
       DP("Capping number of teams to team limit %d\n",
          DeviceData[DeviceId].BlocksPerGrid);
       CudaBlocksPerGrid = DeviceData[DeviceId].BlocksPerGrid;
     }
 
     INFO(OMP_INFOTYPE_PLUGIN_KERNEL, DeviceId,
          "Launching kernel %s with %d blocks and %d threads in %s mode\n",