[SYCL] Specialization constants: fix scope, allow redifinition and AOT.

- Specialization constants are now per-program, as the implemented spec
  requires.
- Program's specialization constant set is added to the program
  compilation cache key, it gets recompiled on new specialization
  constant combination.

Signed-off-by: Konstantin S Bobrovsky <[email protected]>

Author: kbobrovs

sycl/include/CL/sycl/detail/spec_constant_impl.hpp

@@ -9,38 +9,44 @@
 #pragma once
 
 #include <CL/sycl/detail/defines.hpp>
+#include <CL/sycl/detail/util.hpp>
+#include <CL/sycl/stl.hpp>
 
 #include <iostream>
+#include <map>
 
 __SYCL_INLINE_NAMESPACE(cl) {
 namespace sycl {
 namespace detail {
 
-// Represents a specialization constant in SYCL runtime.
+// Represents a specialization constant value in SYCL runtime.
 class spec_constant_impl {
 public:
-  spec_constant_impl(unsigned int ID) : ID(ID), Size(0), Bytes{0} {}
+  spec_constant_impl() = default;
 
-  spec_constant_impl(unsigned int ID, size_t Size, const void *Val) : ID(ID) {
-    set(Size, Val);
-  }
+  spec_constant_impl(size_t Size, const void *Val) { set(Size, Val); }
 
   void set(size_t Size, const void *Val);
 
-  unsigned int getID() const { return ID; }
   size_t getSize() const { return Size; }
   const unsigned char *getValuePtr() const { return Bytes; }
   bool isSet() const { return Size != 0; }
 
 private:
-  unsigned int ID; // specialization constant's ID (equals to SPIRV ID)
-  size_t Size;     // size of its value
+  size_t Size; // size of its value
   // TODO invent more flexible approach to support values of arbitrary type:
   unsigned char Bytes[8]; // memory to hold the value bytes
 };
 
 std::ostream &operator<<(std::ostream &Out, const spec_constant_impl &V);
 
+// Used to define specialization constant registry. Must be ordered map, since
+// the order of entries matters in stableSerializeSpecConstRegistry.
+using SpecConstRegistryT = std::map<string_class, spec_constant_impl>;
+
+void stableSerializeSpecConstRegistry(const SpecConstRegistryT &Reg,
+                                      SerializedObj &Dst);
+
 } // namespace detail
 } // namespace sycl
 } // __SYCL_INLINE_NAMESPACE(cl)

sycl/include/CL/sycl/detail/util.hpp

@@ -11,6 +11,7 @@
 #ifndef __SYCL_DEVICE_ONLY
 
 #include <CL/sycl/detail/defines.hpp>
+#include <CL/sycl/stl.hpp>
 
 #include <cstring>
 #include <mutex>
@@ -52,12 +53,7 @@ struct CmpCStr {
   }
 };
 
-// Interface to iterate via C strings.
-class CStringIterator {
-public:
-  // Get the next string. Returns next string's pointer or nullptr.
-  virtual const char *next() = 0;
-};
+using SerializedObj = sycl::vector_class<unsigned char>;
 
 } // namespace detail
 } // namespace sycl

sycl/source/detail/kernel_program_cache.hpp

@@ -8,9 +8,11 @@
 
 #pragma once
 
+#include <CL/sycl/detail/common.hpp>
 #include <CL/sycl/detail/locked.hpp>
 #include <CL/sycl/detail/os_util.hpp>
 #include <CL/sycl/detail/pi.hpp>
+#include <CL/sycl/detail/util.hpp>
 #include <detail/platform_impl.hpp>
 
 #include <atomic>
