[Libomptarget] Deinitialize AMDGPU global state more intentionally

[jhuber6]

Summary

A previous patch made the destruction of the HSA plugin more
deterministic. However, there were still other global values that are not
handled this way. When attempting to call a destructor kernel, the
device would have already been uninitialized and we could not find the
appropriate kernel to call. This is because they were stored in global
containers that had their destructors called already. Merges this global
state into the rest of the info state by putting those global values
inside of the global pointer already allocated and deallocated by the
constructor and destructor. This should allow the AMDGPU plugin to
correctly identify the destructors if we were to run them.

[JonChesterfield]

Lots of noise in the diff. Could you do a precommit one where you move the struct / functions around in the file without changing them then rebase this patch?

[jhuber6]

Trying to clean up.

That being said, I've been doing some more thinking about how to handle this and I don't think this is the right way to go. We can set an explicit order on the destruction of the plugins, but that doesn't necessarily mean that the memory or libraries depending on that plugin will still be valid. I think this is the core issue behind why I was seeing weird behaviour porting CUDA to this method. For the CUDA RTL, the CUDA library may already be unloaded by the time we try to uninititalize the data and i don't think there's any way to guaruntee this one way or the other.

I'm thinking one reasonable way to solve this is to instead handle device destructors via atexit(), I believe that function is supposed to be called before standard user exit, e.g. after main() and exit(). That would allow us to run those functions before the destructors start unwinding presumably. Then to solve the problem that Jon's original patch aimed to solve, we simply make accessing internal plugin state illegal and explicitly dlclose the plugin at the start to prevent this. For this we would need to then remove the is_valid_binary_info calls, this should be easily replaced with some data structure that saves them instead of trying to derive them all over again. There may be some problems with this approach as well, but we really need an intelligent way to handle this stuff that won't blow up.

[JonChesterfield]

Strongly against going with atexit. The lifetime of these things needs to be sanely nested without use after free errors and we're only going to get that with explicit control flow.

I suspect the root cause of at least some problems is we don't bother to call dlclose and we should do so. I think we can guarantee relative lifetime of hsa/cuda vs the plugin. The edge between host executable and libomptarget is more alarming, but I expect the fix to that is to stop using global destructors in libomptarget and emit a call to tear it down from the main executable instead.

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

private:
   struct Lock {
     Lock(pthread_mutex_t *M) : M(M) { pthread_mutex_lock(M); }
     ~Lock() { pthread_mutex_unlock(M); }
     pthread_mutex_t *M;
   };
 };
 pthread_mutex_t KernelArgPool::Mutex = PTHREAD_MUTEX_INITIALIZER;
 
-std::unordered_map<std::string /*kernel*/, std::unique_ptr<KernelArgPool>>
-    KernelArgPoolMap;
-
 /// Use a single entity to encode a kernel and a set of flags
 struct KernelTy {
   llvm::omp::OMPTgtExecModeFlags ExecutionMode;
   int16_t ConstWGSize;
   int32_t DeviceId;
   void *CallStackAddr = nullptr;
   const char *Name;
 
   KernelTy(llvm::omp::OMPTgtExecModeFlags ExecutionMode, int16_t ConstWgSize,
            int32_t DeviceId, void *CallStackAddr, const char *Name,
            uint32_t KernargSegmentSize,
-           hsa_amd_memory_pool_t &KernArgMemoryPool)
+           hsa_amd_memory_pool_t &KernArgMemoryPool,
+           std::unordered_map<std::string, std::unique_ptr<KernelArgPool>>
+               &KernelArgPoolMap)
       : ExecutionMode(ExecutionMode), ConstWGSize(ConstWgSize),
         DeviceId(DeviceId), CallStackAddr(CallStackAddr), Name(Name) {
     DP("Construct kernelinfo: ExecMode %d\n", ExecutionMode);
 
     std::string N(Name);
     if (KernelArgPoolMap.find(N) == KernelArgPoolMap.end()) {
       KernelArgPoolMap.insert(
           std::make_pair(N, std::unique_ptr<KernelArgPool>(new KernelArgPool(
                                 KernargSegmentSize, KernArgMemoryPool))));
     }
   }
 };
 
