[SYCL] Choose image with inlined default values if default value is set explicitly (#12626)

againull · web-flow · commit 73d34739bff5 · 2024-02-09T08:24:57.000-08:00
If -fsycl-add-default-spec-consts-image option is used then DPCPP
generates a device image where default values of specialization
constants are inlined (in addition to regular device image).

Currently we always choose regular device image if somebody calls
`set_specialization_constant ` API.
This PR improves this behavior: if `set_specialization_constant` sets
the value equal to default value then we can choose the device image
where default values of specialization constants are inlined.
diff --git a/sycl/include/sycl/kernel_bundle.hpp b/sycl/include/sycl/kernel_bundle.hpp
@@ -205,6 +205,8 @@ class __SYCL_EXPORT kernel_bundle_plain {
   void get_specialization_constant_impl(const char *SpecName,
                                         void *Value) const noexcept;
 
+  // \returns a bool value which indicates if specialization constant was set to
+  // a value different from default value.
   bool is_specialization_constant_set(const char *SpecName) const noexcept;
 
   detail::KernelBundleImplPtr impl;
diff --git a/sycl/source/detail/device_image_impl.hpp b/sycl/source/detail/device_image_impl.hpp
@@ -50,6 +50,8 @@ class device_image_impl {
     unsigned int CompositeOffset = 0;
     unsigned int Size = 0;
     unsigned int BlobOffset = 0;
+    // Indicates if the specialization constant was set to a value which is
+    // different from the default value.
     bool IsSet = false;
   };
 
@@ -61,7 +63,8 @@ class device_image_impl {
                     sycl::detail::pi::PiProgram Program)
       : MBinImage(BinImage), MContext(std::move(Context)),
         MDevices(std::move(Devices)), MState(State), MProgram(Program),
-        MKernelIDs(std::move(KernelIDs)) {
+        MKernelIDs(std::move(KernelIDs)),
+        MSpecConstsDefValBlob(getSpecConstsDefValBlob()) {
     updateSpecConstSymMap();
   }
 
@@ -74,6 +77,7 @@ class device_image_impl {
       : MBinImage(BinImage), MContext(std::move(Context)),
         MDevices(std::move(Devices)), MState(State), MProgram(Program),
         MKernelIDs(std::move(KernelIDs)), MSpecConstsBlob(SpecConstsBlob),
+        MSpecConstsDefValBlob(getSpecConstsDefValBlob()),
         MSpecConstSymMap(SpecConstMap) {}
 
   bool has_kernel(const kernel_id &KernelIDCand) const noexcept {
@@ -152,6 +156,21 @@ class device_image_impl {
     std::vector<SpecConstDescT> &Descs =
         MSpecConstSymMap[std::string{SpecName}];
     for (SpecConstDescT &Desc : Descs) {
+      // If there is a default value of the specialization constant and it is
+      // the same as the value which is being set then do nothing, runtime is
+      // going to handle this case just like if only the default value of the
+      // specialization constant was provided.
+      if (MSpecConstsDefValBlob.size() &&
+          (std::memcmp(MSpecConstsDefValBlob.begin() + Desc.BlobOffset,
+                       static_cast<const char *>(Value) + Desc.CompositeOffset,
+                       Desc.Size) == 0)) {
+        // Now we have default value, so reset to false.
+        Desc.IsSet = false;
+        continue;
+      }
+
+      // Value of the specialization constant is set to a value which is
+      // different from the default value.
       Desc.IsSet = true;
       std::memcpy(MSpecConstsBlob.data() + Desc.BlobOffset,
                   static_cast<const char *>(Value) + Desc.CompositeOffset,
@@ -161,19 +180,20 @@ class device_image_impl {
 
   void get_specialization_constant_raw_value(const char *SpecName,
                                              void *ValueRet) const noexcept {
-    assert(is_specialization_constant_set(SpecName));
+    bool IsSet = 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);
-
+    assert(IsSet || MSpecConstsDefValBlob.size());
     // 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) {
-
+      auto Blob =
+          IsSet ? MSpecConstsBlob.data() : MSpecConstsDefValBlob.begin();
       std::memcpy(static_cast<char *>(ValueRet) + Desc.CompositeOffset,
-                  MSpecConstsBlob.data() + Desc.BlobOffset, Desc.Size);
+                  Blob + Desc.BlobOffset, Desc.Size);
     }
   }
 
@@ -293,16 +313,30 @@ class device_image_impl {
   }
 
 private:
+  // Get the specialization constant default value blob.
+  ByteArray getSpecConstsDefValBlob() const {
+    if (!MBinImage)
+      return ByteArray(nullptr, 0);
+
+    // Get default values for specialization constants.
+    const RTDeviceBinaryImage::PropertyRange &SCDefValRange =
+        MBinImage->getSpecConstantsDefaultValues();
+    if (!SCDefValRange.size())
+      return ByteArray(nullptr, 0);
+
+    ByteArray DefValDescriptors =
+        DeviceBinaryProperty(*SCDefValRange.begin()).asByteArray();
+    // First 8 bytes are consumed by the size of the property.
+    DefValDescriptors.dropBytes(8);
+    return DefValDescriptors;
+  }
+
   void updateSpecConstSymMap() {
     if (MBinImage) {
       const RTDeviceBinaryImage::PropertyRange &SCRange =
           MBinImage->getSpecConstants();
       using SCItTy = RTDeviceBinaryImage::PropertyRange::ConstIterator;
 
-      // get default values for specialization constants
-      const RTDeviceBinaryImage::PropertyRange &SCDefValRange =
-          MBinImage->getSpecConstantsDefaultValues();
-
       // This variable is used to calculate spec constant value offset in a
       // flat byte array.
       unsigned BlobOffset = 0;
@@ -341,16 +375,13 @@ class device_image_impl {
       }
       MSpecConstsBlob.resize(BlobOffset);
 
-      bool HasDefaultValues = SCDefValRange.begin() != SCDefValRange.end();
-
-      if (HasDefaultValues) {
-        ByteArray DefValDescriptors =
-            DeviceBinaryProperty(*SCDefValRange.begin()).asByteArray();
-        assert(DefValDescriptors.size() - 8 == MSpecConstsBlob.size() &&
+      if (MSpecConstsDefValBlob.size()) {
+        assert(MSpecConstsDefValBlob.size() == MSpecConstsBlob.size() &&
                "Specialization constant default value blob do not have the "
                "expected size.");
-        std::uninitialized_copy(&DefValDescriptors[8],
-                                &DefValDescriptors[8] + MSpecConstsBlob.size(),
+        std::uninitialized_copy(MSpecConstsDefValBlob.begin(),
+                                MSpecConstsDefValBlob.begin() +
+                                    MSpecConstsBlob.size(),
                                 MSpecConstsBlob.data());
       }
     }
@@ -372,6 +403,9 @@ class device_image_impl {
   // Binary blob which can have values of all specialization constants in the
   // image
   std::vector<unsigned char> MSpecConstsBlob;
+  // Binary blob which can have default values of all specialization constants
+  // in the image.
+  const ByteArray MSpecConstsDefValBlob;
   // Buffer containing binary blob which can have values of all specialization
   // constants in the image, it is using for storing non-native specialization
   // constants
diff --git a/sycl/test-e2e/SpecConstants/2020/image_selection.cpp b/sycl/test-e2e/SpecConstants/2020/image_selection.cpp
@@ -23,12 +23,33 @@
 // RUN: env SYCL_PI_TRACE=-1 %{run} %t3.out | FileCheck --match-full-lines --check-prefix=CHECK-MIX %s
 // clang-format on
 
+// Check the behaviour when -fsycl-add-default-spec-consts-image option is used
+// and default value is explicitly set with the same value - we are supposed to
+// choose images with inlined values in this case.
+
+// clang-format off
+// RUN: %clangxx  -fsycl-add-default-spec-consts-image -fsycl -fsycl-targets=spir64_gen -Xsycl-target-backend=spir64_gen %gpu_aot_target_opts %s -o %t3.out
+// RUN: env SYCL_PI_TRACE=-1 %{run} %t3.out | FileCheck --match-full-lines --check-prefix=CHECK-DEFAULT-EXPLICIT-SET %s
+// clang-format on
+
+// Check the behaviour when -fsycl-add-default-spec-consts-image option is used
+// and value of specialization constant is changed to new value and then back to
+// the default value - we are supposed to choose images with inlined values in
+// this case.
+
+// clang-format off
+// RUN: %clangxx  -fsycl-add-default-spec-consts-image -fsycl -fsycl-targets=spir64_gen -Xsycl-target-backend=spir64_gen %gpu_aot_target_opts %s -o %t3.out
+// RUN: env SYCL_PI_TRACE=-1 %{run} %t3.out | FileCheck --match-full-lines --check-prefix=CHECK-DEFAULT-BACK-TO-DEFAULT %s
+// clang-format on
+
 #include <sycl/sycl.hpp>
 
 constexpr sycl::specialization_id<int> int_id(3);
 
 class Kernel1;
 class Kernel2;
+class Kernel3;
+class Kernel4;
 
 int main() {
   sycl::queue Q;
@@ -189,5 +210,64 @@ int main() {
   else
     std::cout << "Default value of specialization constant was used."
               << std::endl;
+
+  // Test that if user calls set_specialization_constant with the value equal to
+  // default then we choose image with inlined default values of specialization
+  // constants. We are verifying that by checking the 4th parameter is set to
+  // zero.
+  // CHECK-DEFAULT-EXPLICIT-SET: Default value was explicitly set
+  // CHECK-DEFAULT-EXPLICIT-SET: ---> piextKernelSetArgMemObj(
+  // CHECK-DEFAULT-EXPLICIT-SET-NEXT:	<unknown> : {{.*}}
+  // CHECK-DEFAULT-EXPLICIT-SET-NEXT:	<unknown> : {{.*}}
+  // CHECK-DEFAULT-EXPLICIT-SET-NEXT:	<unknown> : {{.*}}
+  // CHECK-DEFAULT-EXPLICIT-SET-NEXT:	<unknown> : 0
+  // CHECK-DEFAULT-EXPLICIT-SET-NEXT: ) ---> 	pi_result : PI_SUCCESS
+  // CHECK-DEFAULT-EXPLICIT-SET: Default value of specialization constant was used.
+  std::cout << "Default value was explicitly set" << std::endl;
+  Q.submit([&](sycl::handler &cgh) {
+     cgh.set_specialization_constant<int_id>(3);
+
+     cgh.single_task<Kernel3>([=](sycl::kernel_handler h) {
+       auto SpecConst = h.get_specialization_constant<int_id>();
+       *Res = SpecConst == 3 ? 0 : 1;
+     });
+   }).wait();
+
+  if (*Res)
+    std::cout << "New specialization constant value was set." << std::endl;
+  else
+    std::cout << "Default value of specialization constant was used."
+              << std::endl;
+
+  // Test that if user sets new value of specialization constant and then
+  // changes it back to default value then we choose image with inlined default
+  // values of specialization constants. We are verifying that by checking the
+  // 4th parameter is set to zero.
+  // CHECK-DEFAULT-BACK-TO-DEFAULT: Changed to new value and then default value was explicitly set
+  // CHECK-DEFAULT-BACK-TO-DEFAULT: ---> piextKernelSetArgMemObj(
+  // CHECK-DEFAULT-BACK-TO-DEFAULT-NEXT:	<unknown> : {{.*}}
+  // CHECK-DEFAULT-BACK-TO-DEFAULT-NEXT:	<unknown> : {{.*}}
+  // CHECK-DEFAULT-BACK-TO-DEFAULT-NEXT:	<unknown> : {{.*}}
+  // CHECK-DEFAULT-BACK-TO-DEFAULT-NEXT:	<unknown> : 0
+  // CHECK-DEFAULT-BACK-TO-DEFAULT-NEXT: ) ---> 	pi_result : PI_SUCCESS
+  // CHECK-DEFAULT-BACK-TO-DEFAULT: Default value of specialization constant was used.
+  std::cout << "Changed to new value and then default value was explicitly set"
+            << std::endl;
+  Q.submit([&](sycl::handler &cgh) {
+     cgh.set_specialization_constant<int_id>(4);
+     cgh.set_specialization_constant<int_id>(3);
+
+     cgh.single_task<Kernel4>([=](sycl::kernel_handler h) {
+       auto SpecConst = h.get_specialization_constant<int_id>();
+       *Res = SpecConst == 3 ? 0 : 1;
+     });
+   }).wait();
+
+  if (*Res)
+    std::cout << "New specialization constant value was set." << std::endl;
+  else
+    std::cout << "Default value of specialization constant was used."
+              << std::endl;
+
   return 0;
 }
