From ed0f21811544320f829124efbb6a38ee12eb9155 Mon Sep 17 00:00:00 2001
From: Jon Chesterfield <jonathanchesterfield@gmail.com>
Date: Thu, 28 Jul 2022 20:00:01 +0100
Subject: [PATCH] [openmp][amdgpu] Tear down amdgpu plugin accurately

Moves DeviceInfo global to heap to accurately control lifetime.
Moves calls from libomptarget to deinit_plugin later, plugins need to stay
alive until very shortly before libomptarget is destructed.

Leaving the deinit_plugin calls where initially inserted hits use after
free from the dynamic_module.c offloading test (verified with valgrind
 that the new location is sound with respect to this)

Reviewed By: tianshilei1992

Differential Revision: https://reviews.llvm.org/D130714
---
 .../libomptarget/plugins/amdgpu/src/rtl.cpp   | 22 +++++++++++------
 openmp/libomptarget/src/rtl.cpp               | 24 +++++++++++++------
 2 files changed, 32 insertions(+), 14 deletions(-)

diff --git a/openmp/libomptarget/plugins/amdgpu/src/rtl.cpp b/openmp/libomptarget/plugins/amdgpu/src/rtl.cpp
index 3da571f92766..3eb1d0ffa630 100644
--- a/openmp/libomptarget/plugins/amdgpu/src/rtl.cpp
+++ b/openmp/libomptarget/plugins/amdgpu/src/rtl.cpp
@@ -1113,10 +1113,21 @@ public:
 
 pthread_mutex_t SignalPoolT::mutex = PTHREAD_MUTEX_INITIALIZER;
 
-// Putting accesses to DeviceInfo global behind a function call prior
-// to changing to use init_plugin/deinit_plugin calls
-static RTLDeviceInfoTy DeviceInfoState;
-static RTLDeviceInfoTy& DeviceInfo() { return DeviceInfoState; }
+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 {
 
@@ -2051,9 +2062,6 @@ int32_t __tgt_rtl_is_valid_binary_info(__tgt_device_image *image,
   return true;
 }
 
-int32_t __tgt_rtl_init_plugin() { return OFFLOAD_SUCCESS; }
-int32_t __tgt_rtl_deinit_plugin() { return OFFLOAD_SUCCESS; }
-
 int __tgt_rtl_number_of_devices() {
   // If the construction failed, no methods are safe to call
   if (DeviceInfo().ConstructionSucceeded) {
diff --git a/openmp/libomptarget/src/rtl.cpp b/openmp/libomptarget/src/rtl.cpp
index 9acf5f3e19c3..aef6fc6ec82e 100644
--- a/openmp/libomptarget/src/rtl.cpp
+++ b/openmp/libomptarget/src/rtl.cpp
@@ -67,6 +67,23 @@ __attribute__((constructor(101))) void init() {
 
 __attribute__((destructor(101))) void deinit() {
   DP("Deinit target library!\n");
+
+  for (auto *R : PM->RTLs.UsedRTLs) {
+    // Plugins can either destroy their local state using global variables
+    // or attribute(destructor) functions or by implementing deinit_plugin
+    // The hazard with plugin local destructors is they may be called before
+    // or after this destructor. If the plugin is destroyed using global
+    // state before this library finishes calling into it the plugin is
+    // likely to crash. If good fortune means the plugin outlives this
+    // library then there may be no crash.
+    // Using deinit_plugin and no global destructors from the plugin works.
+    if (R->deinit_plugin) {
+      if (R->deinit_plugin() != OFFLOAD_SUCCESS) {
+        DP("Failure deinitializing RTL %s!\n", R->RTLName.c_str());
+      }
+    }
+  }
+
   delete PM;
 
 #ifdef OMPTARGET_PROFILE_ENABLED
@@ -567,13 +584,6 @@ void RTLsTy::unregisterLib(__tgt_bin_desc *Desc) {
   PM->TblMapMtx.unlock();
 
   // TODO: Write some RTL->unload_image(...) function?
-  for (auto *R : UsedRTLs) {
-    if (R->deinit_plugin) {
-      if (R->deinit_plugin() != OFFLOAD_SUCCESS) {
-        DP("Failure deinitializing RTL %s!\n", R->RTLName.c_str());
-      }
-    }
-  }
 
   DP("Done unregistering library!\n");
 }