[libomptarget][cuda] Handle missing _v2 symbols gracefully

[JonChesterfield]

Summary

[libomptarget][cuda] Handle missing _v2 symbols gracefully

Follow on from D95367. Dlsym the _v2 symbols if present, otherwise use the
unsuffixed version. Builds a hashtable for the check, can revise for zero
heap allocations later if necessary.

[grokos]

openmp/libomptarget/plugins/cuda/dynamic_cuda/cuda.cpp

102–103

Maybe add a debug message here in order for the user to know which version was loaded eventually?

[jdoerfert]

LG, with the nit @grokos pointed out.

[JonChesterfield]

• Add debug print

[JonChesterfield]

openmp/libomptarget/plugins/cuda/dynamic_cuda/cuda.cpp

102–103

Good call! Thanks

openmp/libomptarget/plugins/cuda/dynamic_cuda/cuda.h

typedef enum CUfunction_attribute_enum {
   CU_FUNC_ATTRIBUTE_MAX_THREADS_PER_BLOCK = 0,
 } CUfunction_attribute;
 
 typedef enum CUctx_flags_enum {
   CU_CTX_SCHED_BLOCKING_SYNC = 0x04,
   CU_CTX_SCHED_MASK = 0x07,
 } CUctx_flags;
 
-#define cuMemFree cuMemFree_v2
-#define cuMemAlloc cuMemAlloc_v2
-#define cuMemcpyDtoH cuMemcpyDtoH_v2
-#define cuMemcpyHtoD cuMemcpyHtoD_v2
-#define cuStreamDestroy cuStreamDestroy_v2
-#define cuModuleGetGlobal cuModuleGetGlobal_v2
-#define cuMemcpyDtoHAsync cuMemcpyDtoHAsync_v2
-#define cuMemcpyDtoDAsync cuMemcpyDtoDAsync_v2
-#define cuMemcpyHtoDAsync cuMemcpyHtoDAsync_v2
-#define cuDevicePrimaryCtxRelease cuDevicePrimaryCtxRelease_v2
-#define cuDevicePrimaryCtxSetFlags cuDevicePrimaryCtxSetFlags_v2
-
 CUresult cuCtxGetDevice(CUdevice *);
 CUresult cuDeviceGet(CUdevice *, int);
 CUresult cuDeviceGetAttribute(int *, CUdevice_attribute, CUdevice);
 CUresult cuDeviceGetCount(int *);
 CUresult cuFuncGetAttribute(int *, CUfunction_attribute, CUfunction);
 
 CUresult cuGetErrorString(CUresult, const char **);
 CUresult cuInit(unsigned);

openmp/libomptarget/plugins/cuda/dynamic_cuda/cuda.cpp

 //===--- cuda/dynamic_cuda/cuda.pp ------------------------------- 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Implement subset of cuda api by calling into cuda library via dlopen
 // Does the dlopen/dlsym calls as part of the call to cuInit
 //
 //===----------------------------------------------------------------------===//
 
 #include "cuda.h"
 #include "Debug.h"
 #include "dlwrap.h"
 
+#include <string>
+#include <unordered_map>
+
 #include <dlfcn.h>
 
 DLWRAP_INTERNAL(cuInit, 1);
 
 DLWRAP(cuCtxGetDevice, 1);
 DLWRAP(cuDeviceGet, 2);
 DLWRAP(cuDeviceGetAttribute, 3);
 DLWRAP(cuDeviceGetCount, 1);

 #endif
 
 #define TARGET_NAME CUDA
 #define DEBUG_PREFIX "Target " GETNAME(TARGET_NAME) " RTL"
 
 static bool checkForCUDA() {
   // return true if dlopen succeeded and all functions found
 
+  // Prefer _v2 versions of functions if found in the library
+  std::unordered_map<std::string, const char *> TryFirst = {
+      {"cuMemAlloc", "cuMemAlloc_v2"},
+      {"cuMemFree", "cuMemFree_v2"},
+      {"cuMemcpyDtoH", "cuMemcpyDtoH_v2"},
+      {"cuMemcpyHtoD", "cuMemcpyHtoD_v2"},
+      {"cuStreamDestroy", "cuStreamDestroy_v2"},
+      {"cuModuleGetGlobal", "cuModuleGetGlobal_v2"},
+      {"cuMemcpyDtoHAsync", "cuMemcpyDtoHAsync_v2"},
+      {"cuMemcpyDtoDAsync", "cuMemcpyDtoDAsync_v2"},
+      {"cuMemcpyHtoDAsync", "cuMemcpyHtoDAsync_v2"},
+      {"cuDevicePrimaryCtxRelease", "cuDevicePrimaryCtxRelease_v2"},
+      {"cuDevicePrimaryCtxSetFlags", "cuDevicePrimaryCtxSetFlags_v2"},
+  };
+
   const char *CudaLib = DYNAMIC_CUDA_PATH;
   void *DynlibHandle = dlopen(CudaLib, RTLD_NOW);
   if (!DynlibHandle) {
     DP("Unable to load library '%s': %s!\n", CudaLib, dlerror());
     return false;
   }
 
   for (size_t I = 0; I < dlwrap::size(); I++) {
     const char *Sym = dlwrap::symbol(I);
 
+    auto It = TryFirst.find(Sym);
+    if (It != TryFirst.end()) {
+      const char *First = It->second;
+      void *P = dlsym(DynlibHandle, First);
+      if (P) {
+        DP("Implementing %s with dlsym(%s) -> %p\n", Sym, First, P);
+        *dlwrap::pointer(I) = P;
+        continue;
+      }
+    }
+
     void *P = dlsym(DynlibHandle, Sym);
     if (P == nullptr) {
       DP("Unable to find '%s' in '%s'!\n", Sym, CudaLib);
       return false;
     }
+    DP("Implementing %s with dlsym(%s) -> %p\n", Sym, Sym, P);
 
     *dlwrap::pointer(I) = P;
   }
 
   return true;
 }
 
 CUresult cuInit(unsigned X) {
   // Note: Called exactly once from cuda rtl.cpp in a global constructor so
   // does not need to handle being called repeatedly or concurrently
   if (!checkForCUDA()) {
     return CUDA_ERROR_INVALID_VALUE;
   }
   return dlwrap_cuInit(X);
 }