[SYCL] Add support for set(get)_specialization_constant

sycl/include/CL/sycl/detail/kernel_desc.hpp

@@ -17,6 +17,13 @@ __SYCL_INLINE_NAMESPACE(cl) {
 namespace sycl {
 namespace detail {
 
+// This guard is needed because the libsycl.so can be compiled with C++ <=14
+// while the code requires C++17. This code is not supposed to be used by the
+// libsycl.so so it should not be a problem.
+#if __cplusplus > 201402L
+template <auto &S> struct specialization_id_name_generator {};
+#endif
+
 #ifndef __SYCL_DEVICE_ONLY__
 #define _Bool bool
 #endif
@@ -49,6 +56,14 @@ template <class Name> struct SpecConstantInfo {
   static constexpr const char *getName() { return ""; }
 };
 
+#if __cplusplus >= 201703L
+// Translates SYCL 2020 specialization constant type to its name.
+template <auto &SpecName> const char *get_spec_constant_symbolic_ID() {
+  return __builtin_unique_stable_name(
+      specialization_id_name_generator<SpecName>);
+}
+#endif
+
 #ifndef __SYCL_UNNAMED_LAMBDA__
 template <class KernelNameType> struct KernelInfo {
   static constexpr unsigned getNumParams() { return 0; }

sycl/include/CL/sycl/detail/sycl_fe_intrins.hpp

@@ -37,13 +37,12 @@ SYCL_EXTERNAL T __sycl_getCompositeSpecConstantValue(const char *ID);
 // are not available.
 template <typename T>
 SYCL_EXTERNAL T __sycl_getScalar2020SpecConstantValue(const char *SymbolicID,
-                                                      void *DefaultValue,
-                                                      void *RTBuffer);
+                                                      const void *DefaultValue,
+                                                      const void *RTBuffer);
 
 template <typename T>
-SYCL_EXTERNAL T __sycl_getComposite2020SpecConstantValue(const char *SymbolicID,
-                                                         void *DefaultValue,
-                                                         void *RTBuffer);
+SYCL_EXTERNAL T __sycl_getComposite2020SpecConstantValue(
+    const char *SymbolicID, const void *DefaultValue, const void *RTBuffer);
 
 // Request a fixed-size allocation in local address space at kernel scope.
 extern "C" SYCL_EXTERNAL __attribute__((opencl_local)) std::uint8_t *

sycl/include/CL/sycl/handler.hpp

@@ -1048,7 +1048,7 @@ class __SYCL_EXPORT handler {
   }
 
   template <auto &SpecName>
-  typename std::remove_reference_t<decltype(SpecName)>::type
+  typename std::remove_reference_t<decltype(SpecName)>::value_type
   get_specialization_constant() const {
 
     std::shared_ptr<detail::kernel_bundle_impl> KernelBundleImplPtr =

sycl/include/CL/sycl/kernel_bundle.hpp

@@ -168,6 +168,16 @@ class __SYCL_EXPORT kernel_bundle_plain {
   // \returns an iterator to the last device image kernel_bundle contains
   const device_image_plain *end() const;
 
+  bool has_specialization_constant_impl(const char *SpecName) const noexcept;
+
+  void set_specialization_constant_impl(const char *SpecName,
+                                        void *Value) noexcept;
+
+  void get_specialization_constant_impl(const char *SpecName, void *Value) const
+      noexcept;
+
+  bool is_specialization_constant_set(const char *SpecName) const noexcept;
+
   detail::KernelBundleImplPtr impl;
 };
 
@@ -247,9 +257,8 @@ class kernel_bundle : public detail::kernel_bundle_plain {
   /// \returns true if any device image in the kernel_bundle uses specialization
   /// constant whose address is SpecName
   template <auto &SpecName> bool has_specialization_constant() const noexcept {
-    throw sycl::runtime_error(
-        "kernel_bundle::has_specialization_constant is not implemented yet",
-        PI_INVALID_OPERATION);
+    const char *SpecSymName = detail::get_spec_constant_symbolic_ID<SpecName>();
+    return has_specialization_constant_impl(SpecSymName);
   }
 
   /// Sets the value of the specialization constant whose address is SpecName
@@ -259,20 +268,27 @@ class kernel_bundle : public detail::kernel_bundle_plain {
             typename = detail::enable_if_t<_State == bundle_state::input>>
   void set_specialization_constant(
       typename std::remove_reference_t<decltype(SpecName)>::value_type Value) {
-    (void)Value;
-    throw sycl::runtime_error(
-        "kernel_bundle::set_specialization_constant is not implemented yet",
-        PI_INVALID_OPERATION);
+    const char *SpecSymName = detail::get_spec_constant_symbolic_ID<SpecName>();
+    set_specialization_constant_impl(SpecSymName, &Value);
   }
 
-  /// The value of the specialization constant whose address is SpecName for
-  /// this kernel bundle.
+  /// \returns the value of the specialization constant whose address is
+  /// SpecName for this kernel bundle.
   template <auto &SpecName>
   typename std::remove_reference_t<decltype(SpecName)>::value_type
   get_specialization_constant() const {
-    throw sycl::runtime_error(
-        "kernel_bundle::get_specialization_constant is not implemented yet",
-        PI_INVALID_OPERATION);
+    const char *SpecSymName = detail::get_spec_constant_symbolic_ID<SpecName>();
+    if (!is_specialization_constant_set(SpecSymName))
+      return SpecName.getDefaultValue();
+
+    using SCType =
+        typename std::remove_reference_t<decltype(SpecName)>::value_type;
+
+    std::array<char *, sizeof(SCType)> RetValue;
+
+    get_specialization_constant_impl(SpecSymName, RetValue.data());
+
+    return *reinterpret_cast<SCType *>(RetValue.data());
   }
 #endif
 

sycl/include/CL/sycl/kernel_handler.hpp

@@ -8,27 +8,22 @@
 
 #pragma once
 
-__SYCL_INLINE_NAMESPACE(cl) {
-namespace sycl {
-namespace detail {
-
-// This guard is needed because the libsycl.so can compiled with C++ <=14
-// while the code requires C++17. This code is not supposed to be used by the
-// libsycl.so so it should not be a problem.
-#if __cplusplus > 201402L
-template <auto &S> struct specialization_id_name_generator {};
-#endif
+#include <CL/sycl/detail/kernel_desc.hpp>
+#include <CL/sycl/exception.hpp>
 
-} // namespace detail
+#include <type_traits>
 
+__SYCL_INLINE_NAMESPACE(cl) {
+namespace sycl {
 /// Reading the value of a specialization constant
 ///
 /// \ingroup sycl_api
 class kernel_handler {
 public:
 #if __cplusplus > 201402L
   template <auto &S>
-  typename std::remove_reference_t<decltype(S)> get_specialization_constant() {
+  typename std::remove_reference_t<decltype(S)>::value_type
+  get_specialization_constant() {
 #ifdef __SYCL_DEVICE_ONLY__
     return getSpecializationConstantOnDevice<S>();
 #else
@@ -48,16 +43,20 @@ class kernel_handler {
   }
 
 #ifdef __SYCL_DEVICE_ONLY__
-  template <auto &S, typename T = std::remove_reference_t<decltype(S)>,
-            std::enable_if_t<std::is_fundamental_v<T>> * = nullptr>
+  template <
+      auto &S,
+      typename T = typename std::remove_reference_t<decltype(S)>::value_type,
+      std::enable_if_t<std::is_fundamental_v<T>> * = nullptr>
   T getSpecializationConstantOnDevice() {
     const char *SymbolicID = __builtin_unique_stable_name(
         detail::specialization_id_name_generator<S>);
     return __sycl_getScalar2020SpecConstantValue<T>(
         SymbolicID, &S, MSpecializationConstantsBuffer);
   }
-  template <auto &S, typename T = std::remove_reference_t<decltype(S)>,
-            std::enable_if_t<std::is_compound_v<T>> * = nullptr>
+  template <
+      auto &S,
+      typename T = typename std::remove_reference_t<decltype(S)>::value_type,
+      std::enable_if_t<std::is_compound_v<T>> * = nullptr>
   T getSpecializationConstantOnDevice() {
     const char *SymbolicID = __builtin_unique_stable_name(
         detail::specialization_id_name_generator<S>);

sycl/include/CL/sycl/specialization_id.hpp

@@ -28,6 +28,9 @@ template <typename T> class specialization_id {
   specialization_id &operator=(specialization_id &&rhs) = delete;
 
 private:
+  template <bundle_state State> friend class kernel_bundle;
+  T getDefaultValue() const noexcept { return MDefaultValue; }
+
   T MDefaultValue;
 };
 

sycl/source/detail/device_image_impl.hpp

@@ -39,7 +39,9 @@ class device_image_impl {
                     std::vector<kernel_id> KernelIDs, RT::PiProgram Program)
       : MBinImage(BinImage), MContext(std::move(Context)),
         MDevices(std::move(Devices)), MState(State), MProgram(Program),
-        MKernelIDs(std::move(KernelIDs)) {}
+        MKernelIDs(std::move(KernelIDs)) {
+    updateSpecConstSymMap();
+  }
 
   bool has_kernel(const kernel_id &KernelIDCand) const noexcept {
     return std::binary_search(MKernelIDs.begin(), MKernelIDs.end(),
@@ -60,7 +62,11 @@ class device_image_impl {
   }
 
   bool has_specialization_constants() const noexcept {
-    return !MSpecConstsBlob.empty();
+    // Lock the mutex to prevent when one thread in the middle of writing a
+    // new value while another thread is reading the value to pass it to
+    // JIT compiler.
+    const std::lock_guard<std::mutex> SpecConstLock(MSpecConstAccessMtx);
+    return !MSpecConstSymMap.empty();
   }
 
   bool all_specialization_constant_native() const noexcept {
@@ -72,45 +78,69 @@ class device_image_impl {
   // for this spec const should be.
   struct SpecConstDescT {
     unsigned int ID = 0;
-    unsigned int Offset = 0;
+    unsigned int CompositeOffset = 0;
+    unsigned int Size = 0;
+    unsigned int BlobOffset = 0;
     bool IsSet = false;
   };
 
-  bool has_specialization_constant(unsigned int SpecID) const noexcept {
-    return std::any_of(MSpecConstDescs.begin(), MSpecConstDescs.end(),
-                       [SpecID](const SpecConstDescT &SpecConstDesc) {
-                         return SpecConstDesc.ID == SpecID;
-                       });
-  }
-
-  void set_specialization_constant_raw_value(unsigned int SpecID,
-                                             const void *Value,
-                                             size_t ValueSize) noexcept {
-    for (const SpecConstDescT &SpecConstDesc : MSpecConstDescs)
-      if (SpecConstDesc.ID == SpecID) {
-        // Lock the mutex to prevent when one thread in the middle of writing a
-        // new value while another thread is reading the value to pass it to
-        // JIT compiler.
-        const std::lock_guard<std::mutex> SpecConstLock(MSpecConstAccessMtx);
-        std::memcpy(MSpecConstsBlob.data() + SpecConstDesc.Offset, Value,
-                    ValueSize);
-        return;
-      }
+  bool has_specialization_constant(const char *SpecName) const noexcept {
+    // Lock the mutex to prevent when one thread in the middle of writing a
+    // new value while another thread is reading the value to pass it to
+    // JIT compiler.
+    const std::lock_guard<std::mutex> SpecConstLock(MSpecConstAccessMtx);
+    return MSpecConstSymMap.count(SpecName) != 0;
   }
 
-  void get_specialization_constant_raw_value(unsigned int SpecID,
-                                             void *ValueRet,
-                                             size_t ValueSize) const noexcept {
-    for (const SpecConstDescT &SpecConstDesc : MSpecConstDescs)
-      if (SpecConstDesc.ID == SpecID) {
-        // Lock the mutex to prevent when one thread in the middle of writing a
-        // new value while another thread is reading the value to pass it to
-        // JIT compiler.
-        const std::lock_guard<std::mutex> SpecConstLock(MSpecConstAccessMtx);
-        std::memcpy(ValueRet, MSpecConstsBlob.data() + SpecConstDesc.Offset,
-                    ValueSize);
-        return;
-      }
+  void set_specialization_constant_raw_value(const char *SpecName,
+                                             const void *Value) noexcept {
+    // Lock the mutex to prevent when one thread in the middle of writing a
+    // new value while another thread is reading the value to pass it to
+    // JIT compiler.
+    const std::lock_guard<std::mutex> SpecConstLock(MSpecConstAccessMtx);
+
+    if (MSpecConstSymMap.count(std::string{SpecName}) == 0)
+      return;
+
+    std::vector<SpecConstDescT> &Descs =
+        MSpecConstSymMap[std::string{SpecName}];
+    for (SpecConstDescT &Desc : Descs) {
+      Desc.IsSet = true;
+      std::memcpy(MSpecConstsBlob.data() + Desc.BlobOffset,
+                  static_cast<const char *>(Value) + Desc.CompositeOffset,
+                  Desc.Size);
+    }
+  }
+
+  void get_specialization_constant_raw_value(const char *SpecName,
+                                             void *ValueRet) const noexcept {
+    assert(is_specialization_constant_set(SpecName));
+    // Lock the mutex to prevent when one thread in the middle of writing a
+    // new value while another thread is reading the value to pass it to
+    // JIT compiler.
+    const std::lock_guard<std::mutex> SpecConstLock(MSpecConstAccessMtx);
+
+    // operator[] can't be used here, since it's not marked as const
+    const std::vector<SpecConstDescT> &Descs =
+        MSpecConstSymMap.at(std::string{SpecName});
+    for (const SpecConstDescT &Desc : Descs) {
+
+      std::memcpy(static_cast<char *>(ValueRet) + Desc.CompositeOffset,
+                  MSpecConstsBlob.data() + Desc.BlobOffset, Desc.Size);
+    }
+  }
+
+  bool is_specialization_constant_set(const char *SpecName) const noexcept {
+    // Lock the mutex to prevent when one thread in the middle of writing a
+    // new value while another thread is reading the value to pass it to
+    // JIT compiler.
+    const std::lock_guard<std::mutex> SpecConstLock(MSpecConstAccessMtx);
+    if (MSpecConstSymMap.count(std::string{SpecName}) == 0)
+      return false;
+
+    const std::vector<SpecConstDescT> &Descs =
+        MSpecConstSymMap.at(std::string{SpecName});
+    return Descs.front().IsSet;
   }
 
   bundle_state get_state() const noexcept { return MState; }
@@ -137,8 +167,13 @@ class device_image_impl {
     return MSpecConstsBlob;
   }
 
-  std::vector<SpecConstDescT> &get_spec_const_offsets_ref() noexcept {
-    return MSpecConstDescs;
+  const std::map<std::string, std::vector<SpecConstDescT>> &
+  get_spec_const_data_ref() const noexcept {
+    return MSpecConstSymMap;
+  }
+
+  std::mutex &get_spec_const_data_lock() noexcept {
+    return MSpecConstAccessMtx;
   }
 
   ~device_image_impl() {
@@ -150,6 +185,49 @@ class device_image_impl {
   }
 
 private:
+  void updateSpecConstSymMap() {
+    if (MBinImage) {
+      const pi::DeviceBinaryImage::PropertyRange &SCRange =
+          MBinImage->getSpecConstants();
+      using SCItTy = pi::DeviceBinaryImage::PropertyRange::ConstIterator;
+
+      // This variable is used to calculate spec constant value offset in a
+      // flat byte array.
+      unsigned BlobOffset = 0;
+      for (SCItTy SCIt : SCRange) {
+        const char *SCName = (*SCIt)->Name;
+
+        pi::ByteArray Descriptors =
+            pi::DeviceBinaryProperty(*SCIt).asByteArray();
+        assert(Descriptors.size() > 8 && "Unexpected property size");
+
+        // Expected layout is vector of 3-component tuples (flattened into a
+        // vector of scalars), where each tuple consists of: ID of a scalar spec
+        // constant, (which might be a member of the composite); offset, which
+        // is used to calculate location of scalar member within the composite
+        // or zero for scalar spec constants; size of a spec constant
+        constexpr size_t NumElements = 3;
+        assert(((Descriptors.size() - 8) / sizeof(std::uint32_t)) %
+                       NumElements ==
+                   0 &&
+               "unexpected layout of composite spec const descriptors");
+        auto *It = reinterpret_cast<const std::uint32_t *>(&Descriptors[8]);
+        auto *End = reinterpret_cast<const std::uint32_t *>(&Descriptors[0] +
+                                                            Descriptors.size());
+        while (It != End) {
+          // The map is not locked here because updateSpecConstSymMap() is only
+          // supposed to be called from c'tor.
+          MSpecConstSymMap[std::string{SCName}].push_back(
+              SpecConstDescT{/*ID*/ It[0], /*CompositeOffset*/ It[1],
+                             /*Size*/ It[2], BlobOffset});
+          BlobOffset += /*Size*/ It[2];
+          It += NumElements;
+        }
+      }
+      MSpecConstsBlob.resize(BlobOffset);
+    }
+  }
+
   const RTDeviceBinaryImage *MBinImage = nullptr;
   context MContext;
   std::vector<device> MDevices;
@@ -166,8 +244,9 @@ class device_image_impl {
   // Binary blob which can have values of all specialization constants in the
   // image
   std::vector<unsigned char> MSpecConstsBlob;
-  // Contains list of spec ID + their offsets in the MSpecConstsBlob
-  std::vector<SpecConstDescT> MSpecConstDescs;
+  // Contains map of spec const names to their descriptions + offsets in
+  // the MSpecConstsBlob
+  std::map<std::string, std::vector<SpecConstDescT>> MSpecConstSymMap;
 };
 
 } // namespace detail