[SYCL] Fix piextProgramSetSpecializationConstant usage (#6125)

This entry point is only meant to be called when native specialization
constants are supported, so AOT only targets don't support it.

Still implement it in the AOT only plugins to have clearer error if it's
accidentally called.

This is mentioned in #6093, although it doesn't touch the alignment issue.

Author: npmiller

sycl/include/CL/sycl/detail/pi.h

@@ -1283,6 +1283,9 @@ __SYCL_EXPORT pi_result piProgramRelease(pi_program program);
 
 /// Sets a specialization constant to a specific value.
 ///
+/// Note: Only used when specialization constants are natively supported (SPIR-V
+/// binaries), and not when they are emulated (AOT binaries).
+///
 /// \param prog the program object which will use the value
 /// \param spec_id integer ID of the constant
 /// \param spec_size size of the value

sycl/plugins/cuda/pi_cuda.cpp

@@ -3436,6 +3436,15 @@ pi_result cuda_piKernelSetExecInfo(pi_kernel, pi_kernel_exec_info, size_t,
   return PI_SUCCESS;
 }
 
+pi_result cuda_piextProgramSetSpecializationConstant(pi_program, pi_uint32,
+                                                     size_t, const void *) {
+  // This entry point is only used for native specialization constants (SPIR-V),
+  // and the CUDA plugin is AOT only so this entry point is not supported.
+  cl::sycl::detail::pi::die(
+      "Native specialization constants are not supported");
+  return {};
+}
+
 pi_result cuda_piextKernelSetArgPointer(pi_kernel kernel, pi_uint32 arg_index,
                                         size_t arg_size,
                                         const void *arg_value) {
@@ -5066,6 +5075,8 @@ pi_result piPluginInit(pi_plugin *PluginInit) {
   _PI_CL(piKernelRetain, cuda_piKernelRetain)
   _PI_CL(piKernelRelease, cuda_piKernelRelease)
   _PI_CL(piKernelSetExecInfo, cuda_piKernelSetExecInfo)
+  _PI_CL(piextProgramSetSpecializationConstant,
+         cuda_piextProgramSetSpecializationConstant)
   _PI_CL(piextKernelSetArgPointer, cuda_piextKernelSetArgPointer)
   _PI_CL(piextKernelCreateWithNativeHandle,
          cuda_piextKernelCreateWithNativeHandle)

sycl/plugins/hip/pi_hip.cpp

@@ -3375,6 +3375,15 @@ pi_result hip_piKernelSetExecInfo(pi_kernel kernel,
   return PI_SUCCESS;
 }
 
+pi_result hip_piextProgramSetSpecializationConstant(pi_program, pi_uint32,
+                                                    size_t, const void *) {
+  // This entry point is only used for native specialization constants (SPIR-V),
+  // and the HIP plugin is AOT only so this entry point is not supported.
+  cl::sycl::detail::pi::die(
+      "Native specialization constants are not supported");
+  return {};
+}
+
 pi_result hip_piextKernelSetArgPointer(pi_kernel kernel, pi_uint32 arg_index,
                                        size_t arg_size, const void *arg_value) {
   kernel->set_kernel_arg(arg_index, arg_size, arg_value);
@@ -4959,6 +4968,8 @@ pi_result piPluginInit(pi_plugin *PluginInit) {
   _PI_CL(piKernelRetain, hip_piKernelRetain)
   _PI_CL(piKernelRelease, hip_piKernelRelease)
   _PI_CL(piKernelSetExecInfo, hip_piKernelSetExecInfo)
+  _PI_CL(piextProgramSetSpecializationConstant,
+         hip_piextProgramSetSpecializationConstant)
   _PI_CL(piextKernelSetArgPointer, hip_piextKernelSetArgPointer)
   // Event
   _PI_CL(piEventCreate, hip_piEventCreate)

sycl/source/detail/program_manager/program_manager.cpp

@@ -1855,24 +1855,22 @@ device_image_plain ProgramManager::build(const device_image_plain &DeviceImage,
     auto [NativePrg, DeviceCodeWasInCache] = getOrCreatePIProgram(
         Img, Context, Devs[0], CompileOpts + LinkOpts, SpecConsts);
 
-    if (!DeviceCodeWasInCache) {
-      if (InputImpl->get_bin_image_ref()->supportsSpecConstants())
-        enableITTAnnotationsIfNeeded(NativePrg, Plugin);
-
-      {
-        std::lock_guard<std::mutex> Lock{InputImpl->get_spec_const_data_lock()};
-        const std::map<std::string,
-                       std::vector<device_image_impl::SpecConstDescT>>
-            &SpecConstData = InputImpl->get_spec_const_data_ref();
-
-        for (const auto &DescPair : SpecConstData) {
-          for (const device_image_impl::SpecConstDescT &SpecIDDesc :
-               DescPair.second) {
-            if (SpecIDDesc.IsSet) {
-              Plugin.call<PiApiKind::piextProgramSetSpecializationConstant>(
-                  NativePrg, SpecIDDesc.ID, SpecIDDesc.Size,
-                  SpecConsts.data() + SpecIDDesc.BlobOffset);
-            }
+    if (!DeviceCodeWasInCache &&
+        InputImpl->get_bin_image_ref()->supportsSpecConstants()) {
+      enableITTAnnotationsIfNeeded(NativePrg, Plugin);
+
+      std::lock_guard<std::mutex> Lock{InputImpl->get_spec_const_data_lock()};
+      const std::map<std::string,
+                     std::vector<device_image_impl::SpecConstDescT>>
+          &SpecConstData = InputImpl->get_spec_const_data_ref();
+
+      for (const auto &DescPair : SpecConstData) {
+        for (const device_image_impl::SpecConstDescT &SpecIDDesc :
+             DescPair.second) {
+          if (SpecIDDesc.IsSet) {
+            Plugin.call<PiApiKind::piextProgramSetSpecializationConstant>(
+                NativePrg, SpecIDDesc.ID, SpecIDDesc.Size,
+                SpecConsts.data() + SpecIDDesc.BlobOffset);
           }
         }
       }