hook specialization constants into kernel bundle interface

Author: ndgrigorian

dpctl/_backend.pxd

@@ -440,7 +440,9 @@ cdef extern from "syclinterface/dpctl_sycl_kernel_bundle_interface.h":
         const DPCTLSyclDeviceRef Dev,
         const void *IL,
         size_t Length,
-        const char *CompileOpts)
+        const char *CompileOpts,
+        size_t NumSpecConsts,
+        const _spec_const *SpecConsts)
     cdef DPCTLSyclKernelBundleRef DPCTLKernelBundle_CreateFromOCLSource(
         const DPCTLSyclContextRef Ctx,
         const DPCTLSyclDeviceRef Dev,

dpctl/program/_program.pxd

@@ -63,7 +63,10 @@ cpdef create_kernel_bundle_from_source (
     SyclQueue q, unicode source, unicode copts=*
 )
 cpdef create_kernel_bundle_from_spirv (
-    SyclQueue q, const unsigned char[:] IL, unicode copts=*
+    SyclQueue q,
+    const unsigned char[:] IL,
+    unicode copts=*,
+    list specializations=*,
 )
 cpdef create_program_from_source (SyclQueue q, unicode source, unicode copts=*)
 cpdef create_program_from_spirv (

dpctl/program/_program.pyx

@@ -36,6 +36,7 @@ from cpython.buffer cimport (
 )
 from cpython.bytes cimport PyBytes_FromStringAndSize
 from libc.stdint cimport uint32_t
+from libc.stdlib cimport free, malloc
 from libc.string cimport memcmp
 
 import warnings
@@ -469,7 +470,10 @@ cpdef create_kernel_bundle_from_source(SyclQueue q, str src, str copts=""):
 
 
 cpdef create_kernel_bundle_from_spirv(
-    SyclQueue q, const unsigned char[:] IL, str copts=""
+    SyclQueue q,
+    const unsigned char[:] IL,
+    str copts="",
+    list specializations=None,
 ):
     """
         Creates a Sycl interoperability kernel bundle from an SPIR-V binary.
@@ -487,7 +491,9 @@ cpdef create_kernel_bundle_from_spirv(
             copts (str, optional)
                 Optional compilation flags that will be used
                 when compiling the kernel bundle. Default: ``""``.
-
+            specializations (list, optional)
+                A list of :class:`.SpecializationConstant` objects to be used
+                when creating the kernel bundle. Default: ``None``.
         Returns:
             kernel_bundle (:class:`.SyclKernelBundle`)
                 A :class:`.SyclKernelBundle` object wrapping the
@@ -506,11 +512,44 @@ cpdef create_kernel_bundle_from_spirv(
     cdef size_t length = IL.shape[0]
     cdef bytes bCOpts = copts.encode("utf8")
     cdef const char *COpts = <const char*>bCOpts
-    KBref = DPCTLKernelBundle_CreateFromSpirv(
-        CRef, DRef, <const void*>dIL, length, COpts
-    )
-    if KBref is NULL:
-        raise SyclKernelBundleCompilationError()
+    cdef size_t num_spconsts
+    cdef _spec_const *spconsts
+    cdef SpecializationConstant spconst
+
+    if specializations is not None:
+        num_spconsts = len(specializations)
+        spconsts = <_spec_const *>(
+            malloc(num_spconsts * sizeof(_spec_const))
+        )
+        if spconsts == NULL:
+            raise MemoryError(
+                "Failed to allocate memory for specialization constants."
+            )
+        for i, spconst in enumerate(specializations):
+            if not isinstance(spconst, SpecializationConstant):
+                free(spconsts)
+                raise TypeError(
+                    "All items in specializations must be of type "
+                    f"`SpecializationConstant`, got {type(spconst)}"
+                )
+            spconsts[i] = spconst._spec_const
+    else:
+        num_spconsts = 0
+        spconsts = NULL
+    try:
+        KBref = DPCTLKernelBundle_CreateFromSpirv(
+            CRef,
+            DRef,
+            <const void*>dIL,
+            length, COpts,
+            num_spconsts,
+            spconsts,
+        )
+        if KBref is NULL:
+            raise SyclKernelBundleCompilationError()
+    finally:
+        if spconsts != NULL:
+            free(spconsts)
 
     return SyclKernelBundle._create(KBref)
 

libsyclinterface/include/syclinterface/dpctl_sycl_kernel_bundle_interface.h

@@ -58,6 +58,8 @@ typedef struct DPCTLSpecConstTy
  * @param    Length         The size of the IL binary in bytes.
  * @param    CompileOpts    Optional compiler flags used when compiling the
  *                          SPIR-V binary.
+ * @param    NumSpecConsts  The number of specialization constants.
+ * @param    SpecConsts     An array of specialization constants.
  * @return   A new SyclKernelBundleRef pointer if the kernel_bundle creation
  * succeeded, else returns NULL.
  * @ingroup KernelBundleInterface
@@ -68,7 +70,9 @@ DPCTLKernelBundle_CreateFromSpirv(__dpctl_keep const DPCTLSyclContextRef Ctx,
                                   __dpctl_keep const DPCTLSyclDeviceRef Dev,
                                   __dpctl_keep const void *IL,
                                   size_t Length,
-                                  const char *CompileOpts);
+                                  const char *CompileOpts,
+                                  size_t NumSpecConsts,
+                                  const DPCTLSpecConst *SpecConsts);
 
 /*!
  * @brief Create a Sycl kernel bundle from an OpenCL kernel source string.

libsyclinterface/source/dpctl_sycl_kernel_bundle_interface.cpp

@@ -31,6 +31,7 @@
 #include "dpctl_error_handlers.h"
 #include "dpctl_sycl_type_casters.hpp"
 #include <CL/cl.h> /* OpenCL headers     */
+#include <cstdint>
 #include <sstream>
 #include <stddef.h>
 #include <sycl/backend/opencl.hpp>
@@ -170,6 +171,21 @@ std::string _GetErrorCode_ocl_impl(cl_int code)
     }
 }
 
+typedef cl_int (*clSetProgramSpecializationConstantFT)(cl_program,
+                                                       cl_uint,
+                                                       size_t,
+                                                       const void *);
+const char *clSetProgramSpecializationConstant_Name =
+    "clSetProgramSpecializationConstant";
+clSetProgramSpecializationConstantFT get_clSetProgramSpecializationConstant()
+{
+    static auto st_clSetProgramSpecializationConstantF =
+        cl_loader::get().getSymbol<clSetProgramSpecializationConstantFT>(
+            clSetProgramSpecializationConstant_Name);
+
+    return st_clSetProgramSpecializationConstantF;
+}
+
 DPCTLSyclKernelBundleRef
 _CreateKernelBundle_common_ocl_impl(cl_program clProgram,
                                     const context &ctx,
@@ -235,7 +251,9 @@ _CreateKernelBundleWithIL_ocl_impl(const context &ctx,
                                    const device &dev,
                                    const void *IL,
                                    size_t il_length,
-                                   const char *CompileOpts)
+                                   const char *CompileOpts,
+                                   size_t NumSpecConsts,
+                                   const DPCTLSpecConst *SpecConsts)
 {
     auto clCreateProgramWithILF = get_clCreateProgramWithIL();
     if (clCreateProgramWithILF == nullptr) {
@@ -257,6 +275,22 @@ _CreateKernelBundleWithIL_ocl_impl(const context &ctx,
         return nullptr;
     }
 
+    if (SpecConsts != nullptr && NumSpecConsts > 0) {
+        auto clSetProgramSpecConstF = get_clSetProgramSpecializationConstant();
+        if (clSetProgramSpecConstF) {
+            for (size_t i = 0; i < NumSpecConsts; ++i) {
+                clSetProgramSpecConstF(clProgram, SpecConsts[i].id,
+                                       SpecConsts[i].size, SpecConsts[i].value);
+            }
+        }
+        else {
+            error_handler("clSetProgramSpecializationConstant is not available "
+                          "in the OpenCL implementation.",
+                          __FILE__, __func__, __LINE__);
+            return nullptr;
+        }
+    }
+
     return _CreateKernelBundle_common_ocl_impl(clProgram, ctx, dev,
                                                CompileOpts);
 }
@@ -428,7 +462,9 @@ _CreateKernelBundleWithIL_ze_impl(const context &SyclCtx,
                                   const device &SyclDev,
                                   const void *IL,
                                   size_t il_length,
-                                  const char *CompileOpts)
+                                  const char *CompileOpts,
+                                  size_t NumSpecConsts,
+                                  const DPCTLSpecConst *SpecConsts)
 {
     auto zeModuleCreateFn = get_zeModuleCreate();
     if (zeModuleCreateFn == nullptr) {
@@ -444,8 +480,22 @@ _CreateKernelBundleWithIL_ze_impl(const context &SyclCtx,
     ZeDevice = get_native<ze_be>(SyclDev);
 
     // Specialization constants are not supported by DPCTL at the moment
+    std::vector<std::uint32_t> spec_ids;
+    std::vector<const void *> spec_values;
+
+    if (SpecConsts != nullptr && NumSpecConsts > 0) {
+        spec_ids.reserve(NumSpecConsts);
+        spec_values.reserve(NumSpecConsts);
+        for (size_t i = 0; i < NumSpecConsts; ++i) {
+            spec_ids.push_back(SpecConsts[i].id);
+            spec_values.push_back(SpecConsts[i].value);
+        }
+    }
     ze_module_constants_t ZeSpecConstants = {};
-    ZeSpecConstants.numConstants = 0;
+    ZeSpecConstants.numConstants = static_cast<std::uint32_t>(NumSpecConsts);
+    ZeSpecConstants.pConstantIds = spec_ids.empty() ? nullptr : spec_ids.data();
+    ZeSpecConstants.pConstantValues =
+        spec_values.empty() ? nullptr : spec_values.data();
 
     // Populate the Level Zero module descriptions
     ze_module_desc_t ZeModuleDesc = {};
@@ -583,7 +633,9 @@ DPCTLKernelBundle_CreateFromSpirv(__dpctl_keep const DPCTLSyclContextRef CtxRef,
                                   __dpctl_keep const DPCTLSyclDeviceRef DevRef,
                                   __dpctl_keep const void *IL,
                                   size_t length,
-                                  const char *CompileOpts)
+                                  const char *CompileOpts,
+                                  size_t NumSpecConsts,
+                                  const DPCTLSpecConst *SpecConsts)
 {
     DPCTLSyclKernelBundleRef KBRef = nullptr;
     if (!CtxRef) {
@@ -611,12 +663,14 @@ DPCTLKernelBundle_CreateFromSpirv(__dpctl_keep const DPCTLSyclContextRef CtxRef,
     switch (BE) {
     case backend::opencl:
         KBRef = _CreateKernelBundleWithIL_ocl_impl(*SyclCtx, *SyclDev, IL,
-                                                   length, CompileOpts);
+                                                   length, CompileOpts,
+                                                   NumSpecConsts, SpecConsts);
         break;
     case backend::ext_oneapi_level_zero:
 #ifdef DPCTL_ENABLE_L0_PROGRAM_CREATION
         KBRef = _CreateKernelBundleWithIL_ze_impl(*SyclCtx, *SyclDev, IL,
-                                                  length, CompileOpts);
+                                                  length, CompileOpts,
+                                                  NumSpecConsts, SpecConsts);
         break;
 #endif
     default: