apply review comments

Author: ndgrigorian

dpctl/program/__init__.py

@@ -21,6 +21,7 @@
 
 """
 
+from . import utils
 from ._program import (
     SpecializationConstant,
     SyclKernel,

dpctl/program/_program.pyx

@@ -36,7 +36,7 @@ from cpython.buffer cimport (
     PyObject_GetBuffer,
 )
 from cpython.bytes cimport PyBytes_FromStringAndSize
-from libc.stdint cimport uint32_t
+from libc.stdint cimport UINT32_MAX, uint32_t
 from libc.stdlib cimport free, malloc
 from libc.string cimport memcmp, memcpy
 
@@ -337,13 +337,23 @@ cdef class SpecializationConstant:
                 f"{len(args) + 1}."
             )
 
+        if spec_id < 0 or spec_id > UINT32_MAX:
+            raise ValueError(
+                "Specialization constant ID must fit in a 32-bit unsigned "
+                f"integer (0 <= id <= {UINT32_MAX}), got {spec_id}"
+            )
         self._spec_const.id = <uint32_t>spec_id
 
         if len(args) == 2:
             if (
                 isinstance(args[0], numbers.Integral) and
                 isinstance(args[1], numbers.Integral)
             ):
+                if args[0] <= 0:
+                    raise ValueError(
+                        "Size must be a positive integer when constructing "
+                        f"from a pointer and size, got {args[0]}"
+                    )
                 target_obj = PyBytes_FromStringAndSize(
                     <const char *><size_t>args[1], <Py_ssize_t>args[0]
                 )
@@ -380,6 +390,12 @@ cdef class SpecializationConstant:
                 "Failed to get buffer view for the provided object."
             )
 
+        if _local_buffer.len <= 0:
+            PyBuffer_Release(&(_local_buffer))
+            raise ValueError(
+                "Specialization constant value must not be zero-sized."
+            )
+
         self._spec_const.size = <size_t>_local_buffer.len
         copied_data = malloc(self._spec_const.size)
 
@@ -403,7 +419,7 @@ cdef class SpecializationConstant:
 
     def __eq__(self, other):
         if not isinstance(other, SpecializationConstant):
-            return False
+            return NotImplemented
         cdef SpecializationConstant _other = <SpecializationConstant>other
         if (
             self._spec_const.id != _other._spec_const.id or
@@ -418,6 +434,16 @@ cdef class SpecializationConstant:
             self._spec_const.size
         ) == 0
 
+    def __hash__(self):
+        cdef bytes value_bytes
+        if self._spec_const.value != NULL:
+            value_bytes = (<char *>self._spec_const.value)[
+                :self._spec_const.size
+            ]
+        else:
+            value_bytes = b""
+        return hash((self._spec_const.id, value_bytes))
+
     @property
     def id(self):
         """Returns the specialization ID for this specialization constant."""
@@ -551,13 +577,14 @@ cpdef create_kernel_bundle_from_spirv(
             raise MemoryError(
                 "Failed to allocate memory for specialization constants."
             )
-        for i, spconst in enumerate(specializations):
-            if not isinstance(spconst, SpecializationConstant):
+        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)}"
+                    f"`SpecializationConstant`, got {type(_spconst)}"
                 )
+            spconst = <SpecializationConstant>_spconst
             spconsts[i] = spconst._spec_const
     else:
         num_spconsts = 0

dpctl/program/utils/_utils.py

@@ -117,9 +117,13 @@ def parse_spirv_specializations(
             # parsing, so we can stop parsing at this point
             break
         elif opcode == SpirvOpCode.OpTypeBool:
+            if word_count < 2:
+                raise ValueError(f"Truncated OpTypeBool at word index {i}")
             result_id = int(words[i + 1])
             types[result_id] = {"dtype": "?", "itemsize": 1}
         elif opcode == SpirvOpCode.OpTypeInt:
+            if word_count < 4:
+                raise ValueError(f"Truncated OpTypeInt at word index {i}")
             result_id = int(words[i + 1])
             width = int(words[i + 2])
             signed = int(words[i + 3])
@@ -129,37 +133,62 @@ def parse_spirv_specializations(
                 "itemsize": width // 8,
             }
         elif opcode == SpirvOpCode.OpTypeFloat:
+            if word_count < 3:
+                raise ValueError(f"Truncated OpTypeFloat at word index {i}")
             result_id = int(words[i + 1])
             width = int(words[i + 2])
             types[result_id] = {
                 "dtype": f"f{width // 8}",
                 "itemsize": width // 8,
             }
         elif opcode == SpirvOpCode.OpSpecConstant:
+            if word_count < 3:
+                raise ValueError(f"Truncated OpSpecConstant at word index {i}")
             type_id = int(words[i + 1])
             result_id = int(words[i + 2])
             constants[result_id] = type_id
             literal_words = words[i + 3 : i + word_count]
             defaults[result_id] = literal_words.tobytes()
         elif opcode == SpirvOpCode.OpSpecConstantTrue:
+            if word_count < 3:
+                raise ValueError(
+                    f"Truncated OpSpecConstantTrue at word index {i}"
+                )
             type_id = int(words[i + 1])
             result_id = int(words[i + 2])
             constants[result_id] = type_id
             defaults[result_id] = True
         elif opcode == SpirvOpCode.OpSpecConstantFalse:
+            if word_count < 3:
+                raise ValueError(
+                    f"Truncated OpSpecConstantFalse at word index {i}"
+                )
             type_id = int(words[i + 1])
             result_id = int(words[i + 2])
             constants[result_id] = type_id
             defaults[result_id] = False
         elif opcode == SpirvOpCode.OpDecorate:
+            if word_count < 3:
+                raise ValueError(f"Truncated OpDecorate at word index {i}")
             target_id = int(words[i + 1])
             decoration = int(words[i + 2])
             if decoration == SpirvDecoration.SpecId:
+                if word_count < 4:
+                    raise ValueError(
+                        f"Truncated OpDecorate SpecId at word index {i}"
+                    )
                 ids[target_id] = int(words[i + 3])
         elif opcode == SpirvOpCode.OpName:
+            if word_count < 2:
+                raise ValueError(f"Truncated OpName at word index {i}")
             target_id = int(words[i + 1])
             name_bytes = words[i + 2 : i + word_count].tobytes()
-            names[target_id] = name_bytes.split(b"\x00", 1)[0].decode("utf-8")
+            try:
+                names[target_id] = name_bytes.split(b"\x00", 1)[0].decode(
+                    "utf-8"
+                )
+            except UnicodeDecodeError:
+                raise ValueError(f"Invalid UTF-8 in OpName at word index {i}")
 
         i += word_count
 

libsyclinterface/include/syclinterface/dpctl_sycl_kernel_bundle_interface.h

@@ -65,14 +65,14 @@ typedef struct DPCTLSpecConstTy
  * @ingroup KernelBundleInterface
  */
 DPCTL_API
-__dpctl_give DPCTLSyclKernelBundleRef
-DPCTLKernelBundle_CreateFromSpirv(__dpctl_keep const DPCTLSyclContextRef Ctx,
-                                  __dpctl_keep const DPCTLSyclDeviceRef Dev,
-                                  __dpctl_keep const void *IL,
-                                  size_t Length,
-                                  const char *CompileOpts,
-                                  size_t NumSpecConsts,
-                                  const DPCTLSpecConst *SpecConsts);
+__dpctl_give DPCTLSyclKernelBundleRef DPCTLKernelBundle_CreateFromSpirv(
+    __dpctl_keep const DPCTLSyclContextRef Ctx,
+    __dpctl_keep const DPCTLSyclDeviceRef Dev,
+    __dpctl_keep const void *IL,
+    size_t Length,
+    const char *CompileOpts,
+    size_t NumSpecConsts,
+    __dpctl_keep const DPCTLSpecConst *SpecConsts);
 
 /*!
  * @brief Create a Sycl kernel bundle from an OpenCL kernel source string.

libsyclinterface/source/dpctl_sycl_kernel_bundle_interface.cpp

@@ -293,6 +293,12 @@ _CreateKernelBundleWithIL_ocl_impl(const context &ctx,
         return nullptr;
     }
 
+    if (NumSpecConsts > 0 && SpecConsts == nullptr) {
+        error_handler("NumSpecConsts > 0 but SpecConsts is NULL.", __FILE__,
+                      __func__, __LINE__);
+        return nullptr;
+    }
+
     if (SpecConsts != nullptr && NumSpecConsts > 0) {
         auto clSetProgramSpecConstF = get_clSetProgramSpecializationConstant();
         if (clSetProgramSpecConstF) {
@@ -520,6 +526,12 @@ _CreateKernelBundleWithIL_ze_impl(const context &SyclCtx,
     backend_traits<ze_be>::return_type<device> ZeDevice;
     ZeDevice = get_native<ze_be>(SyclDev);
 
+    if (NumSpecConsts > 0 && SpecConsts == nullptr) {
+        error_handler("NumSpecConsts > 0 but SpecConsts is NULL.", __FILE__,
+                      __func__, __LINE__);
+        return nullptr;
+    }
+
     std::vector<std::uint32_t> spec_ids;
     std::vector<const void *> spec_values;
 
@@ -532,7 +544,7 @@ _CreateKernelBundleWithIL_ze_impl(const context &SyclCtx,
         }
     }
     ze_module_constants_t ZeSpecConstants = {};
-    ZeSpecConstants.numConstants = static_cast<std::uint32_t>(NumSpecConsts);
+    ZeSpecConstants.numConstants = static_cast<std::uint32_t>(spec_ids.size());
     ZeSpecConstants.pConstantIds = spec_ids.empty() ? nullptr : spec_ids.data();
     ZeSpecConstants.pConstantValues =
         spec_values.empty() ? nullptr : spec_values.data();
@@ -699,25 +711,30 @@ DPCTLKernelBundle_CreateFromSpirv(__dpctl_keep const DPCTLSyclContextRef CtxRef,
     context *SyclCtx = unwrap<context>(CtxRef);
     device *SyclDev = unwrap<device>(DevRef);
     // get the backend type
-    auto BE = SyclCtx->get_platform().get_backend();
-    switch (BE) {
-    case backend::opencl:
-        KBRef = _CreateKernelBundleWithIL_ocl_impl(*SyclCtx, *SyclDev, IL,
-                                                   length, CompileOpts,
-                                                   NumSpecConsts, SpecConsts);
-        break;
-    case backend::ext_oneapi_level_zero:
+    try {
+        auto BE = SyclCtx->get_platform().get_backend();
+        switch (BE) {
+        case backend::opencl:
+            KBRef = _CreateKernelBundleWithIL_ocl_impl(
+                *SyclCtx, *SyclDev, IL, 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,
-                                                  NumSpecConsts, SpecConsts);
-        break;
+            KBRef = _CreateKernelBundleWithIL_ze_impl(
+                *SyclCtx, *SyclDev, IL, length, CompileOpts, NumSpecConsts,
+                SpecConsts);
+            break;
 #endif
-    default:
-        std::ostringstream os;
-        os << "Backend " << BE << " is not supported";
-        error_handler(os.str(), __FILE__, __func__, __LINE__);
-        break;
+        default:
+            std::ostringstream os;
+            os << "Backend " << BE << " is not supported";
+            error_handler(os.str(), __FILE__, __func__, __LINE__);
+            break;
+        }
+    } catch (std::exception const &e) {
+        error_handler(e, __FILE__, __func__, __LINE__);
+        return nullptr;
     }
     return KBRef;
 }