Apply similar fix to hip adapter.

Author: aarongreig

unified-runtime/source/adapters/hip/adapter.cpp

@@ -10,16 +10,14 @@
 
 #include "adapter.hpp"
 #include "common.hpp"
-#include "logger/ur_logger.hpp"
 
-#include <atomic>
 #include <ur_api.h>
 
-struct ur_adapter_handle_t_ {
-  std::atomic<uint32_t> RefCount = 0;
-  logger::Logger &logger;
-  ur_adapter_handle_t_();
-};
+#include <memory>
+
+namespace ur::hip {
+std::shared_ptr<ur_adapter_handle_t_> adapter;
+}
 
 class ur_legacy_sink : public logger::Sink {
 public:
@@ -42,7 +40,7 @@ class ur_legacy_sink : public logger::Sink {
 ur_adapter_handle_t_::ur_adapter_handle_t_()
     : logger(
           logger::get_logger("hip", /*default_log_level*/ logger::Level::ERR)) {
-
+  Platform = std::make_unique<ur_platform_handle_t_>();
   if (std::getenv("UR_LOG_HIP") != nullptr)
     return;
 
@@ -52,13 +50,16 @@ ur_adapter_handle_t_::ur_adapter_handle_t_()
   }
 }
 
-ur_adapter_handle_t_ adapter{};
-
 UR_APIEXPORT ur_result_t UR_APICALL urAdapterGet(
     uint32_t, ur_adapter_handle_t *phAdapters, uint32_t *pNumAdapters) {
   if (phAdapters) {
-    adapter.RefCount++;
-    *phAdapters = &adapter;
+    static std::once_flag InitFlag;
+    std::call_once(InitFlag, [=]() {
+      ur::hip::adapter = std::make_shared<ur_adapter_handle_t_>();
+    });
+
+    ur::hip::adapter->RefCount++;
+    *phAdapters = ur::hip::adapter.get();
   }
   if (pNumAdapters) {
     *pNumAdapters = 1;
@@ -69,12 +70,12 @@ UR_APIEXPORT ur_result_t UR_APICALL urAdapterGet(
 
 UR_APIEXPORT ur_result_t UR_APICALL urAdapterRelease(ur_adapter_handle_t) {
   // No state to clean up so we don't need to check for 0 references
-  adapter.RefCount--;
+  ur::hip::adapter->RefCount--;
   return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urAdapterRetain(ur_adapter_handle_t) {
-  adapter.RefCount++;
+  ur::hip::adapter->RefCount++;
   return UR_RESULT_SUCCESS;
 }
 
@@ -96,7 +97,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urAdapterGetInfo(ur_adapter_handle_t,
   case UR_ADAPTER_INFO_BACKEND:
     return ReturnValue(UR_ADAPTER_BACKEND_HIP);
   case UR_ADAPTER_INFO_REFERENCE_COUNT:
-    return ReturnValue(adapter.RefCount.load());
+    return ReturnValue(ur::hip::adapter->RefCount.load());
   case UR_ADAPTER_INFO_VERSION:
     return ReturnValue(uint32_t{1});
   default:

unified-runtime/source/adapters/hip/adapter.hpp

@@ -8,6 +8,24 @@
 //
 //===----------------------------------------------------------------------===//
 
-struct ur_adapter_handle_t_;
+#ifndef UR_HIP_ADAPTER_HPP_INCLUDED
+#define UR_HIP_ADAPTER_HPP_INCLUDED
 
-extern ur_adapter_handle_t_ adapter;
+#include "logger/ur_logger.hpp"
+#include "platform.hpp"
+
+#include <atomic>
+#include <memory>
+
+struct ur_adapter_handle_t_ {
+  std::atomic<uint32_t> RefCount = 1;
+  logger::Logger &logger;
+  std::unique_ptr<ur_platform_handle_t_> Platform;
+  ur_adapter_handle_t_();
+};
+
+namespace ur::hip {
+extern std::shared_ptr<ur_adapter_handle_t_> adapter;
+}
+
+#endif // UR_HIP_ADAPTER_HPP_INCLUDED

unified-runtime/source/adapters/hip/context.hpp

@@ -11,6 +11,7 @@
 
 #include <set>
 
+#include "adapter.hpp"
 #include "common.hpp"
 #include "device.hpp"
 #include "platform.hpp"
@@ -85,16 +86,14 @@ struct ur_context_handle_t_ {
     void operator()() { Function(UserData); }
   };
 
+  std::shared_ptr<ur_adapter_handle_t_> Adapter;
   std::vector<ur_device_handle_t> Devices;
 
   std::atomic_uint32_t RefCount;
 
   ur_context_handle_t_(const ur_device_handle_t *Devs, uint32_t NumDevices)
-      : Devices{Devs, Devs + NumDevices}, RefCount{1} {
-    for (auto &Dev : Devices) {
-      urDeviceRetain(Dev);
-    }
-  };
+      : Adapter(ur::hip::adapter), Devices{Devs, Devs + NumDevices},
+        RefCount{1} {};
 
   ~ur_context_handle_t_() {}
 

unified-runtime/source/adapters/hip/device.cpp

@@ -1174,7 +1174,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceCreateWithNativeHandle(
 
   // Get list of platforms
   uint32_t NumPlatforms = 0;
-  ur_adapter_handle_t AdapterHandle = &adapter;
+  ur_adapter_handle_t AdapterHandle = ur::hip::adapter.get();
   ur_result_t Result =
       urPlatformGet(&AdapterHandle, 1, 0, nullptr, &NumPlatforms);
   if (Result != UR_RESULT_SUCCESS)

unified-runtime/source/adapters/hip/platform.cpp

@@ -36,7 +36,7 @@ urPlatformGetInfo(ur_platform_handle_t, ur_platform_info_t propName,
     return ReturnValue("");
   }
   case UR_PLATFORM_INFO_ADAPTER: {
-    return ReturnValue(&adapter);
+    return ReturnValue(ur::hip::adapter.get());
   }
   default:
     return UR_RESULT_ERROR_INVALID_ENUMERATION;
@@ -56,7 +56,6 @@ urPlatformGet(ur_adapter_handle_t *, uint32_t, uint32_t NumEntries,
   try {
     static std::once_flag InitFlag;
     static uint32_t NumPlatforms = 1;
-    static ur_platform_handle_t_ Platform;
 
     UR_ASSERT(phPlatforms || pNumPlatforms, UR_RESULT_ERROR_INVALID_VALUE);
     UR_ASSERT(!phPlatforms || NumEntries > 0, UR_RESULT_ERROR_INVALID_VALUE);
@@ -87,18 +86,20 @@ urPlatformGet(ur_adapter_handle_t *, uint32_t, uint32_t NumEntries,
 
               // Use the default stream to record base event counter
               UR_CHECK_ERROR(hipEventRecord(EvBase, 0));
-              Platform.Devices.emplace_back(
-                  new ur_device_handle_t_{Device, EvBase, &Platform, i});
+              ur::hip::adapter->Platform->Devices.emplace_back(
+                  new ur_device_handle_t_{Device, EvBase,
+                                          ur::hip::adapter->Platform.get(), i});
 
-              ScopedDevice Active(Platform.Devices.front().get());
+              ScopedDevice Active(
+                  ur::hip::adapter->Platform->Devices.front().get());
             }
           } catch (const std::bad_alloc &) {
             // Signal out-of-memory situation
-            Platform.Devices.clear();
+            ur::hip::adapter->Platform->Devices.clear();
             Err = UR_RESULT_ERROR_OUT_OF_HOST_MEMORY;
           } catch (ur_result_t CatchErr) {
             // Clear and rethrow to allow retry
-            Platform.Devices.clear();
+            ur::hip::adapter->Platform->Devices.clear();
             Err = CatchErr;
             throw CatchErr;
           } catch (...) {
@@ -113,7 +114,7 @@ urPlatformGet(ur_adapter_handle_t *, uint32_t, uint32_t NumEntries,
     }
 
     if (phPlatforms != nullptr) {
-      *phPlatforms = &Platform;
+      *phPlatforms = ur::hip::adapter->Platform.get();
     }
 
     return Result;

unified-runtime/source/adapters/hip/usm.cpp

@@ -199,7 +199,7 @@ urUSMGetMemAllocInfo(ur_context_handle_t hContext, const void *pMem,
 
       // hip backend has only one platform containing all devices
       ur_platform_handle_t platform;
-      ur_adapter_handle_t AdapterHandle = &adapter;
+      ur_adapter_handle_t AdapterHandle = ur::hip::adapter.get();
       UR_CHECK_ERROR(urPlatformGet(&AdapterHandle, 1, 1, &platform, nullptr));
 
       // get the device from the platform