[OpenMP][CUDA] Cache the maximal number of threads per block (per kernel)

[jdoerfert]

Summary

Instead of calling cuFuncGetAttribute with
CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK for every kernel invocation,
we can do it for the first one and cache the result as part of the
KernelInfo struct. The only functional change is that we now expect
cuFuncGetAttribute to succeed and otherwise propagate the error.
Ignoring any error seems like a slippery slope...

[JonChesterfield]

LGTM

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

892

tab

[tianshilei1992]

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

79

Do we need to change the front end? This data structure should be generated by the FE, right?

[tianshilei1992]

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

79

Oh, my bad. It is not. The pointer is replaced during the initialization.

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 <vector>
 
 #include "omptargetplugin.h"
 

 struct KernelTy {
   CUfunction Func;
 
   // execution mode of kernel
   // 0 - SPMD mode (without master warp)
   // 1 - Generic mode (with master warp)
   int8_t ExecutionMode;
 
+  /// Maximal number of threads per block for this kernel.
+  int MaxThreadsPerBlock = 0;
+
   KernelTy(CUfunction _Func, int8_t _ExecutionMode)
       : Func(_Func), ExecutionMode(_ExecutionMode) {}
 };
 
 /// Device environment data
 /// Manually sync with the deviceRTL side for now, move to a dedicated header
 /// file later.
 struct omptarget_device_environmentTy {
int dataDelete(const int DeviceId, void *TgtPtr) const {
 
     Err = cuMemFree((CUdeviceptr)TgtPtr);
     if (!checkResult(Err, "Error returned from cuMemFree\n"))
       return OFFLOAD_FAIL;
 
     return OFFLOAD_SUCCESS;
   }
 
-  int runTargetTeamRegion(const int DeviceId, const void *TgtEntryPtr,
-                          void **TgtArgs, ptrdiff_t *TgtOffsets,
-                          const int ArgNum, const int TeamNum,
-                          const int ThreadLimit,
+  int runTargetTeamRegion(const int DeviceId, void *TgtEntryPtr, void **TgtArgs,
+                          ptrdiff_t *TgtOffsets, const int ArgNum,
+                          const int TeamNum, const int ThreadLimit,
                           const unsigned int LoopTripCount,
                           __tgt_async_info *AsyncInfo) const {
     CUresult Err = cuCtxSetCurrent(DeviceData[DeviceId].Context);
     if (!checkResult(Err, "Error returned from cuCtxSetCurrent\n"))
       return OFFLOAD_FAIL;
 
     // All args are references.
     std::vector<void *> Args(ArgNum);
     std::vector<void *> Ptrs(ArgNum);
 
     for (int I = 0; I < ArgNum; ++I) {
       Ptrs[I] = (void *)((intptr_t)TgtArgs[I] + TgtOffsets[I]);
       Args[I] = &Ptrs[I];
     }
 
-    const KernelTy *KernelInfo =
-        reinterpret_cast<const KernelTy *>(TgtEntryPtr);
+    KernelTy *KernelInfo = reinterpret_cast<KernelTy *>(TgtEntryPtr);
 
-    unsigned int CudaThreadsPerBlock;
+    int CudaThreadsPerBlock;
     if (ThreadLimit > 0) {
       DP("Setting CUDA threads per block to requested %d\n", ThreadLimit);
       CudaThreadsPerBlock = ThreadLimit;
       // Add master warp if necessary
       if (KernelInfo->ExecutionMode == GENERIC) {
         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;
     }
 
-    int KernelLimit;
-    Err = cuFuncGetAttribute(&KernelLimit,
+    if (!KernelInfo->MaxThreadsPerBlock) {
+      Err = cuFuncGetAttribute(&KernelInfo->MaxThreadsPerBlock,
                                CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK,
                                KernelInfo->Func);
-    if (Err == CUDA_SUCCESS && KernelLimit < CudaThreadsPerBlock) {
-      DP("Threads per block capped at kernel limit %d\n", KernelLimit);
-      CudaThreadsPerBlock = KernelLimit;
+      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 (KernelInfo->ExecutionMode == SPMD) {
           // We have a combined construct, i.e. `target teams distribute
           // parallel for [simd]`. We launch so many teams so that each thread