[SYCL] Expand device_global map and make initialization order agnostic (#5902)

This commit makes the following changes:
 * Adds the unique identifier to the device_global map entries.
 * Makes the map of USM allocations for device_globals private and
   add context to the map key.
 * Makes the order of initialization of device_global entries in the map
   order agnostic, i.e. the runtime will not care if the integration
   header or the device-image initialization happens first.
 * Fixes the export of the `add` function for the map, as used by the
   integration header.

Signed-off-by: Steffen Larsen <[email protected]>

Author: steffenlarsen

sycl/include/CL/sycl/detail/device_global_map.hpp

@@ -13,7 +13,7 @@ namespace sycl {
 namespace detail {
 namespace device_global_map {
 
-void add(void *DeviceGlobalPtr, const char *UniqueId);
+__SYCL_EXPORT void add(const void *DeviceGlobalPtr, const char *UniqueId);
 
 } // namespace device_global_map
 } // namespace detail

sycl/source/detail/device_global_map.cpp

@@ -13,9 +13,9 @@ namespace sycl {
 namespace detail {
 namespace device_global_map {
 
-void add(void *DeviceGlobalPtr, const char *UniqueId) {
-  detail::ProgramManager::getInstance().addDeviceGlobalEntry(DeviceGlobalPtr,
-                                                             UniqueId);
+__SYCL_EXPORT void add(const void *DeviceGlobalPtr, const char *UniqueId) {
+  detail::ProgramManager::getInstance().addOrInitDeviceGlobalEntry(
+      DeviceGlobalPtr, UniqueId);
 }
 
 } // namespace device_global_map

sycl/source/detail/device_global_map_entry.hpp

@@ -9,7 +9,8 @@
 #pragma once
 
 #include <cstdint>
-#include <unordered_map>
+#include <map>
+#include <mutex>
 
 __SYCL_INLINE_NAMESPACE(cl) {
 namespace sycl {
@@ -19,23 +20,39 @@ namespace detail {
 class device_impl;
 
 struct DeviceGlobalMapEntry {
+  // The unique identifier of the device_global.
+  std::string MUniqueId;
   // Pointer to the device_global on host.
-  void *MDeviceGlobalPtr;
+  const void *MDeviceGlobalPtr;
   // Size of the underlying type in the device_global.
   std::uint32_t MDeviceGlobalTSize;
   // True if the device_global has been decorated with device_image_scope
   bool MIsDeviceImageScopeDecorated;
-  // Map between devices and corresponding USM allocations for the
-  // device_global. This should always be empty if MIsDeviceImageScopeDecorated
-  // is true.
-  std::unordered_map<std::shared_ptr<device_impl>, void *> MDeviceToUSMPtrMap;
 
-  // Constructor only initializes with the pointer to the device_global as the
-  // additional information is loaded after.
-  DeviceGlobalMapEntry(void *DeviceGlobalPtr)
-      : MDeviceGlobalPtr(DeviceGlobalPtr), MDeviceGlobalTSize(0),
-        MIsDeviceImageScopeDecorated(false) {}
+  // Constructor for only initializing ID and pointer. The other members will
+  // be initialized later.
+  DeviceGlobalMapEntry(std::string UniqueId, const void *DeviceGlobalPtr)
+      : MUniqueId(UniqueId), MDeviceGlobalPtr(DeviceGlobalPtr),
+        MDeviceGlobalTSize(0), MIsDeviceImageScopeDecorated(false) {}
+
+  // Constructor for only initializing ID, type size, and device image scope
+  // flag. The pointer to the device global will be initialized later.
+  DeviceGlobalMapEntry(std::string UniqueId, std::uint32_t DeviceGlobalTSize,
+                       bool IsDeviceImageScopeDecorated)
+      : MUniqueId(UniqueId), MDeviceGlobalPtr(nullptr),
+        MDeviceGlobalTSize(DeviceGlobalTSize),
+        MIsDeviceImageScopeDecorated(IsDeviceImageScopeDecorated) {}
+
+  // Initialize the pointer to the associated device_global.
+  void initialize(const void *DeviceGlobalPtr) {
+    assert(DeviceGlobalPtr && "Device global pointer cannot be null");
+    assert(!MDeviceGlobalPtr &&
+           "Device global pointer has already been initialized.");
+    MDeviceGlobalPtr = DeviceGlobalPtr;
+  }
 
+  // Initialize the device_global's element type size and the flag signalling
+  // if the device_global has the device_image_scope property.
   void initialize(std::uint32_t DeviceGlobalTSize,
                   bool IsDeviceImageScopeDecorated) {
     assert(DeviceGlobalTSize != 0 && "Device global initialized with 0 size.");
@@ -44,6 +61,14 @@ struct DeviceGlobalMapEntry {
     MDeviceGlobalTSize = DeviceGlobalTSize;
     MIsDeviceImageScopeDecorated = IsDeviceImageScopeDecorated;
   }
+
+private:
+  // Map from a device and a context to the associated USM allocation for the
+  // device_global. This should always be empty if MIsDeviceImageScopeDecorated
+  // is true.
+  std::map<std::pair<const device_impl *, const context_impl *>, void *>
+      MDeviceToUSMPtrMap;
+  std::mutex MDeviceToUSMPtrMapMutex;
 };
 
 } // namespace detail

sycl/source/detail/program_manager/program_manager.cpp

@@ -1217,7 +1217,20 @@ void ProgramManager::addImages(pi_device_binaries DeviceBinary) {
               *reinterpret_cast<const std::uint32_t *>(&DeviceGlobalInfo[8]);
           const std::uint32_t DeviceImageScopeDecorated =
               *reinterpret_cast<const std::uint32_t *>(&DeviceGlobalInfo[12]);
-          Entry->second.initialize(TypeSize, DeviceImageScopeDecorated);
+
+          auto ExistingDeviceGlobal = m_DeviceGlobals.find(DeviceGlobal->Name);
+          if (ExistingDeviceGlobal != m_DeviceGlobals.end()) {
+            // If it has already been registered we update the information.
+            ExistingDeviceGlobal->second->initialize(TypeSize,
+                                                     DeviceImageScopeDecorated);
+          } else {
+            // If it has not already been registered we create a new entry.
+            // Note: Pointer to the device global is not available here, so it
+            //       cannot be set until registration happens.
+            auto EntryUPtr = std::make_unique<DeviceGlobalMapEntry>(
+                DeviceGlobal->Name, TypeSize, DeviceImageScopeDecorated);
+            m_DeviceGlobals.emplace(DeviceGlobal->Name, std::move(EntryUPtr));
+          }
         }
       }
       m_DeviceImages[KSId].reset(new std::vector<RTDeviceBinaryImageUPtr>());
@@ -1469,13 +1482,23 @@ kernel_id ProgramManager::getBuiltInKernelID(const std::string &KernelName) {
   return KernelID->second;
 }
 
-void ProgramManager::addDeviceGlobalEntry(void *DeviceGlobalPtr,
-                                          const char *UniqueId) {
+void ProgramManager::addOrInitDeviceGlobalEntry(const void *DeviceGlobalPtr,
+                                                const char *UniqueId) {
   std::lock_guard<std::mutex> DeviceGlobalsGuard(m_DeviceGlobalsMutex);
 
-  assert(m_DeviceGlobals.find(UniqueId) == m_DeviceGlobals.end() &&
-         "Device global has already been registered.");
-  m_DeviceGlobals.insert({UniqueId, DeviceGlobalMapEntry(DeviceGlobalPtr)});
+  auto ExistingDeviceGlobal = m_DeviceGlobals.find(UniqueId);
+  if (ExistingDeviceGlobal != m_DeviceGlobals.end()) {
+    // Update the existing information and add the entry to the pointer map.
+    ExistingDeviceGlobal->second->initialize(DeviceGlobalPtr);
+    m_Ptr2DeviceGlobal.insert(
+        {DeviceGlobalPtr, ExistingDeviceGlobal->second.get()});
+    return;
+  }
+
+  auto EntryUPtr =
+      std::make_unique<DeviceGlobalMapEntry>(UniqueId, DeviceGlobalPtr);
+  auto NewEntry = m_DeviceGlobals.emplace(UniqueId, std::move(EntryUPtr));
+  m_Ptr2DeviceGlobal.insert({DeviceGlobalPtr, NewEntry.first->second.get()});
 }
 
 std::vector<device_image_plain>

sycl/source/detail/program_manager/program_manager.hpp

@@ -183,8 +183,10 @@ class ProgramManager {
   // built-in kernel name.
   kernel_id getBuiltInKernelID(const std::string &KernelName);
 
-  // The function inserts a device_global entry into the device_global map.
-  void addDeviceGlobalEntry(void *DeviceGlobalPtr, const char *UniqueId);
+  // The function inserts or initializes a device_global entry into the
+  // device_global map.
+  void addOrInitDeviceGlobalEntry(const void *DeviceGlobalPtr,
+                                  const char *UniqueId);
 
   // The function returns a vector of SYCL device images that are compiled with
   // the required state and at least one device from the passed list of devices.
@@ -387,10 +389,12 @@ class ProgramManager {
   using KernelNameWithOSModule = std::pair<std::string, OSModuleHandle>;
   std::set<KernelNameWithOSModule> m_KernelUsesAssert;
 
-  // Map between device_global unique ids and associated information.
-  std::unordered_map<std::string, DeviceGlobalMapEntry> m_DeviceGlobals;
+  // Maps between device_global identifiers and associated information.
+  std::unordered_map<std::string, std::unique_ptr<DeviceGlobalMapEntry>>
+      m_DeviceGlobals;
+  std::unordered_map<const void *, DeviceGlobalMapEntry *> m_Ptr2DeviceGlobal;
 
-  /// Protects m_DeviceGlobals.
+  /// Protects m_DeviceGlobals and m_Ptr2DeviceGlobal.
   std::mutex m_DeviceGlobalsMutex;
 };
 } // namespace detail

sycl/test/abi/sycl_symbols_linux.dump

@@ -3895,6 +3895,7 @@ _ZN2cl4sycl6detail16AccessorImplHostD1Ev
 _ZN2cl4sycl6detail16AccessorImplHostD2Ev
 _ZN2cl4sycl6detail17HostProfilingInfo3endEv
 _ZN2cl4sycl6detail17HostProfilingInfo5startEv
+_ZN2cl4sycl6detail17device_global_map3addEPKvPKc
 _ZN2cl4sycl6detail18convertChannelTypeE22_pi_image_channel_type
 _ZN2cl4sycl6detail18convertChannelTypeENS0_18image_channel_typeE
 _ZN2cl4sycl6detail18get_kernel_id_implENSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE

sycl/test/abi/sycl_symbols_windows.dump

@@ -1090,6 +1090,7 @@
 ?acospi@__host_std@cl@@YA?AVhalf@half_impl@detail@sycl@2@V34562@@Z
 ?acospi@__host_std@cl@@YAMM@Z
 ?acospi@__host_std@cl@@YANN@Z
+?add@device_global_map@detail@sycl@cl@@YAXPEBXPEBD@Z
 ?addHostAccessorAndWait@detail@sycl@cl@@YAXPEAVAccessorImplHost@123@@Z
 ?addOrReplaceAccessorProperties@SYCLMemObjT@detail@sycl@cl@@QEAAXAEBVproperty_list@34@@Z
 ?addReduction@handler@sycl@cl@@AEAAXAEBV?$shared_ptr@$$CBX@std@@@Z