Add get_address, get_size, and attempt to design a copy-API.

jopperm · jopperm · commit d45415dce958 · 2025-01-13T09:17:40.000Z
Signed-off-by: Julian Oppermann &lt;julian.oppermann@codeplay.com&gt;
diff --git a/sycl/include/sycl/kernel_bundle.hpp b/sycl/include/sycl/kernel_bundle.hpp
@@ -206,6 +206,30 @@ class __SYCL_EXPORT kernel_bundle_plain {
     return ext_oneapi_has_device_global(detail::string_view{name}, dev);
   }
 
+  void *ext_oneapi_get_device_global_address(const std::string &name,
+                                             const device &dev) {
+    return ext_oneapi_get_device_global_address(detail::string_view{name}, dev);
+  }
+
+  size_t ext_oneapi_get_device_global_size(const std::string &name,
+                                           const device &dev) {
+    return ext_oneapi_get_device_global_size(detail::string_view{name}, dev);
+  }
+
+  template <typename T>
+  event ext_oneapi_copy_to_device_global(const std::string &dest, const T &src,
+                                         const queue &queue) {
+    return ext_oneapi_copy_to_device_global(detail::string_view{dest}, &src,
+                                            sizeof(T), queue);
+  }
+
+  template <typename T>
+  event ext_oneapi_copy_from_device_global(T &dest, const std::string &src,
+                                           const queue &queue) {
+    return ext_oneapi_copy_from_device_global(&dest, detail::string_view{src},
+                                              sizeof(T), queue);
+  }
+
 protected:
   // \returns a kernel object which represents the kernel identified by
   // kernel_id passed
@@ -237,6 +261,16 @@ class __SYCL_EXPORT kernel_bundle_plain {
 
   bool ext_oneapi_has_device_global(detail::string_view name,
                                     const device &dev);
+  void *ext_oneapi_get_device_global_address(detail::string_view name,
+                                             const device &dev);
+  size_t ext_oneapi_get_device_global_size(detail::string_view name,
+                                           const device &dev);
+  event ext_oneapi_copy_to_device_global(detail::string_view dest,
+                                         const void *src, size_t num_bytes,
+                                         const queue &queue);
+  event ext_oneapi_copy_from_device_global(void *dest, detail::string_view src,
+                                           size_t num_bytes,
+                                           const queue &queue);
 };
 
 } // namespace detail
diff --git a/sycl/source/detail/kernel_bundle_impl.hpp b/sycl/source/detail/kernel_bundle_impl.hpp
@@ -674,6 +674,8 @@ class kernel_bundle_impl {
                                                 KernelNames, Language);
   }
 
