[Openmp-commits] [PATCH] D95415: [libomptarget][cuda] Handle missing _v2 symbols gracefully

[Jon Chesterfield via Phabricator via Openmp-commits]

JonChesterfield created this revision.
JonChesterfield added reviewers: jdoerfert, tianshilei1992, kkwli0.
Herald added a subscriber: yaxunl.
JonChesterfield requested review of this revision.
Herald added a project: OpenMP.
Herald added a subscriber: openmp-commits.

[libomptarget][cuda] Handle missing _v2 symbols gracefully

Follow on from D95367 <https://reviews.llvm.org/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.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D95415

Files:
  openmp/libomptarget/plugins/cuda/dynamic_cuda/cuda.cpp
  openmp/libomptarget/plugins/cuda/dynamic_cuda/cuda.h


Index: openmp/libomptarget/plugins/cuda/dynamic_cuda/cuda.h
===================================================================
--- openmp/libomptarget/plugins/cuda/dynamic_cuda/cuda.h
+++ openmp/libomptarget/plugins/cuda/dynamic_cuda/cuda.h
@@ -49,18 +49,6 @@
   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);
Index: openmp/libomptarget/plugins/cuda/dynamic_cuda/cuda.cpp
===================================================================
--- openmp/libomptarget/plugins/cuda/dynamic_cuda/cuda.cpp
+++ openmp/libomptarget/plugins/cuda/dynamic_cuda/cuda.cpp
@@ -15,6 +15,9 @@
 #include "Debug.h"
 #include "dlwrap.h"
 
+#include <string>
+#include <unordered_map>
+
 #include <dlfcn.h>
 
 DLWRAP_INTERNAL(cuInit, 1);
@@ -67,6 +70,21 @@
 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) {
@@ -77,6 +95,15 @@
   for (size_t I = 0; I < dlwrap::size(); I++) {
     const char *Sym = dlwrap::symbol(I);
 
+    auto It = TryFirst.find(Sym);
+    if (It != TryFirst.end()) {
+      void *P = dlsym(DynlibHandle, It->second);
+      if (P) {
+        *dlwrap::pointer(I) = P;
+        continue;
+      }
+    }
+
     void *P = dlsym(DynlibHandle, Sym);
     if (P == nullptr) {
       DP("Unable to find '%s' in '%s'!\n", Sym, CudaLib);


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D95415.319190.patch
Type: text/x-patch
Size: 2753 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/openmp-commits/attachments/20210126/29243612/attachment.bin>