UR updates

Author: jchlanda

sycl/plugins/level_zero/pi_level_zero.cpp

@@ -1061,10 +1061,11 @@ pi_result piKernelSetExecInfo(pi_kernel Kernel, pi_kernel_exec_info ParamName,
 }
 
 pi_result piextProgramSetSpecializationConstant(pi_program Prog,
+                                                pi_kernel Kernel,
                                                 pi_uint32 SpecID, size_t Size,
                                                 const void *SpecValue) {
-  return pi2ur::piextProgramSetSpecializationConstant(Prog, SpecID, Size,
-                                                      SpecValue);
+  return pi2ur::piextProgramSetSpecializationConstant(Prog, Kernel, SpecID,
+                                                      Size, SpecValue);
 }
 
 // Command buffer extension

sycl/plugins/native_cpu/pi_native_cpu.cpp

@@ -1183,7 +1183,8 @@ pi_result piKernelSetExecInfo(pi_kernel, pi_kernel_exec_info, size_t,
   return PI_SUCCESS;
 }
 
-pi_result piextProgramSetSpecializationConstant(pi_program, pi_uint32, size_t,
+pi_result piextProgramSetSpecializationConstant(pi_program, pi_kernel,
+                                                pi_uint32, size_t,
                                                 const void *) {
   DIE_NO_IMPLEMENTATION;
 }

sycl/plugins/opencl/pi_opencl.cpp

@@ -2212,9 +2212,11 @@ pi_result piKernelSetExecInfo(pi_kernel kernel, pi_kernel_exec_info param_name,
 }
 
 pi_result piextProgramSetSpecializationConstant(pi_program prog,
+                                                pi_kernel kernel,
                                                 pi_uint32 spec_id,
                                                 size_t spec_size,
                                                 const void *spec_value) {
+  std::ignore = kernel;
   cl_program ClProg = cast<cl_program>(prog);
   cl_context Ctx = nullptr;
   size_t RetSize = 0;

sycl/plugins/unified_runtime/pi2ur.hpp

@@ -1912,18 +1912,22 @@ piProgramBuild(pi_program Program, pi_uint32 NumDevices,
 }
 
 inline pi_result piextProgramSetSpecializationConstant(pi_program Program,
+                                                       pi_kernel Kernel,
                                                        pi_uint32 SpecID,
                                                        size_t Size,
                                                        const void *SpecValue) {
   ur_program_handle_t UrProgram =
       reinterpret_cast<ur_program_handle_t>(Program);
+  ur_kernel_handle_t UrKernel = reinterpret_cast<ur_kernel_handle_t>(Kernel);
   uint32_t Count = 1;
   ur_specialization_constant_info_t SpecConstant{};
   SpecConstant.id = SpecID;
   SpecConstant.size = Size;
   SpecConstant.pValue = SpecValue;
-  HANDLE_ERRORS(
-      urProgramSetSpecializationConstants(UrProgram, Count, &SpecConstant));
+  HANDLE_ERRORS(Kernel ? urKernelSetSpecializationConstants(UrKernel, Count,
+                                                            &SpecConstant)
+                       : urProgramSetSpecializationConstants(UrProgram, Count,
+                                                             &SpecConstant));
 
   return PI_SUCCESS;
 }

sycl/plugins/unified_runtime/pi_unified_runtime.cpp

@@ -127,9 +127,10 @@ __SYCL_EXPORT pi_result piProgramBuild(
 }
 
 __SYCL_EXPORT pi_result piextProgramSetSpecializationConstant(
-    pi_program Prog, pi_uint32 SpecID, size_t Size, const void *SpecValue) {
-  return pi2ur::piextProgramSetSpecializationConstant(Prog, SpecID, Size,
-                                                      SpecValue);
+    pi_program Prog, pi_kernel Kernel, pi_uint32 SpecID, size_t Size,
+    const void *SpecValue) {
+  return pi2ur::piextProgramSetSpecializationConstant(Prog, Kernel, SpecID,
+                                                      Size, SpecValue);
 }
 
 __SYCL_EXPORT pi_result

sycl/plugins/unified_runtime/ur/adapters/cuda/device.cpp

@@ -993,7 +993,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
                        strnlen(AddressBuffer, AddressBufferSize - 1) + 1);
   }
   case UR_DEVICE_INFO_KERNEL_SET_SPECIALIZATION_CONSTANTS:
-    return ReturnValue(false);
+    return ReturnValue(true);
     // TODO: Investigate if this information is available on CUDA.
   case UR_DEVICE_INFO_GPU_EU_COUNT:
   case UR_DEVICE_INFO_GPU_EU_SIMD_WIDTH:

sycl/plugins/unified_runtime/ur/adapters/cuda/kernel.cpp

