Add ext_oneapi_has_device_global.

Signed-off-by: Julian Oppermann <[email protected]>

Author: jopperm

sycl/include/sycl/kernel_bundle.hpp

@@ -201,6 +201,11 @@ class __SYCL_EXPORT kernel_bundle_plain {
     return ext_oneapi_get_kernel(detail::string_view{name});
   }
 
+  bool ext_oneapi_has_device_global(const std::string &name,
+                                    const device &dev) {
+    return ext_oneapi_has_device_global(detail::string_view{name}, dev);
+  }
+
 protected:
   // \returns a kernel object which represents the kernel identified by
   // kernel_id passed
@@ -229,6 +234,9 @@ class __SYCL_EXPORT kernel_bundle_plain {
 private:
   bool ext_oneapi_has_kernel(detail::string_view name);
   kernel ext_oneapi_get_kernel(detail::string_view name);
+
+  bool ext_oneapi_has_device_global(detail::string_view name,
+                                    const device &dev);
 };
 
 } // namespace detail
@@ -449,6 +457,17 @@ class kernel_bundle : public detail::kernel_bundle_plain,
     return detail::kernel_bundle_plain::ext_oneapi_get_kernel(name);
   }
 
+  /////////////////////////
+  // ext_oneapi_has_device_global
+  //  only true if created from source and has this global for the given device
+  /////////////////////////
+  template <bundle_state _State = State,
+            typename = std::enable_if_t<_State == bundle_state::executable>>
+  bool ext_oneapi_has_device_global(const std::string &name,
+                                    const device &dev) {
+    return detail::kernel_bundle_plain::ext_oneapi_has_device_global(name, dev);
+  }
+
 private:
   kernel_bundle(detail::KernelBundleImplPtr Impl)
       : kernel_bundle_plain(std::move(Impl)) {}

sycl/source/detail/kernel_bundle_impl.hpp

