[SYCL][RTC] Initial support for device globals

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

Author: jopperm

sycl-jit/jit-compiler/lib/rtc/DeviceCompilation.cpp

@@ -503,8 +503,9 @@ jit_compiler::performPostLink(std::unique_ptr<llvm::Module> Module,
       /*IROutputOnly=*/false, EmitOnlyKernelsAsEntryPoints);
   assert(Splitter->hasMoreSplits());
 
-  // TODO: Call `verifyNoCrossModuleDeviceGlobalUsage` if device globals shall
-  //       be processed.
+  if (auto Err = Splitter->verifyNoCrossModuleDeviceGlobalUsage()) {
+    return std::move(Err);
+  }
 
   // TODO: This allocation assumes that there are no further splits required,
   //       i.e. there are no mixed SYCL/ESIMD modules.
@@ -547,7 +548,7 @@ jit_compiler::performPostLink(std::unique_ptr<llvm::Module> Module,
     GlobalBinImageProps PropReq{
         /*EmitKernelParamInfo=*/true, /*EmitProgramMetadata=*/true,
         /*EmitExportedSymbols=*/true, /*EmitImportedSymbols=*/true,
-        /*DeviceGlobals=*/false};
+        /*DeviceGlobals=*/true};
     PropertySetRegistry Properties =
         computeModuleProperties(MDesc.getModule(), MDesc.entries(), PropReq);
     // TODO: Manually add `compile_target` property as in

sycl/source/detail/jit_compiler.cpp

@@ -1150,14 +1150,19 @@ sycl_device_binaries jit_compiler::createDeviceBinaryImage(
     }
 
     for (const auto &FPS : DevImgInfo.Properties) {
+      bool IsDeviceGlobalsPropSet =
+          FPS.Name == __SYCL_PROPERTY_SET_SYCL_DEVICE_GLOBALS;
       PropertySetContainer PropSet{FPS.Name.c_str()};
       for (const auto &FPV : FPS.Values) {
         if (FPV.IsUIntValue) {
           PropSet.addProperty(
               PropertyContainer{FPV.Name.c_str(), FPV.UIntValue});
         } else {
+          std::string PrefixedName =
+              (IsDeviceGlobalsPropSet ? OffloadEntryPrefix : "") +
+              FPV.Name.c_str();
           PropSet.addProperty(PropertyContainer{
-              FPV.Name.c_str(), FPV.Bytes.begin(), FPV.Bytes.size(),
+              PrefixedName.c_str(), FPV.Bytes.begin(), FPV.Bytes.size(),
               sycl_property_type::SYCL_PROPERTY_TYPE_BYTE_ARRAY});
         }
       }

sycl/source/detail/kernel_bundle_impl.hpp

@@ -518,8 +518,45 @@ class kernel_bundle_impl {
         }
       }
 
-      return std::make_shared<kernel_bundle_impl>(
+      // 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;
+      for (const auto &RawImg : PM.getRawDeviceImages(KernelIDs)) {
+        for (const auto &DeviceGlobalProp : RawImg->getDeviceGlobals()) {
+          auto [It, Ins] = UniqueDeviceGlobalIDs.insert(DeviceGlobalProp->Name);
+          if (Ins) {
+            DeviceGlobalIDs.push_back(*It);
+          }
+        }
+      }
+
+      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.
+        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));
+        }
+
+        // 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());
+      }
+
+      return ExecBundle;
     }
 
     ur_program_handle_t UrProgram = nullptr;
@@ -960,6 +997,7 @@ class kernel_bundle_impl {
   std::vector<std::string> KernelNames;
   std::string Prefix;
   include_pairs_t IncludePairs;
+  std::vector<std::unique_ptr<std::byte[]>> DeviceGlobals;
 };
 
 } // namespace detail

sycl/test-e2e/KernelCompiler/kernel_compiler_sycl_jit.cpp

@@ -58,6 +58,29 @@ void ff_templated(T *ptr, T *unused) {
 }
 )===";
 
