[SYCL] Fix spec constants in object and executable kernel bundles (#3853)

Allow set specialization constants values to be retrieved from
kernel_bundles in object or executable states.

Author: Alexander Batashev

sycl/source/detail/device_image_impl.hpp

@@ -36,6 +36,18 @@ namespace detail {
 // of specialization constants for it
 class device_image_impl {
 public:
+  // The struct maps specialization ID to offset in the binary blob where value
+  // for this spec const should be.
+  struct SpecConstDescT {
+    unsigned int ID = 0;
+    unsigned int CompositeOffset = 0;
+    unsigned int Size = 0;
+    unsigned int BlobOffset = 0;
+    bool IsSet = false;
+  };
+
+  using SpecConstMapT = std::map<std::string, std::vector<SpecConstDescT>>;
+
   device_image_impl(const RTDeviceBinaryImage *BinImage, context Context,
                     std::vector<device> Devices, bundle_state State,
                     std::vector<kernel_id> KernelIDs, RT::PiProgram Program)
@@ -45,6 +57,16 @@ class device_image_impl {
     updateSpecConstSymMap();
   }
 
+  device_image_impl(const RTDeviceBinaryImage *BinImage, context Context,
+                    std::vector<device> Devices, bundle_state State,
+                    std::vector<kernel_id> KernelIDs, RT::PiProgram Program,
+                    const SpecConstMapT &SpecConstMap,
+                    const std::vector<unsigned char> &SpecConstsBlob)
+      : MBinImage(BinImage), MContext(std::move(Context)),
+        MDevices(std::move(Devices)), MState(State), MProgram(Program),
+        MKernelIDs(std::move(KernelIDs)), MSpecConstsBlob(SpecConstsBlob),
+        MSpecConstSymMap(SpecConstMap) {}
+
   bool has_kernel(const kernel_id &KernelIDCand) const noexcept {
     return std::binary_search(MKernelIDs.begin(), MKernelIDs.end(),
                               KernelIDCand, LessByNameComp{});
@@ -76,16 +98,6 @@ class device_image_impl {
     return false;
   }
 
-  // The struct maps specialization ID to offset in the binary blob where value
-  // for this spec const should be.
-  struct SpecConstDescT {
-    unsigned int ID = 0;
-    unsigned int CompositeOffset = 0;
-    unsigned int Size = 0;
-    unsigned int BlobOffset = 0;
-    bool IsSet = false;
-  };
-
   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
@@ -182,8 +194,7 @@ class device_image_impl {
     return MSpecConstsBuffer;
   }
 
-  const std::map<std::string, std::vector<SpecConstDescT>> &
-  get_spec_const_data_ref() const noexcept {
+  const SpecConstMapT &get_spec_const_data_ref() const noexcept {
     return MSpecConstSymMap;
   }
 

sycl/source/detail/kernel_bundle_impl.hpp

@@ -55,6 +55,8 @@ static bool checkAllDevicesHaveAspect(const std::vector<device> &Devices,
 // objects.
 class kernel_bundle_impl {
 
+  using SpecConstMapT = std::map<std::string, std::vector<unsigned char>>;
+
   void common_ctor_checks(bundle_state State) {
     const bool AllDevicesInTheContext =
         checkAllDevicesAreInContext(MDevices, MContext);
@@ -105,6 +107,8 @@ class kernel_bundle_impl {
                      bundle_state TargetState)
       : MContext(InputBundle.get_context()), MDevices(std::move(Devs)) {
 
+    MSpecConstValues = getSyclObjImpl(InputBundle)->get_spec_const_map_ref();
+
     const std::vector<device> &InputBundleDevices =
         getSyclObjImpl(InputBundle)->get_devices();
     const bool AllDevsAssociatedWithInputBundle =
@@ -207,6 +211,14 @@ class kernel_bundle_impl {
 
     MDeviceImages = detail::ProgramManager::getInstance().link(
         std::move(DeviceImages), MDevices, PropList);
+
+    for (const kernel_bundle<bundle_state::object> &Bundle : ObjectBundles) {
+      const KernelBundleImplPtr BundlePtr = getSyclObjImpl(Bundle);
+      for (const std::pair<const std::string, std::vector<unsigned char>>
+               &SpecConst : BundlePtr->MSpecConstValues) {
+        MSpecConstValues[SpecConst.first] = SpecConst.second;
+      }
+    }
   }
 
   kernel_bundle_impl(context Ctx, std::vector<device> Devs,
@@ -258,8 +270,6 @@ class kernel_bundle_impl {
 
     std::sort(MDeviceImages.begin(), MDeviceImages.end(),
               LessByHash<device_image_plain>{});
-    const auto DevImgIt =
-        std::unique(MDeviceImages.begin(), MDeviceImages.end());
 
     if (get_bundle_state() == bundle_state::input) {
       // Copy spec constants values from the device images to be removed.
@@ -285,6 +295,9 @@ class kernel_bundle_impl {
                     MergeSpecConstants);
     }
 
+    const auto DevImgIt =
+        std::unique(MDeviceImages.begin(), MDeviceImages.end());
+
     // Remove duplicate device images.
     MDeviceImages.erase(DevImgIt, MDeviceImages.end());
 
@@ -459,13 +472,17 @@ class kernel_bundle_impl {
                : detail::getSyclObjImpl(MDeviceImages[0])->get_state();
   }
 
+  const SpecConstMapT &get_spec_const_map_ref() const noexcept {
+    return MSpecConstValues;
+  }
+
 private:
   context MContext;
   std::vector<device> MDevices;
   std::vector<device_image_plain> MDeviceImages;
   // This map stores values for specialization constants, that are missing
   // from any device image.
-  std::map<std::string, std::vector<unsigned char>> MSpecConstValues;
+  SpecConstMapT MSpecConstValues;
 };
 
 } // namespace detail

sycl/source/detail/program_manager/program_manager.cpp

@@ -1405,7 +1405,9 @@ ProgramManager::compile(const device_image_plain &DeviceImage,
 
   DeviceImageImplPtr ObjectImpl = std::make_shared<detail::device_image_impl>(
       InputImpl->get_bin_image_ref(), InputImpl->get_context(), Devs,
-      bundle_state::object, InputImpl->get_kernel_ids_ref(), Prog);
+      bundle_state::object, InputImpl->get_kernel_ids_ref(), Prog,
+      InputImpl->get_spec_const_data_ref(),
+      InputImpl->get_spec_const_blob_ref());
 
   std::vector<pi_device> PIDevices;
   PIDevices.reserve(Devs.size());
@@ -1652,7 +1654,9 @@ device_image_plain ProgramManager::build(const device_image_plain &DeviceImage,
 
   DeviceImageImplPtr ExecImpl = std::make_shared<detail::device_image_impl>(
       InputImpl->get_bin_image_ref(), Context, Devs, bundle_state::executable,
-      InputImpl->get_kernel_ids_ref(), ResProgram);
+      InputImpl->get_kernel_ids_ref(), ResProgram,
+      InputImpl->get_spec_const_data_ref(),
+      InputImpl->get_spec_const_blob_ref());
 
   return createSyclObjFromImpl<device_image_plain>(ExecImpl);
 }

sycl/test/on-device/basic_tests/specialization_constants/host_apis.cpp

@@ -58,6 +58,8 @@ int main() {
   KernelBundle.set_specialization_constant<SpecConst2>(1);
   {
     auto ExecBundle = sycl::build(KernelBundle);
+    assert(ExecBundle.get_specialization_constant<SpecConst1>() == 1);
+    assert(ExecBundle.get_specialization_constant<SpecConst2>() == 1);
     sycl::buffer<int, 1> Buf{sycl::range{1}};
     Q.submit([&](sycl::handler &CGH) {
       CGH.use_kernel_bundle(ExecBundle);