[Openmp-commits] [PATCH] D121308: [OpenMP][CUDA] Fix an issue that multiple `CUmodule` are could be overwritten

[Shilei Tian via Phabricator via Openmp-commits]

tianshilei1992 created this revision.
tianshilei1992 added reviewers: jdoerfert, JonChesterfield, jhuber6.
Herald added subscribers: guansong, yaxunl.
Herald added a project: All.
tianshilei1992 requested review of this revision.
Herald added subscribers: openmp-commits, sstefan1.
Herald added a project: OpenMP.

This patch fixes the issue introduced in 14de0820e87f <https://reviews.llvm.org/rG14de0820e87f552993cf9b2c5c31f67e12fb0644> and D120089 <https://reviews.llvm.org/D120089>, that
if dynamic libraries are used, the `CUmodule` array could be overwritten.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D121308

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


Index: openmp/libomptarget/plugins/cuda/src/rtl.cpp
===================================================================
--- openmp/libomptarget/plugins/cuda/src/rtl.cpp
+++ openmp/libomptarget/plugins/cuda/src/rtl.cpp
@@ -354,7 +354,7 @@
   std::vector<std::unique_ptr<EventPoolTy>> EventPool;
 
   std::vector<DeviceDataTy> DeviceData;
-  std::vector<CUmodule> Modules;
+  std::vector<std::vector<CUmodule>> Modules;
 
   /// Vector of flags indicating the initalization status of all associated
   /// devices.
@@ -777,25 +777,30 @@
     if (UseMemoryManager)
       MemoryManagers[DeviceId].release();
 
-    // Close module
-    if (CUmodule &M = Modules[DeviceId])
-      checkResult(cuModuleUnload(M), "Error returned from cuModuleUnload\n");
-
     StreamPool[DeviceId].reset();
     EventPool[DeviceId].reset();
 
-    // Destroy context
     DeviceDataTy &D = DeviceData[DeviceId];
-    if (D.Context) {
-      if (checkResult(cuCtxSetCurrent(D.Context),
-                      "Error returned from cuCtxSetCurrent\n")) {
-        CUdevice Device;
-        if (checkResult(cuCtxGetDevice(&Device),
-                        "Error returned from cuCtxGetDevice\n"))
-          checkResult(cuDevicePrimaryCtxRelease(Device),
-                      "Error returned from cuDevicePrimaryCtxRelease\n");
-      }
-    }
+    if (!checkResult(cuCtxSetCurrent(D.Context),
+                     "Error returned from cuCtxSetCurrent\n"))
+      return OFFLOAD_FAIL;
+
+    // Unload all modules.
+    for (auto &M : Modules[DeviceId])
+      if (!checkResult(cuModuleUnload(M),
+                       "Error returned from cuModuleUnload\n"))
+        return OFFLOAD_FAIL;
+
+    // Destroy context.
+    CUdevice Device;
+    if (!checkResult(cuCtxGetDevice(&Device),
+                     "Error returned from cuCtxGetDevice\n"))
+      return OFFLOAD_FAIL;
+
+    if (!checkResult(cuDevicePrimaryCtxRelease(Device),
+                     "Error returned from cuDevicePrimaryCtxRelease\n"))
+      return OFFLOAD_FAIL;
+
     return OFFLOAD_SUCCESS;
   }
 
@@ -818,7 +823,7 @@
 
     DP("CUDA module successfully loaded!\n");
 
-    Modules[DeviceId] = Module;
+    Modules[DeviceId].push_back(Module);
 
     // Find the symbols in the module by name.
     const __tgt_offload_entry *HostBegin = Image->EntriesBegin;


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D121308.414148.patch
Type: text/x-patch
Size: 2318 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/openmp-commits/attachments/20220309/9c5c3bc4/attachment-0001.bin>