@@ -9,6 +9,7 @@
 #include "kernel.hpp"
 #include "memory.hpp"
 #include "sampler.hpp"
+#include "ur_api.h"
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urKernelCreate(ur_program_handle_t hProgram, const char *pKernelName,
@@ -374,3 +375,35 @@ urKernelSetArgSampler(ur_kernel_handle_t hKernel, uint32_t argIndex,
   }
   return Result;
 }
+
+UR_APIEXPORT ur_result_t UR_APICALL urKernelSetSpecializationConstants(
+    ur_kernel_handle_t Kernel, uint32_t,
+    const ur_specialization_constant_info_t *SpecConstants) {
+  CUdeviceptr DPtr = 0;
+  size_t Bytes = 0;
+  // NOTE: GlobalNamePrefix is generated by CUDANativeSpecConstantsPass and
+  // must not be changed.
+  constexpr char GlobalNamePrefix[] = "sycl_specialization_constants_kernel_";
+  std::string Name{GlobalNamePrefix};
+  Name.append(Kernel->getName());
+  const auto ResGetGlobal = cuModuleGetGlobal(
+      &DPtr, &Bytes, Kernel->get_program()->get(), Name.c_str());
+  UR_CHECK_ERROR(ResGetGlobal);
+  // NOTE: The size of the symbol here - 1 - is important (the value is invalid
+  // and would result in a failure from `cuMemcpyHtoD`), instead, it's used as
+  // a flag to communicate to the plugin that even though the implicit kernel
+  // argument is present (hence the call to
+  // `cuda_piextProgramSetSpecializationConstant`), it has no uses and there is
+  // no need to set up the symbol.
+  // See: `CUDANativeSpecConstantsPass::setUpPlaceholderEntries` for how pi
+  // handles it.
+  if (Bytes == 1) {
+    return UR_RESULT_SUCCESS;
+  }
+
+  UR_ASSERT(DPtr && Bytes, UR_RESULT_ERROR_INVALID_VALUE);
+
+  const auto ResMemcpy =
+      cuMemcpyHtoD(DPtr, SpecConstants->pValue, SpecConstants->size);
+  return UR_CHECK_ERROR(ResMemcpy);
+}

sycl/source/detail/program_impl.cpp

@@ -549,7 +549,7 @@ void program_impl::flush_spec_constants(
           Descriptors.consume<uint32_t, uint32_t, uint32_t>();
 
       Ctx->getPlugin()->call<PiApiKind::piextProgramSetSpecializationConstant>(
-          NativePrg, Id, Size, SC.getValuePtr() + Offset);
+          NativePrg, nullptr, Id, Size, SC.getValuePtr() + Offset);
     }
   }
 }

sycl/source/detail/program_manager/program_manager.cpp

@@ -60,7 +60,7 @@ enableITTAnnotationsIfNeeded(const sycl::detail::pi::PiProgram &Prog,
   if (SYCLConfig<INTEL_ENABLE_OFFLOAD_ANNOTATIONS>::get() != nullptr) {
     constexpr char SpecValue = 1;
     Plugin->call<PiApiKind::piextProgramSetSpecializationConstant>(
-        Prog, ITTSpecConstId, sizeof(char), &SpecValue);
+        Prog, nullptr, ITTSpecConstId, sizeof(char), &SpecValue);
   }
 }
 
@@ -2095,7 +2095,7 @@ setSpecializationConstants(const std::shared_ptr<device_image_impl> &InputImpl,
     for (const device_image_impl::SpecConstDescT &SpecIDDesc : SpecConstDescs) {
       if (SpecIDDesc.IsSet) {
         Plugin->call<PiApiKind::piextProgramSetSpecializationConstant>(
-            Prog, SpecIDDesc.ID, SpecIDDesc.Size,
+            Prog, nullptr, SpecIDDesc.ID, SpecIDDesc.Size,
             SpecConsts.data() + SpecIDDesc.BlobOffset);
       }
     }

sycl/source/detail/scheduler/commands.cpp

@@ -2271,10 +2271,10 @@ static pi_result SetKernelParamsAndLaunch(
           SpecConstsBuffer ? &SpecConstsBuffer : nullptr;
 
       // Call into set spec constant for pi cuda kernels.
-      if (Queue->getPlugin().getBackend() == backend::ext_oneapi_cuda) {
+      if (Queue->getDeviceImplPtr()->getBackend() == backend::ext_oneapi_cuda) {
         static unsigned SpecID = 0;
         auto &Blob = DeviceImageImpl->get_spec_const_blob_ref();
-        Plugin.call<PiApiKind::piextProgramSetSpecializationConstant>(
+        Plugin->call<PiApiKind::piextProgramSetSpecializationConstant>(
             DeviceImageImpl->get_program_ref(), Kernel, SpecID++, Blob.size(),
             Blob.data());
       }