+  // Utility methods for kernel_compiler functionality
+private:
   std::string adjust_kernel_name(const std::string &Name,
                                  syclex::source_language Lang) {
     // Once name demangling support is in, we won't need this.
@@ -685,6 +687,35 @@ class kernel_bundle_impl {
     return isMangled ? Name : "__sycl_kernel_" + Name;
   }
 
+  std::string mangle_device_global_name(const std::string &Name) {
+    // TODO: Support device globals declared in namespaces.
+    return "_Z" + std::to_string(Name.length()) + Name;
+  }
+
+  const DeviceGlobalMapEntry *get_device_global_entry(const std::string &Name,
+                                                      const device &Dev) {
+    if (Language != syclex::source_language::sycl_jit || Prefix.empty()) {
+      throw sycl::exception(make_error_code(errc::invalid),
+                            "Querying device globals by name is only available "
+                            "in kernel_bundles successfully built from "
+                            "kernel_bundle<bundle_state:ext_oneapi_source> "
+                            "with 'sycl_jit' source language.");
+    }
+
+    if (!ext_oneapi_has_device_global(Name, Dev)) {
+      throw sycl::exception(make_error_code(errc::invalid),
+                            "device global '" + Name +
+                                "' not found in kernel_bundle");
+    }
+
+    std::vector<DeviceGlobalMapEntry *> Entries =
+        ProgramManager::getInstance().getDeviceGlobalEntries(
+            {Prefix + mangle_device_global_name(Name)});
+    assert(Entries.size() == 1);
+    return Entries.front();
+  }
+
+public:
   bool ext_oneapi_has_kernel(const std::string &Name) {
     auto it = std::find(KernelNames.begin(), KernelNames.end(),
                         adjust_kernel_name(Name, Language));
@@ -746,18 +777,58 @@ class kernel_bundle_impl {
     return detail::createSyclObjFromImpl<kernel>(KernelImpl);
   }
 
-  std::string mangle_device_global_name(const std::string &Name) {
-    // TODO: Support device globals declared in namespaces.
-    return "_Z" + std::to_string(Name.length()) + Name;
-  }
-
   bool ext_oneapi_has_device_global(const std::string &Name,
-                                    [[maybe_unused]] const device &Dev) {
+                                    const device &Dev) {
+    if (!std::any_of(
+            MDevices.begin(), MDevices.end(),
+            [&Dev](const device &DevCand) { return Dev == DevCand; })) {
+      // TODO: device_image::has_kernel(id, device) checks the  device if the
+      //       given device is a sub-device.
+      return false;
+    }
+
     std::string MangledName = mangle_device_global_name(Name);
     return std::find(DeviceGlobalNames.begin(), DeviceGlobalNames.end(),
                      MangledName) != DeviceGlobalNames.end();
   }
 
+  void *ext_oneapi_get_device_global_address(const std::string &Name,
+                                             const device &Dev) {
+    return const_cast<void *>(
+        get_device_global_entry(Name, Dev)->MDeviceGlobalPtr);
+  }
+
+  size_t ext_oneapi_get_device_global_size(const std::string &Name,
+                                           const device &Dev) {
+    return get_device_global_entry(Name, Dev)->MDeviceGlobalTSize;
+  }
+
+  event ext_oneapi_copy_to_device_global(const std::string &Dest,
+                                         const void *Src, size_t NumBytes,
+                                         const queue &Queue) {
+    const auto *Entry = get_device_global_entry(Dest, Queue.get_device());
+    if (NumBytes != Entry->MDeviceGlobalTSize) {
+      throw sycl::exception(make_error_code(errc::invalid),
+                            "Incompatible type size for device global '" +
+                                Dest + "'");
+    }
+    return syclex::detail::SYCL_JIT_memcpy_to_device_global(
+        Entry, Src, NumBytes, /*Offset=*/0, Queue, /*DepEvents=*/{});
+  }
+
+  event ext_oneapi_copy_from_device_global(void *Dest, const std::string &Src,
+                                           size_t NumBytes,
+                                           const queue &Queue) {
+    const auto *Entry = get_device_global_entry(Src, Queue.get_device());
+    if (NumBytes != Entry->MDeviceGlobalTSize) {
+      throw sycl::exception(make_error_code(errc::invalid),
+                            "Incompatible type size for device global '" + Src +
+                                "'");
+    }
+    return syclex::detail::SYCL_JIT_memcpy_from_device_global(
+        Dest, Entry, NumBytes, /*Offset=*/0, Queue, /*DepEvents=*/{});
+  }
+
   bool empty() const noexcept { return MDeviceImages.empty(); }
 
   backend get_backend() const noexcept {
diff --git a/sycl/source/detail/kernel_compiler/kernel_compiler_sycl.cpp b/sycl/source/detail/kernel_compiler/kernel_compiler_sycl.cpp
@@ -288,6 +288,8 @@ bool SYCL_Compilation_Available() {
 #if SYCL_EXT_JIT_ENABLE
 #include "../jit_compiler.hpp"
 #include <atomic>
+#include <detail/device_global_map_entry.hpp>
+#include <detail/queue_impl.hpp>
 #endif
 
 namespace sycl {
@@ -323,6 +325,41 @@ std::pair<sycl_device_binaries, std::string> SYCL_JIT_to_SPIRV(
 #endif
 }
 
+event SYCL_JIT_memcpy_to_device_global(
+    [[maybe_unused]] const DeviceGlobalMapEntry *Dest,
+    [[maybe_unused]] const void *Src, [[maybe_unused]] size_t NumBytes,
+    [[maybe_unused]] size_t Offset, [[maybe_unused]] const queue &Queue,
+    [[maybe_unused]] const std::vector<event> &DepEvents) {
+#if SYCL_EXT_JIT_ENABLE
+  const std::shared_ptr<queue_impl> &QueueImplPtr = getSyclObjImpl(Queue);
+  return QueueImplPtr->memcpyToDeviceGlobal(
+      QueueImplPtr, const_cast<void *>(Dest->MDeviceGlobalPtr), Src,
+      Dest->MIsDeviceImageScopeDecorated, NumBytes, Offset, DepEvents,
+      /*CallerNeedsEvent=*/true);
+#else
+  throw sycl::exception(sycl::errc::invalid,
+                        "runtime-compiled device global support not available");
+#endif
+}
+
+event SYCL_JIT_memcpy_from_device_global(
+    [[maybe_unused]] void *Dest,
+    [[maybe_unused]] const DeviceGlobalMapEntry *Src,
+    [[maybe_unused]] size_t NumBytes, [[maybe_unused]] size_t Offset,
+    [[maybe_unused]] const queue &Queue,
+    [[maybe_unused]] const std::vector<event> &DepEvents) {
+#if SYCL_EXT_JIT_ENABLE
+  const std::shared_ptr<queue_impl> &QueueImplPtr = getSyclObjImpl(Queue);
+  return QueueImplPtr->memcpyFromDeviceGlobal(
+      QueueImplPtr, Dest, const_cast<void *>(Src->MDeviceGlobalPtr),
+      Src->MIsDeviceImageScopeDecorated, NumBytes, Offset, DepEvents,
+      /*CallerNeedsEvent=*/true);
+#else
+  throw sycl::exception(sycl::errc::invalid,
+                        "runtime-compiled device global support not available");
+#endif
+}
+
 } // namespace detail
 } // namespace ext::oneapi::experimental
 } // namespace _V1
diff --git a/sycl/source/detail/kernel_compiler/kernel_compiler_sycl.hpp b/sycl/source/detail/kernel_compiler/kernel_compiler_sycl.hpp
@@ -11,6 +11,7 @@
 #include <sycl/detail/defines_elementary.hpp>
 #include <sycl/detail/export.hpp> // __SYCL_EXPORT
 #include <sycl/device.hpp>
+#include <sycl/queue.hpp>
 
 #include <detail/compiler.hpp> // sycl_device_binaries
 
@@ -20,6 +21,11 @@
 
 namespace sycl {
 inline namespace _V1 {
+
+namespace detail {
+struct DeviceGlobalMapEntry;
+} // namespace detail
+
 namespace ext::oneapi::experimental {
 namespace detail {
 
@@ -42,6 +48,17 @@ SYCL_JIT_to_SPIRV(const std::string &Source, include_pairs_t IncludePairs,
 
 bool SYCL_JIT_Compilation_Available();
 
+event SYCL_JIT_memcpy_to_device_global(const DeviceGlobalMapEntry *Dest,
+                                       const void *Src, size_t NumBytes,
+                                       size_t Offset, const queue &Queue,
+                                       const std::vector<event> &DepEvents);
+
+event SYCL_JIT_memcpy_from_device_global(void *Dest,
+                                         const DeviceGlobalMapEntry *Src,
+                                         size_t NumBytes, size_t Offset,
+                                         const queue &Queue,
+                                         const std::vector<event> &DepEvents);
+
 } // namespace detail
 } // namespace ext::oneapi::experimental
 
diff --git a/sycl/source/kernel_bundle.cpp b/sycl/source/kernel_bundle.cpp
@@ -129,6 +129,30 @@ bool kernel_bundle_plain::ext_oneapi_has_device_global(detail::string_view name,
   return impl->ext_oneapi_has_device_global(name.data(), dev);
 }
 
+void *kernel_bundle_plain::ext_oneapi_get_device_global_address(
+    detail::string_view name, const device &dev) {
+  return impl->ext_oneapi_get_device_global_address(name.data(), dev);
+}
+
+size_t
+kernel_bundle_plain::ext_oneapi_get_device_global_size(detail::string_view name,
+                                                       const device &dev) {
+  return impl->ext_oneapi_get_device_global_size(name.data(), dev);
+}
+
+event kernel_bundle_plain::ext_oneapi_copy_to_device_global(
+    detail::string_view dest, const void *src, size_t num_bytes,
+    const queue &queue) {
+  return impl->ext_oneapi_copy_to_device_global(dest.data(), src, num_bytes,
+                                                queue);
+}
+
+event kernel_bundle_plain::ext_oneapi_copy_from_device_global(
+    void *dest, detail::string_view src, size_t num_bytes, const queue &queue) {
+  return impl->ext_oneapi_copy_from_device_global(dest, src.data(), num_bytes,
+                                                  queue);
+}
+
 //////////////////////////////////
 ///// sycl::detail free functions
 //////////////////////////////////
diff --git a/sycl/test-e2e/KernelCompiler/kernel_compiler_sycl_jit.cpp b/sycl/test-e2e/KernelCompiler/kernel_compiler_sycl_jit.cpp
