[SYCL][RTC] Query kernels by source code name

Signed-off-by: Julian Oppermann <julian.oppermann@codeplay.com>

Author: jopperm

sycl/source/detail/compiler.hpp

@@ -68,6 +68,8 @@
 #define __SYCL_PROPERTY_SET_SYCL_VIRTUAL_FUNCTIONS "SYCL/virtual functions"
 /// PropertySetRegistry::SYCL_IMPLICIT_LOCAL_ARG defined in PropertySetIO.h
 #define __SYCL_PROPERTY_SET_SYCL_IMPLICIT_LOCAL_ARG "SYCL/implicit local arg"
+/// PropertySetRegistry::SYCL_REGISTERED_KERNELS defined in PropertySetIO.h
+#define __SYCL_PROPERTY_SET_SYCL_REGISTERED_KERNELS "SYCL/registered kernels"
 
 /// Program metadata tags recognized by the PI backends. For kernels the tag
 /// must appear after the kernel name.

sycl/source/detail/device_binary_image.cpp

@@ -195,6 +195,7 @@ void RTDeviceBinaryImage::init(sycl_device_binary Bin) {
   DeviceRequirements.init(Bin, __SYCL_PROPERTY_SET_SYCL_DEVICE_REQUIREMENTS);
   HostPipes.init(Bin, __SYCL_PROPERTY_SET_SYCL_HOST_PIPES);
   VirtualFunctions.init(Bin, __SYCL_PROPERTY_SET_SYCL_VIRTUAL_FUNCTIONS);
+  RegisteredKernels.init(Bin, __SYCL_PROPERTY_SET_SYCL_REGISTERED_KERNELS);
 
   ImageId = ImageCounter++;
 }

sycl/source/detail/device_binary_image.hpp