@@ -51,7 +53,8 @@ class KernelProgramCache {
   using PiProgramT = std::remove_pointer<RT::PiProgram>::type;
   using PiProgramPtrT = std::atomic<PiProgramT *>;
   using ProgramWithBuildStateT = BuildResult<PiProgramT>;
-  using ProgramCacheT = std::map<OSModuleHandle, ProgramWithBuildStateT>;
+  using ProgramCacheKeyT = std::pair<SerializedObj, KernelSetId>;
+  using ProgramCacheT = std::map<ProgramCacheKeyT, ProgramWithBuildStateT>;
   using ContextPtr = context_impl *;
 
   using PiKernelT = std::remove_pointer<RT::PiKernel>::type;

sycl/source/detail/program_impl.cpp

@@ -231,7 +231,7 @@ void program_impl::build_with_kernel_name(string_class KernelName,
     if (is_cacheable_with_options(BuildOptions)) {
       MProgramAndKernelCachingAllowed = true;
       MProgram = ProgramManager::getInstance().getBuiltPIProgram(
-          Module, get_context(), KernelName);
+          Module, get_context(), KernelName, this);
       const detail::plugin &Plugin = getPlugin();
       Plugin.call<PiApiKind::piProgramRetain>(MProgram);
     } else {
@@ -332,7 +332,7 @@ void program_impl::compile(const string_class &Options) {
   check_device_feature_support<info::device::is_compiler_available>(MDevices);
   vector_class<RT::PiDevice> Devices(get_pi_devices());
   const detail::plugin &Plugin = getPlugin();
-  ProgramManager::getInstance().flushSpecConstants(MProgram, *MContext);
+  ProgramManager::getInstance().flushSpecConstants(*this);
   RT::PiResult Err = Plugin.call_nocheck<PiApiKind::piProgramCompile>(
       MProgram, Devices.size(), Devices.data(), Options.c_str(), 0, nullptr,
       nullptr, nullptr, nullptr);
@@ -351,7 +351,7 @@ void program_impl::build(const string_class &Options) {
   check_device_feature_support<info::device::is_compiler_available>(MDevices);
   vector_class<RT::PiDevice> Devices(get_pi_devices());
   const detail::plugin &Plugin = getPlugin();
-  ProgramManager::getInstance().flushSpecConstants(MProgram, *MContext);
+  ProgramManager::getInstance().flushSpecConstants(*this);
   RT::PiResult Err = Plugin.call_nocheck<PiApiKind::piProgramBuild>(
       MProgram, Devices.size(), Devices.data(), Options.c_str(), nullptr,
       nullptr);
@@ -398,7 +398,7 @@ RT::PiKernel program_impl::get_pi_kernel(const string_class &KernelName) const {
 
   if (is_cacheable()) {
     Kernel = ProgramManager::getInstance().getOrCreateKernel(
-        MProgramModuleHandle, get_context(), KernelName);
+        MProgramModuleHandle, get_context(), KernelName, this);
   } else {
     const detail::plugin &Plugin = getPlugin();
     RT::PiResult Err = Plugin.call_nocheck<PiApiKind::piKernelCreate>(
@@ -470,11 +470,38 @@ vector_class<device> program_impl::get_info<info::program::devices>() const {
 
 void program_impl::set_spec_constant_impl(const char *Name, const void *ValAddr,
                                           size_t ValSize) {
-  spec_constant_impl &SC =
-      ProgramManager::getInstance().resolveSpecConstant(this, Name);
+  // Reuse cached programs lock as opposed to introducing a new lock.
+  auto LockGuard = MContext->getKernelProgramCache().acquireCachedPrograms();
+  spec_constant_impl &SC = SpecConstRegistry[Name];
   SC.set(ValSize, ValAddr);
 }
 
+void program_impl::flush_spec_constants(const RTDeviceBinaryImage &Img,
+                                        RT::PiProgram NativePrg) const {
+  // iterate via all specialization constants the program's image depends on,
+  // and set each to current runtime value (if any)
+  const pi::DeviceBinaryImage::PropertyRange &SCRange = Img.getSpecConstants();
+  ContextImplPtr Ctx = getSyclObjImpl(get_context());
+  using SCItTy = pi::DeviceBinaryImage::PropertyRange::ConstIterator;
+
+  auto LockGuard = Ctx->getKernelProgramCache().acquireCachedPrograms();
+
+  for (SCItTy SCIt = SCRange.begin(); SCIt != SCRange.end(); SCIt++) {
+    const char *SCName = (*SCIt)->Name;
+    auto SCEntry = SpecConstRegistry.find(SCName);
+    if (SCEntry == SpecConstRegistry.end())
+      // spec constant has not been set in user code - SPIRV will use default
+      continue;
+    const spec_constant_impl &SC = SCEntry->second;
+    assert(SC.isSet() && "uninitialized spec constant");
+    pi_device_binary_property SCProp = *SCIt;
+    pi_uint32 ID = pi::DeviceBinaryProperty(SCProp).asUint32();
+    NativePrg = NativePrg ? NativePrg : pi::cast<pi::PiProgram>(get());
+    Ctx->getPlugin().call<PiApiKind::piextProgramSetSpecializationConstant>(
+        NativePrg, ID, SC.getSize(), SC.getValuePtr());
+  }
+}
+
 } // namespace detail
 } // namespace sycl
 } // __SYCL_INLINE_NAMESPACE(cl)

sycl/source/detail/program_impl.hpp

@@ -10,6 +10,7 @@
 #include <CL/sycl/context.hpp>
 #include <CL/sycl/detail/common_info.hpp>
 #include <CL/sycl/detail/kernel_desc.hpp>
+#include <CL/sycl/detail/spec_constant_impl.hpp>
 #include <CL/sycl/device.hpp>
 #include <CL/sycl/program.hpp>
 #include <CL/sycl/stl.hpp>
@@ -292,13 +293,28 @@ class program_impl {
   void set_spec_constant_impl(const char *Name, const void *ValAddr,
                               size_t ValSize);
 
+  /// Takes current values of specialization constants and "injects" them into
+  /// the underlying native program program via specialization constant
+  /// managemment PI APIs. The native program passed as non-null argument
+  /// overrides the MProgram native program field.
+  /// \param Img device binary image corresponding to this program, used to
+  ///        resolve spec constant name to SPIRV integer ID
+  /// \param NativePrg if not null, used as the flush target, otherwise MProgram
+  ///        is used
+  void flush_spec_constants(const RTDeviceBinaryImage &Img,
+                            RT::PiProgram NativePrg = nullptr) const;
+
   /// Returns the OS module handle this program belongs to. A program belongs to
   /// an OS module if it was built from device image(s) belonging to that
   /// module.
   /// TODO Some programs can be linked from images belonging to different
   ///      modules. May need a special fake handle for the resulting program.
   OSModuleHandle getOSModuleHandle() const { return MProgramModuleHandle; }
 
+  void stableSerializeSpecConstRegistry(SerializedObj &Dst) const {
+    detail::stableSerializeSpecConstRegistry(SpecConstRegistry, Dst);
+  }
+
 private:
   // Deligating Constructor used in Implementation.
   program_impl(ContextImplPtr Context, pi_native_handle InteropProgram,
@@ -390,6 +406,12 @@ class program_impl {
   string_class MBuildOptions;
   OSModuleHandle MProgramModuleHandle = OSUtil::ExeModuleHandle;
 
+  // Keeps specialization constant map for this program. Spec constant name
+  // resolution to actual SPIRV integer ID happens at build time, where the
+  // device binary image is available. Access is guarded by this context's
+  // program cache lock.
+  SpecConstRegistryT SpecConstRegistry;
+
   /// Only allow kernel caching for programs constructed with context only (or
   /// device list and context) and built with build_with_kernel_type with
   /// default build options