+auto constexpr DGSource = R"===(
+#include <sycl/sycl.hpp>
+
+namespace syclex = sycl::ext::oneapi::experimental;
+
+syclex::device_global<int> DG;
+
+extern "C" SYCL_EXTERNAL SYCL_EXT_ONEAPI_FUNCTION_PROPERTY(
+    (syclex::single_task_kernel)) void ff_dg_setter(int val) {
+  DG = val;
+}
+
+extern "C" SYCL_EXTERNAL SYCL_EXT_ONEAPI_FUNCTION_PROPERTY(
+    (syclex::single_task_kernel)) void ff_dg_adder(int val) {
+  DG = DG + val;
+}
+
+extern "C" SYCL_EXTERNAL SYCL_EXT_ONEAPI_FUNCTION_PROPERTY(
+    (syclex::single_task_kernel)) void ff_dg_getter(int *val) {
+  *val = DG;
+}
+)===";
+
 auto constexpr ESIMDSource = R"===(
 #include <sycl/sycl.hpp>
 #include <sycl/ext/intel/esimd.hpp>
@@ -219,6 +242,73 @@ int test_build_and_run() {
   return 0;
 }
 
+int test_device_global() {
+  namespace syclex = sycl::ext::oneapi::experimental;
+  using source_kb = sycl::kernel_bundle<sycl::bundle_state::ext_oneapi_source>;
+  using exe_kb = sycl::kernel_bundle<sycl::bundle_state::executable>;
+
+  sycl::queue q;
+  sycl::context ctx = q.get_context();
+
+  bool ok =
+      q.get_device().ext_oneapi_can_compile(syclex::source_language::sycl_jit);
+  if (!ok) {
+    std::cout << "Apparently this device does not support `sycl_jit` source "
+                 "kernel bundle extension: "
+              << q.get_device().get_info<sycl::info::device::name>()
+              << std::endl;
+    return -1;
+  }
+
+  auto modifyDG = [&q](sycl::kernel &k, int val) {
+    q.submit([&](sycl::handler &CGH) {
+      CGH.set_arg(0, val);
+      CGH.single_task(k);
+    });
+    q.wait();
+  };
+
+  auto getDG = [&q](sycl::kernel &k) -> int {
+    int *buf = sycl::malloc_shared<int>(1, q);
+    q.submit([&](sycl::handler &CGH) {
+      CGH.set_arg(0, buf);
+      CGH.single_task(k);
+    });
+    q.wait();
+    int val = *buf;
+    sycl::free(buf, q);
+    return val;
+  };
+
+  source_kb kbSrc = syclex::create_kernel_bundle_from_source(
+      ctx, syclex::source_language::sycl_jit, DGSource);
+
+  exe_kb kbExe1 = syclex::build(kbSrc);
+
+  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");
+
+  assert(getDG(getK) == 0);
+  modifyDG(setK, 42);
+  assert(getDG(getK) == 42);
+  modifyDG(addK, 1);
+  assert(getDG(getK) == 43);
+
+  exe_kb kbExe2 = syclex::build(kbSrc);
+
+  auto setK2 = kbExe2.ext_oneapi_get_kernel("ff_dg_setter");
+  auto getK2 = kbExe2.ext_oneapi_get_kernel("ff_dg_getter");
+
+  // `DG` is private per RTC bundle
+  assert(getDG(getK2) == 0);
+  modifyDG(setK2, -17);
+  assert(getDG(getK2) == -17);
+  assert(getDG(getK) == 43);
+
+  return 0;
+}
+
 int test_esimd() {
   namespace syclex = sycl::ext::oneapi::experimental;
   using source_kb = sycl::kernel_bundle<sycl::bundle_state::ext_oneapi_source>;
@@ -390,8 +480,8 @@ int test_warning() {
 int main(int argc, char **) {
 #ifdef SYCL_EXT_ONEAPI_KERNEL_COMPILER
   int optional_tests = (argc > 1) ? test_warning() : 0;
-  return test_build_and_run() || test_esimd() || test_unsupported_options() ||
-         test_error() || optional_tests;
+  return test_build_and_run() || test_device_global() || test_esimd() ||
+         test_unsupported_options() || test_error() || optional_tests;
 #else
   static_assert(false, "Kernel Compiler feature test macro undefined");
 #endif