@@ -232,6 +232,9 @@ class RTDeviceBinaryImage {
   const PropertyRange &getHostPipes() const { return HostPipes; }
   const PropertyRange &getVirtualFunctions() const { return VirtualFunctions; }
   const PropertyRange &getImplicitLocalArg() const { return ImplicitLocalArg; }
+  const PropertyRange &getRegisteredKernels() const {
+    return RegisteredKernels;
+  }
 
   std::uintptr_t getImageID() const {
     assert(Bin && "Image ID is not available without a binary image.");
@@ -258,6 +261,7 @@ class RTDeviceBinaryImage {
   RTDeviceBinaryImage::PropertyRange HostPipes;
   RTDeviceBinaryImage::PropertyRange VirtualFunctions;
   RTDeviceBinaryImage::PropertyRange ImplicitLocalArg;
+  RTDeviceBinaryImage::PropertyRange RegisteredKernels;
 
   std::vector<ur_program_metadata_t> ProgramMetadataUR;
 

sycl/source/detail/jit_compiler.cpp

@@ -1249,24 +1249,11 @@ std::vector<uint8_t> jit_compiler::encodeReqdWorkGroupSize(
 sycl_device_binaries jit_compiler::compileSYCL(
     const std::string &CompilationID, const std::string &SYCLSource,
     const std::vector<std::pair<std::string, std::string>> &IncludePairs,
-    const std::vector<std::string> &UserArgs, std::string *LogPtr,
-    const std::vector<std::string> &RegisteredKernelNames) {
-
-  // RegisteredKernelNames may contain template specializations, so we just put
-  // them in main() which ensures they are instantiated.
-  std::ostringstream ss;
-  ss << SYCLSource << '\n';
-  ss << "int main() {\n";
-  for (const std::string &KernelName : RegisteredKernelNames) {
-    ss << "  (void)" << KernelName << ";\n";
-  }
-  ss << "  return 0;\n}\n" << std::endl;
-
-  std::string FinalSource = ss.str();
+    const std::vector<std::string> &UserArgs, std::string *LogPtr) {
 
   std::string SYCLFileName = CompilationID + ".cpp";
   ::jit_compiler::InMemoryFile SourceFile{SYCLFileName.c_str(),
-                                          FinalSource.c_str()};
+                                          SYCLSource.c_str()};
 
   std::vector<::jit_compiler::InMemoryFile> IncludeFilesView;
   IncludeFilesView.reserve(IncludePairs.size());

sycl/source/detail/jit_compiler.hpp

@@ -52,8 +52,7 @@ class jit_compiler {
   sycl_device_binaries compileSYCL(
       const std::string &CompilationID, const std::string &SYCLSource,
       const std::vector<std::pair<std::string, std::string>> &IncludePairs,
-      const std::vector<std::string> &UserArgs, std::string *LogPtr,
-      const std::vector<std::string> &RegisteredKernelNames);
+      const std::vector<std::string> &UserArgs, std::string *LogPtr);
 
   void destroyDeviceBinaries(sycl_device_binaries Binaries);
 

sycl/source/detail/kernel_bundle_impl.hpp

@@ -378,11 +378,12 @@ class kernel_bundle_impl {
 
   // oneapi_ext_kernel_compiler
   // 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,
-                     sycl_device_binaries Binaries, std::string Pfx,
-                     syclex::source_language Lang)
+  kernel_bundle_impl(
+      context Ctx, std::vector<device> Devs,
+      const std::vector<kernel_id> &KernelIDs,
+      std::unordered_map<std::string, std::string> &&MangledKernelNames,
+      sycl_device_binaries Binaries, std::string &&Prefix,
+      syclex::source_language Lang)
       : kernel_bundle_impl(std::move(Ctx), std::move(Devs), KernelIDs,
                            bundle_state::executable) {
     assert(Lang == syclex::source_language::sycl_jit);
@@ -392,9 +393,9 @@ class kernel_bundle_impl {
     // loaded via the program manager have `kernel_id`s, they can't be looked up
     // from the (unprefixed) kernel name.
     MIsInterop = true;
-    MKernelNames = std::move(KNames);
+    MMangledKernelNames = std::move(MangledKernelNames);
     MDeviceBinaries = Binaries;
-    MPrefix = std::move(Pfx);
+    MPrefix = std::move(Prefix);
     MLanguage = Lang;
   }
 
@@ -501,15 +502,35 @@ class kernel_bundle_impl {
       // TODO: Support persistent caching.
 
       const std::string &SourceStr = std::get<std::string>(MSource);
+      std::ostringstream SourceExt;
+      if (!RegisteredKernelNames.empty()) {
+        SourceExt << SourceStr << '\n';
+
+        auto EmitEntry =
+            [&SourceExt](const std::string &Name) -> std::ostringstream & {
+          SourceExt << "  {\"" << Name << "\", " << Name << "}";
+          return SourceExt;
+        };
+
+        SourceExt << "[[__sycl_detail__::__registered_kernels__(\n";
+        for (auto It = RegisteredKernelNames.begin(),
+                  SecondToLast = RegisteredKernelNames.end() - 1;
+             It != SecondToLast; ++It) {
+          EmitEntry(*It) << ",\n";
+        }
+        EmitEntry(RegisteredKernelNames.back()) << "\n";
+        SourceExt << ")]];\n";
+      }
+
       auto [Binaries, CompilationID] = syclex::detail::SYCL_JIT_to_SPIRV(
-          SourceStr, MIncludePairs, BuildOptions, LogPtr,
-          RegisteredKernelNames);
+          RegisteredKernelNames.empty() ? SourceStr : SourceExt.str(),
+          MIncludePairs, BuildOptions, LogPtr);
 
       auto &PM = detail::ProgramManager::getInstance();
       PM.addImages(Binaries);
 
       std::vector<kernel_id> KernelIDs;
-      std::vector<std::string> KernelNames;
+      std::unordered_map<std::string, std::string> MangledKernelNames;
       // `jit_compiler::compileSYCL(..)` uses `CompilationID + '$'` as prefix
       // for offload entry names.
       std::string Prefix = CompilationID + '$';
@@ -518,13 +539,38 @@ class kernel_bundle_impl {
         if (KernelName.find(Prefix) == 0) {
           KernelIDs.push_back(KernelID);
           KernelName.remove_prefix(Prefix.length());
-          KernelNames.emplace_back(KernelName);
+          static constexpr std::string_view SYCLKernelMarker{"__sycl_kernel_"};
+          if (KernelName.find(SYCLKernelMarker) == 0) {
+            // extern "C" declaration, register kernel without the marker.
+            std::string_view KernelNameWithoutMarker{KernelName};
+            KernelNameWithoutMarker.remove_prefix(SYCLKernelMarker.length());
+            MangledKernelNames.emplace(KernelNameWithoutMarker, KernelName);
+          } else {
+            // The marker is baked into the mangling, and we cannot easily
+            // adjust it. Register an identity mapping as an escape hatch.
+            // Users shall use `registered_kernel_names` instead, as there's
+            // practically no way to guess the mangled name.
+            MangledKernelNames.emplace(KernelName, KernelName);
+          }
         }
       }
 
-      return std::make_shared<kernel_bundle_impl>(MContext, MDevices, KernelIDs,
-                                                  KernelNames, Binaries, Prefix,
-                                                  MLanguage);
+      // Apply frontend information.
+      for (const auto *RawImg : PM.getRawDeviceImages(KernelIDs)) {
+        for (const sycl_device_binary_property &RKProp :
+             RawImg->getRegisteredKernels()) {
+
+          auto BA = DeviceBinaryProperty(RKProp).asByteArray();
+          auto MangledNameLen = BA.consume<uint64_t>() / 8 /*bits in a byte*/;
+          std::string_view MangledName{
+              reinterpret_cast<const char *>(BA.begin()), MangledNameLen};
+          MangledKernelNames.emplace(RKProp->Name, MangledName);
+        }
+      }
+
+      return std::make_shared<kernel_bundle_impl>(
+          MContext, MDevices, KernelIDs, std::move(MangledKernelNames),
+          Binaries, std::move(Prefix), MLanguage);
     }
 
     ur_program_handle_t UrProgram = nullptr;
@@ -642,6 +688,9 @@ class kernel_bundle_impl {
   }
 
   bool ext_oneapi_has_kernel(const std::string &Name) {
+    if (MLanguage == syclex::source_language::sycl_jit) {
+      return MMangledKernelNames.count(Name);
+    }
     auto it = std::find(MKernelNames.begin(), MKernelNames.end(),
                         adjust_kernel_name(Name, MLanguage));
     return it != MKernelNames.end();
@@ -650,21 +699,25 @@ class kernel_bundle_impl {
   kernel
   ext_oneapi_get_kernel(const std::string &Name,
                         const std::shared_ptr<kernel_bundle_impl> &Self) {
-    if (MKernelNames.empty())
-      throw sycl::exception(make_error_code(errc::invalid),
-                            "'ext_oneapi_get_kernel' is only available in "
-                            "kernel_bundles successfully built from "
-                            "kernel_bundle<bundle_state:ext_oneapi_source>.");
+    if (MLanguage == syclex::source_language::sycl_jit) {
+      if (MMangledKernelNames.empty()) {
+        throw sycl::exception(
+            make_error_code(errc::invalid),
+            "'ext_oneapi_get_kernel' is only available in kernel_bundles "
+            "successfully built from "
+            "kernel_bundle<bundle_state::ext_oneapi_source>.");
+      }
 
-    std::string AdjustedName = adjust_kernel_name(Name, MLanguage);
-    if (!ext_oneapi_has_kernel(Name))
-      throw sycl::exception(make_error_code(errc::invalid),
-                            "kernel '" + AdjustedName +
-                                "' not found in kernel_bundle");
+      auto It = MMangledKernelNames.find(Name);
+      if (It == MMangledKernelNames.end()) {
+        throw sycl::exception(make_error_code(errc::invalid),
+                              "kernel '" + Name +
+                                  "' not found in kernel_bundle");
+      }
 
-    if (MLanguage == syclex::source_language::sycl_jit) {
+      const std::string &MangledName = It->second;
       auto &PM = ProgramManager::getInstance();
-      auto KID = PM.getSYCLKernelID(MPrefix + AdjustedName);
+      auto KID = PM.getSYCLKernelID(MPrefix + MangledName);
 
       for (const auto &DevImgWithDeps : MDeviceImages) {
         const auto &DevImg = DevImgWithDeps.getMain();
@@ -674,7 +727,7 @@ class kernel_bundle_impl {
         const auto &DevImgImpl = getSyclObjImpl(DevImg);
         auto UrProgram = DevImgImpl->get_ur_program_ref();
         auto [UrKernel, CacheMutex, ArgMask] =
-            PM.getOrCreateKernel(MContext, AdjustedName,
+            PM.getOrCreateKernel(MContext, MangledName,
                                  /*PropList=*/{}, UrProgram);
         auto KernelImpl = std::make_shared<kernel_impl>(
             UrKernel, getSyclObjImpl(MContext), DevImgImpl, Self, ArgMask,
@@ -685,6 +738,18 @@ class kernel_bundle_impl {
       assert(false && "Malformed RTC kernel bundle");
     }
 
+    if (MKernelNames.empty())
+      throw sycl::exception(make_error_code(errc::invalid),
+                            "'ext_oneapi_get_kernel' is only available in "
+                            "kernel_bundles successfully built from "
+                            "kernel_bundle<bundle_state:ext_oneapi_source>.");
+
+    std::string AdjustedName = adjust_kernel_name(Name, MLanguage);
+    if (!ext_oneapi_has_kernel(Name))
+      throw sycl::exception(make_error_code(errc::invalid),
+                            "kernel '" + AdjustedName +
+                                "' not found in kernel_bundle");
+
     assert(MDeviceImages.size() > 0);
     const std::shared_ptr<detail::device_image_impl> &DeviceImageImpl =
         detail::getSyclObjImpl(MDeviceImages[0].getMain());
@@ -877,12 +942,11 @@ class kernel_bundle_impl {
   }
 
   bool is_specialization_constant_set(const char *SpecName) const noexcept {
-    bool SetInDevImg =
-        std::any_of(begin(), end(),
-                    [SpecName](const device_image_plain &DeviceImage) {
-                      return getSyclObjImpl(DeviceImage)
-                          ->is_specialization_constant_set(SpecName);
-                    });
+    bool SetInDevImg = std::any_of(
+        begin(), end(), [SpecName](const device_image_plain &DeviceImage) {
+          return getSyclObjImpl(DeviceImage)
+              ->is_specialization_constant_set(SpecName);
+        });
     return SetInDevImg || MSpecConstValues.count(std::string{SpecName}) != 0;
   }
 
@@ -973,6 +1037,7 @@ class kernel_bundle_impl {
   const std::variant<std::string, std::vector<std::byte>> MSource;
   // only kernel_bundles created from source have KernelNames member.
   std::vector<std::string> MKernelNames;
+  std::unordered_map<std::string, std::string> MMangledKernelNames;
   sycl_device_binaries MDeviceBinaries = nullptr;
   std::string MPrefix;
   include_pairs_t MIncludePairs;

sycl/source/detail/kernel_compiler/kernel_compiler_sycl.cpp

@@ -303,19 +303,17 @@ bool SYCL_JIT_Compilation_Available() {
 #endif
 }
 
-std::pair<sycl_device_binaries, std::string> SYCL_JIT_to_SPIRV(
-    [[maybe_unused]] const std::string &SYCLSource,
-    [[maybe_unused]] include_pairs_t IncludePairs,
-    [[maybe_unused]] const std::vector<std::string> &UserArgs,
-    [[maybe_unused]] std::string *LogPtr,
-    [[maybe_unused]] const std::vector<std::string> &RegisteredKernelNames) {
+std::pair<sycl_device_binaries, std::string>
+SYCL_JIT_to_SPIRV([[maybe_unused]] const std::string &SYCLSource,
+                  [[maybe_unused]] include_pairs_t IncludePairs,
+                  [[maybe_unused]] const std::vector<std::string> &UserArgs,
+                  [[maybe_unused]] std::string *LogPtr) {
 #if SYCL_EXT_JIT_ENABLE
   static std::atomic_uintptr_t CompilationCounter;
   std::string CompilationID = "rtc_" + std::to_string(CompilationCounter++);
   sycl_device_binaries Binaries =
       sycl::detail::jit_compiler::get_instance().compileSYCL(
-          CompilationID, SYCLSource, IncludePairs, UserArgs, LogPtr,
-          RegisteredKernelNames);
+          CompilationID, SYCLSource, IncludePairs, UserArgs, LogPtr);
   return std::make_pair(Binaries, std::move(CompilationID));
 #else
   throw sycl::exception(sycl::errc::build,

sycl/source/detail/kernel_compiler/kernel_compiler_sycl.hpp

@@ -37,8 +37,8 @@ std::string userArgsAsString(const std::vector<std::string> &UserArguments);
 
 std::pair<sycl_device_binaries, std::string>
 SYCL_JIT_to_SPIRV(const std::string &Source, include_pairs_t IncludePairs,
-                  const std::vector<std::string> &UserArgs, std::string *LogPtr,
-                  const std::vector<std::string> &RegisteredKernelNames);
+                  const std::vector<std::string> &UserArgs,
+                  std::string *LogPtr);
 
 void SYCL_JIT_destroy(sycl_device_binaries Binaries);
 

sycl/test-e2e/KernelCompiler/kernel_compiler_sycl_jit.cpp

@@ -207,14 +207,13 @@ int test_build_and_run() {
   // directly.
   sycl::kernel k = kbExe2.ext_oneapi_get_kernel("ff_cp");
 
-  // The templated function name will have been mangled. Mapping from original
-  // name to mangled is not yet supported. So we cannot yet do this:
-  // sycl::kernel k2 = kbExe2.ext_oneapi_get_kernel("ff_templated<int>");
-
-  // Instead, we can TEMPORARILY use the mangled name. Once demangling is
-  // supported this might no longer work.
-  sycl::kernel k2 =
-      kbExe2.ext_oneapi_get_kernel("_Z26__sycl_kernel_ff_templatedIiEvPT_S1_");
+  // The templated function name will have been mangled.
+  sycl::kernel k2 = kbExe2.ext_oneapi_get_kernel("ff_templated<int>");
+
+  // We can also use the mangled name. This escape hatch might be removed in the
+  // future.
+  assert(
+      kbExe2.ext_oneapi_has_kernel("_Z26__sycl_kernel_ff_templatedIiEvPT_S1_"));
 
   // Test the kernels.
   test_1(q, k, 37 + 5);  // ff_cp seeds 37. AddEm will add 5 more.