@@ -380,8 +380,9 @@ class kernel_bundle_impl {
   // program manager integration, only for sycl_jit language
   kernel_bundle_impl(context Ctx, std::vector<device> Devs,
                      const std::vector<kernel_id> &KernelIDs,
-                     std::vector<std::string> KNames, std::string Pfx,
-                     syclex::source_language Lang)
+                     const std::vector<std::string> &KNames,
+                     const std::vector<std::string> &DGNames,
+                     const std::string &Pfx, syclex::source_language Lang)
       : kernel_bundle_impl(Ctx, Devs, KernelIDs, bundle_state::executable) {
     assert(Lang == syclex::source_language::sycl_jit);
     // Mark this bundle explicitly as "interop" to ensure that its kernels are
@@ -391,6 +392,7 @@ class kernel_bundle_impl {
     // from the (unprefixed) kernel name.
     MIsInterop = true;
     KernelNames = KNames;
+    DeviceGlobalNames = DGNames;
     Prefix = Pfx;
     Language = Lang;
   }
@@ -509,51 +511,60 @@ class kernel_bundle_impl {
       // `jit_compiler::compileSYCL(..)` uses `CompilationID + '$'` as prefix
       // for offload entry names.
       std::string Prefix = CompilationID + '$';
+      auto PrefixLen = Prefix.length();
       for (const auto &KernelID : PM.getAllSYCLKernelIDs()) {
         std::string_view KernelName{KernelID.get_name()};
         if (KernelName.find(Prefix) == 0) {
           KernelIDs.push_back(KernelID);
-          KernelName.remove_prefix(Prefix.length());
+          KernelName.remove_prefix(PrefixLen);
           KernelNames.emplace_back(KernelName);
         }
       }
 
-      // Create the executable bundle.
-      auto ExecBundle = std::make_shared<kernel_bundle_impl>(
-          MContext, MDevices, KernelIDs, KernelNames, Prefix, Language);
-
       // Determine IDs of all device globals referenced by this bundle's
       // kernels. These IDs are also prefixed.
-      std::set<std::string> UniqueDeviceGlobalIDs;
-      std::vector<std::string> DeviceGlobalIDs;
+      std::set<std::string> DeviceGlobalIDSet;
+      std::vector<std::string> DeviceGlobalIDVec;
+      std::vector<std::string> DeviceGlobalNames;
       for (const auto &RawImg : PM.getRawDeviceImages(KernelIDs)) {
         for (const auto &DeviceGlobalProp : RawImg->getDeviceGlobals()) {
-          auto [It, Ins] = UniqueDeviceGlobalIDs.insert(DeviceGlobalProp->Name);
+          std::string_view DeviceGlobalName{DeviceGlobalProp->Name};
+          assert(DeviceGlobalName.find(Prefix) == 0);
+          auto [It, Ins] = DeviceGlobalIDSet.emplace(DeviceGlobalName);
           if (Ins) {
-            DeviceGlobalIDs.push_back(*It);
+            DeviceGlobalIDVec.emplace_back(DeviceGlobalName);
+            DeviceGlobalName.remove_prefix(PrefixLen);
+            DeviceGlobalNames.emplace_back(DeviceGlobalName);
           }
         }
       }
 
+      // Create the executable bundle.
+      auto ExecBundle = std::make_shared<kernel_bundle_impl>(
+          MContext, MDevices, KernelIDs, KernelNames, DeviceGlobalNames, Prefix,
+          Language);
+
+      // Device globals are usually statically allocated and registered in the
+      // integration footer, which we don't have in the RTC context. Instead, we
+      // dynamically allocate storage tied to the executable kernel bundle.
       for (auto *DeviceGlobalEntry :
-           PM.getDeviceGlobalEntries(DeviceGlobalIDs)) {
-        // Device globals without `device_image_scope` are usually statically
-        // allocated and registered in the integration footer, which we don't
-        // have in the RTC context. Instead, we dynamically allocate storage
-        // tied to the executable kernel bundle.
+           PM.getDeviceGlobalEntries(DeviceGlobalIDVec)) {
+
+        size_t AllocSize = DeviceGlobalEntry->MDeviceGlobalTSize; // init value
         if (!DeviceGlobalEntry->MIsDeviceImageScopeDecorated) {
-          auto Alloc = std::make_unique<std::byte[]>(
-              DeviceGlobalEntry->MDeviceGlobalTSize);
-          PM.addOrInitDeviceGlobalEntry(Alloc.get(),
-                                        DeviceGlobalEntry->MUniqueId.c_str());
-          ExecBundle->DeviceGlobals.push_back(std::move(Alloc));
+          // USM pointer. TODO: it's actually a decorated multi_ptr.
+          AllocSize += sizeof(void *);
         }
+        auto Alloc = std::make_unique<std::byte[]>(AllocSize);
+        std::string_view DeviceGlobalName{DeviceGlobalEntry->MUniqueId};
+        PM.addOrInitDeviceGlobalEntry(Alloc.get(), DeviceGlobalName.data());
+        ExecBundle->DeviceGlobalAllocations.push_back(std::move(Alloc));
 
         // Drop the RTC prefix from the entry's symbol name. Note that the PM
         // still manages this device global under its prefixed name.
-        assert(DeviceGlobalEntry->MUniqueId.find(Prefix) == 0);
-        DeviceGlobalEntry->MUniqueId =
-            DeviceGlobalEntry->MUniqueId.substr(Prefix.length());
+        assert(DeviceGlobalName.find(Prefix) == 0);
+        DeviceGlobalName.remove_prefix(PrefixLen);
+        DeviceGlobalEntry->MUniqueId = DeviceGlobalName;
       }
 
       return ExecBundle;
@@ -735,6 +746,18 @@ 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) {
+    std::string MangledName = mangle_device_global_name(Name);
+    return std::find(DeviceGlobalNames.begin(), DeviceGlobalNames.end(),
+                     MangledName) != DeviceGlobalNames.end();
+  }
+
   bool empty() const noexcept { return MDeviceImages.empty(); }
 
   backend get_backend() const noexcept {
@@ -993,11 +1016,13 @@ class kernel_bundle_impl {
   // Language is for both state::source and state::executable.
   syclex::source_language Language = syclex::source_language::opencl;
   const std::variant<std::string, std::vector<std::byte>> Source;
-  // only kernel_bundles created from source have KernelNames member.
+  // only kernel_bundles created from source have the following members.
   std::vector<std::string> KernelNames;
+  std::vector<std::string> DeviceGlobalNames;
   std::string Prefix;
   include_pairs_t IncludePairs;
-  std::vector<std::unique_ptr<std::byte[]>> DeviceGlobals;
+
+  std::vector<std::unique_ptr<std::byte[]>> DeviceGlobalAllocations;
 };
 
 } // namespace detail

sycl/source/kernel_bundle.cpp

@@ -124,6 +124,11 @@ kernel kernel_bundle_plain::ext_oneapi_get_kernel(detail::string_view name) {
   return impl->ext_oneapi_get_kernel(name.data(), impl);
 }
 
+bool kernel_bundle_plain::ext_oneapi_has_device_global(detail::string_view name,
+                                                       const device &dev) {
+  return impl->ext_oneapi_has_device_global(name.data(), dev);
+}
+
 //////////////////////////////////
 ///// sycl::detail free functions
 //////////////////////////////////

sycl/test-e2e/KernelCompiler/kernel_compiler_sycl_jit.cpp

@@ -285,6 +285,11 @@ int test_device_global() {
 
   exe_kb kbExe1 = syclex::build(kbSrc);
 
+  // Check presence of device global.
+  assert(kbExe1.ext_oneapi_has_device_global("DG", q.get_device()));
+  // Querying a non-existing device global shall not crash.
+  assert(!kbExe1.ext_oneapi_has_device_global("bogus_DG", q.get_device()));
+
   auto setK = kbExe1.ext_oneapi_get_kernel("ff_dg_setter");
   auto addK = kbExe1.ext_oneapi_get_kernel("ff_dg_adder");
   auto getK = kbExe1.ext_oneapi_get_kernel("ff_dg_getter");