[UR][SYCL] Add urUSMContextMemcpyExp API to enable device global support. (#17268)

Also adds a path using the new api to avoid a workaround introduced in
#16565

Author: aarongreig

sycl/source/detail/device_global_map_entry.cpp

@@ -87,6 +87,48 @@ DeviceGlobalMapEntry::getOrAllocateDeviceGlobalUSM(queue_impl &QueueImpl) {
   return NewAlloc;
 }
 
+DeviceGlobalUSMMem &
+DeviceGlobalMapEntry::getOrAllocateDeviceGlobalUSM(const context &Context) {
+  assert(!MIsDeviceImageScopeDecorated &&
+         "USM allocations should not be acquired for device_global with "
+         "device_image_scope property.");
+  const std::shared_ptr<context_impl> &CtxImpl = getSyclObjImpl(Context);
+  const std::shared_ptr<device_impl> &DevImpl =
+      getSyclObjImpl(CtxImpl->getDevices().front());
+  std::lock_guard<std::mutex> Lock(MDeviceToUSMPtrMapMutex);
+
+  auto DGUSMPtr = MDeviceToUSMPtrMap.find({DevImpl.get(), CtxImpl.get()});
+  if (DGUSMPtr != MDeviceToUSMPtrMap.end())
+    return DGUSMPtr->second;
+
+  void *NewDGUSMPtr = detail::usm::alignedAllocInternal(
+      0, MDeviceGlobalTSize, CtxImpl.get(), DevImpl.get(),
+      sycl::usm::alloc::device);
+
+  auto NewAllocIt = MDeviceToUSMPtrMap.emplace(
+      std::piecewise_construct,
+      std::forward_as_tuple(DevImpl.get(), CtxImpl.get()),
+      std::forward_as_tuple(NewDGUSMPtr));
+  assert(NewAllocIt.second &&
+         "USM allocation for device and context already happened.");
+  DeviceGlobalUSMMem &NewAlloc = NewAllocIt.first->second;
+
+  // C++ guarantees members appear in memory in the order they are declared,
+  // so since the member variable that contains the initial contents of the
+  // device_global is right after the usm_ptr member variable we can do
+  // some pointer arithmetic to memcopy over this value to the usm_ptr. This
+  // value inside of the device_global will be zero-initialized if it was not
+  // given a value on construction.
+  MemoryManager::context_copy_usm(
+      reinterpret_cast<const void *>(
+          reinterpret_cast<uintptr_t>(MDeviceGlobalPtr) +
+          sizeof(MDeviceGlobalPtr)),
+      CtxImpl, MDeviceGlobalTSize, NewAlloc.MPtr);
+
+  CtxImpl->addAssociatedDeviceGlobal(MDeviceGlobalPtr);
+  return NewAlloc;
+}
+
 void DeviceGlobalMapEntry::removeAssociatedResources(
     const context_impl *CtxImpl) {
   std::lock_guard<std::mutex> Lock{MDeviceToUSMPtrMapMutex};

sycl/source/detail/device_global_map_entry.hpp

@@ -111,6 +111,11 @@ struct DeviceGlobalMapEntry {
   // Gets or allocates USM memory for a device_global.
   DeviceGlobalUSMMem &getOrAllocateDeviceGlobalUSM(queue_impl &QueueImpl);
 
+  // This overload allows the allocation to be initialized without a queue. The
+  // UR adapter in use must report true for
+  // UR_DEVICE_INFO_USM_CONTEXT_MEMCPY_SUPPORT_EXP to take advantage of this.
+  DeviceGlobalUSMMem &getOrAllocateDeviceGlobalUSM(const context &Context);
+
   // Removes resources for device_globals associated with the context.
   void removeAssociatedResources(const context_impl *CtxImpl);
 

sycl/source/detail/kernel_bundle_impl.hpp

@@ -665,13 +665,21 @@ class kernel_bundle_impl {
                             "'device_image_scope' property");
     }
 
-    // TODO: Add context-only initialization via `urUSMContextMemcpyExp` instead
-    // of using a throw-away queue.
-    queue InitQueue{MContext, Dev};
-    auto &USMMem =
-        Entry->getOrAllocateDeviceGlobalUSM(*getSyclObjImpl(InitQueue));
-    InitQueue.wait_and_throw();
-    return USMMem.getPtr();
+    const auto &DeviceImpl = getSyclObjImpl(Dev);
+    bool SupportContextMemcpy = false;
+    DeviceImpl->getAdapter()->call<UrApiKind::urDeviceGetInfo>(
+        DeviceImpl->getHandleRef(),
+        UR_DEVICE_INFO_USM_CONTEXT_MEMCPY_SUPPORT_EXP,
+        sizeof(SupportContextMemcpy), &SupportContextMemcpy, nullptr);
+    if (SupportContextMemcpy) {
+      return Entry->getOrAllocateDeviceGlobalUSM(MContext).getPtr();
+    } else {
+      queue InitQueue{MContext, Dev};
+      auto &USMMem =
+          Entry->getOrAllocateDeviceGlobalUSM(*getSyclObjImpl(InitQueue));
+      InitQueue.wait_and_throw();
+      return USMMem.getPtr();
+    }
   }
 
   size_t ext_oneapi_get_device_global_size(const std::string &Name) const {

sycl/source/detail/memory_manager.cpp

@@ -891,6 +891,16 @@ void MemoryManager::copy_usm(const void *SrcMem, queue_impl &SrcQueue,
                                                DepEvents.data(), OutEvent);
 }
 
+void MemoryManager::context_copy_usm(const void *SrcMem, ContextImplPtr Context,
+                                     size_t Len, void *DstMem) {
+  if (!SrcMem || !DstMem)
+    throw exception(make_error_code(errc::invalid),
+                    "NULL pointer argument in memory copy operation.");
+  const AdapterPtr &Adapter = Context->getAdapter();
+  Adapter->call<UrApiKind::urUSMContextMemcpyExp>(Context->getHandleRef(),
+                                                  DstMem, SrcMem, Len);
+}
+
 void MemoryManager::fill_usm(void *Mem, queue_impl &Queue, size_t Length,
                              const std::vector<unsigned char> &Pattern,
                              std::vector<ur_event_handle_t> DepEvents,

sycl/source/detail/memory_manager.hpp

@@ -140,6 +140,9 @@ class MemoryManager {
                        void *DstMem, std::vector<ur_event_handle_t> DepEvents,
                        ur_event_handle_t *OutEvent);
 
+  static void context_copy_usm(const void *SrcMem, ContextImplPtr Context,
+                               size_t Len, void *DstMem);
+
   static void fill_usm(void *DstMem, queue_impl &Queue, size_t Len,
                        const std::vector<unsigned char> &Pattern,
                        std::vector<ur_event_handle_t> DepEvents,