-/// List that contains all the kernels.
-/// FIXME: we may need this to be per device and per library.
-std::list<KernelTy> KernelsList;
-
 template <typename Callback> static hsa_status_t findAgents(Callback CB) {
 
   hsa_status_t Err =
       hsa::iterate_agents([&](hsa_agent_t Agent) -> hsa_status_t {
         hsa_device_type_t DeviceType;
         // get_info fails iff HSA runtime not yet initialized
         hsa_status_t Err =
             hsa_agent_get_info(Agent, HSA_AGENT_INFO_DEVICE, &DeviceType);
public:
   bool ConstructionSucceeded = false;
 
   // load binary populates symbol tables and mutates various global state
   // run uses those symbol tables
   std::shared_timed_mutex LoadRunLock;
 
   int NumberOfDevices = 0;
 
+  /// List that contains all the kernels.
+  /// FIXME: we may need this to be per device and per library.
+  std::list<KernelTy> KernelsList;
+  std::unordered_map<std::string /*kernel*/, std::unique_ptr<KernelArgPool>>
+      KernelArgPoolMap;
+
   // GPU devices
   std::vector<hsa_agent_t> HSAAgents;
   std::vector<HSAQueueScheduler> HSAQueueSchedulers; // one per gpu
 
   // CPUs
   std::vector<hsa_agent_t> CPUAgents;
 
   // Device properties
public:
   }
 
   // Clear entries table for a device
   void clearOffloadEntriesTable(int DeviceId) {
     assert(DeviceId < (int32_t)FuncGblEntries.size() &&
            "Unexpected device id!");
     FuncGblEntries[DeviceId].emplace_back();
     FuncOrGblEntryTy &E = FuncGblEntries[DeviceId].back();
-    // KernelArgPoolMap.clear();
     E.Entries.clear();
     E.Table.EntriesBegin = E.Table.EntriesEnd = 0;
   }
 
   hsa_status_t addDeviceMemoryPool(hsa_amd_memory_pool_t MemoryPool,
                                    unsigned int DeviceId) {
     assert(DeviceId < DeviceFineGrainedMemoryPools.size() && "Error here.");
     uint32_t GlobalFlags = 0;
public:
   }
 };
 
 pthread_mutex_t SignalPoolT::mutex = PTHREAD_MUTEX_INITIALIZER;
 
 static RTLDeviceInfoTy *DeviceInfoState = nullptr;
 static RTLDeviceInfoTy &DeviceInfo() { return *DeviceInfoState; }
 
-int32_t __tgt_rtl_init_plugin() {
-  DeviceInfoState = new RTLDeviceInfoTy;
-  return (DeviceInfoState && DeviceInfoState->ConstructionSucceeded)
-             ? OFFLOAD_SUCCESS
-             : OFFLOAD_FAIL;
-}
-
-int32_t __tgt_rtl_deinit_plugin() {
-  if (DeviceInfoState)
-    delete DeviceInfoState;
-  return OFFLOAD_SUCCESS;
-}
-
 namespace {
 
 int32_t dataRetrieve(int32_t DeviceId, void *HstPtr, void *TgtPtr, int64_t Size,
                      __tgt_async_info *AsyncInfo) {
   assert(AsyncInfo && "AsyncInfo is nullptr");
   assert(DeviceId < DeviceInfo().NumberOfDevices && "Device ID too large");
   // Return success if we are not copying back to host from target.
   if (!HstPtr)
int32_t dataSubmit(int32_t DeviceId, void *TgtPtr, void *HstPtr, int64_t Size,
   // Return success if we are not doing host to target.
   if (!HstPtr)
     return OFFLOAD_SUCCESS;
 
   DP("Submit data %ld bytes, (hst:%016llx) -> (tgt:%016llx).\n", Size,
      (long long unsigned)(Elf64_Addr)HstPtr,
      (long long unsigned)(Elf64_Addr)TgtPtr);
   Err = DeviceInfo().freesignalpoolMemcpyH2D(TgtPtr, HstPtr, (size_t)Size,
                                              DeviceId);
   if (Err != HSA_STATUS_SUCCESS) {
     DP("Error when copying data from host to device. Pointers: "
        "host = 0x%016lx, device = 0x%016lx, size = %lld\n",
        (Elf64_Addr)HstPtr, (Elf64_Addr)TgtPtr, (unsigned long long)Size);
     return OFFLOAD_FAIL;
   }
   return OFFLOAD_SUCCESS;
 }
// Run on the device.
     Packet->kernel_object = KernelInfoEntry.kernel_object;
     Packet->kernarg_address = 0;     // use the block allocator
     Packet->reserved2 = 0;           // impl writes id_ here
     Packet->completion_signal = {0}; // may want a pool of signals
 
     KernelArgPool *ArgPool = nullptr;
     void *KernArg = nullptr;
     {
-      auto It = KernelArgPoolMap.find(std::string(KernelInfo->Name));
-      if (It != KernelArgPoolMap.end()) {
+      auto It =
+          DeviceInfo().KernelArgPoolMap.find(std::string(KernelInfo->Name));
+      if (It != DeviceInfo().KernelArgPoolMap.end()) {
         ArgPool = (It->second).get();
       }
     }
     if (!ArgPool) {
       DP("Warning: No ArgPool for %s on device %d\n", KernelInfo->Name,
          DeviceId);
     }
     {
bool IsImageCompatibleWithEnv(const char *ImgInfo, std::string EnvInfo) {
   //   set to ANY
   DP("Compatible: Target IDs are compatible \t[Image: %s]\t:\t[Environment: "
      "%s]\n",
      ImgTID.data(), EnvTID.data());
   return true;
 }
 
 extern "C" {
+
+int32_t __tgt_rtl_init_plugin() {
+  DeviceInfoState = new RTLDeviceInfoTy;
+  return (DeviceInfoState && DeviceInfoState->ConstructionSucceeded)
+             ? OFFLOAD_SUCCESS
+             : OFFLOAD_FAIL;
+}
+
+int32_t __tgt_rtl_deinit_plugin() {
+  if (DeviceInfoState)
+    delete DeviceInfoState;
+  return OFFLOAD_SUCCESS;
+}
+
 int32_t __tgt_rtl_is_valid_binary(__tgt_device_image *Image) {
   return elfMachineIdIsAmdgcn(Image);
 }
 
 int32_t __tgt_rtl_is_valid_binary_info(__tgt_device_image *image,
                                        __tgt_image_info *info) {
   if (!__tgt_rtl_is_valid_binary(image))
     return false;
if (Err != HSA_STATUS_SUCCESS) {
 
         void *Ptr = DSS.first.get();
         if (DeviceStateBytes != DSS.second) {
           DP("Inconsistent sizes of device_State unsupported\n");
           return NULL;
         }
 
         // write ptr to device memory so it can be used by later kernels
-        Err = DeviceInfo().freesignalpoolMemcpyH2D(StatePtr, &Ptr, sizeof(void *),
-                                                 DeviceId);
+        Err = DeviceInfo().freesignalpoolMemcpyH2D(StatePtr, &Ptr,
+                                                   sizeof(void *), DeviceId);
         if (Err != HSA_STATUS_SUCCESS) {
           DP("memcpy install of state_ptr failed\n");
           return NULL;
         }
       }
     }
   }
 
if (E->size) {
 
       DeviceInfo().addOffloadEntry(DeviceId, Entry);
 
       if (DeviceInfo().RequiresFlags & OMP_REQ_UNIFIED_SHARED_MEMORY &&
           E->flags & OMP_DECLARE_TARGET_LINK) {
         // If unified memory is present any target link variables
         // can access host addresses directly. There is no longer a
         // need for device copies.
-        Err = DeviceInfo().freesignalpoolMemcpyH2D(Varptr, E->addr,
-                                                 sizeof(void *), DeviceId);
+        Err = DeviceInfo().freesignalpoolMemcpyH2D(Varptr, E->addr, sizeof(void *),
+                                                 DeviceId);
         if (Err != HSA_STATUS_SUCCESS)
           DP("Error when copying USM\n");
         DP("Copy linked variable host address (" DPxMOD ")"
            "to device address (" DPxMOD ")\n",
            DPxPTR(*((void **)E->addr)), DPxPTR(Varptr));
       }
 
       continue;
for (__tgt_offload_entry *E = HostBegin; E != HostEnd; ++E) {
     } else {
       DP("Loading global exec_mode '%s' - symbol missing, using default "
          "value "
          "GENERIC (1)\n",
          ExecModeName);
     }
     check("Loading computation property", Err);
 
-    KernelsList.push_back(KernelTy(ExecModeVal, WGSizeVal, DeviceId,
-                                   CallStackAddr, E->name, KernargSegmentSize,
-                                   DeviceInfo().KernArgPool));
+    DeviceInfo().KernelsList.push_back(
+        KernelTy(ExecModeVal, WGSizeVal, DeviceId, CallStackAddr, E->name,
+                 KernargSegmentSize, DeviceInfo().KernArgPool,
+                 DeviceInfo().KernelArgPoolMap));
     __tgt_offload_entry Entry = *E;
-    Entry.addr = (void *)&KernelsList.back();
+    Entry.addr = (void *)&DeviceInfo().KernelsList.back();
     DeviceInfo().addOffloadEntry(DeviceId, Entry);
     DP("Entry point %ld maps to %s\n", E - HostBegin, E->name);
   }
 
   return DeviceInfo().getOffloadEntriesTable(DeviceId);
 }
 
 void *__tgt_rtl_data_alloc(int DeviceId, int64_t Size, void *, int32_t Kind) {