Fix possible crashes with kernel objects

Getting a named kernel from a kernel bundle that does not contain it could cause a crash.

It was also possible for a kernel bundle wrapping a NULL reference, which could then be submitted and crash

Author: ndgrigorian

dpctl/compiler/_compiler.pyx

@@ -270,18 +270,45 @@ cdef class SyclKernelBundle:
         return self._kernel_bundle_ref
 
     cpdef SyclKernel get_sycl_kernel(self, str kernel_name):
-        name = kernel_name.encode("utf8")
+        """get_sycl_kernel(kernel_name)
+
+        Returns the kernel with the given name defined in this kernel bundle.
+
+        Args:
+            kernel_name (str):
+                Name of the kernel, as it appears in the kernel bundle. For a
+                bundle created from a SPIR-V binary this is the name of the
+                entry point, which for a kernel written in SYCL is the mangled
+                name of its kernel class, e.g.
+                ``'_ZTS20BasicSpecConstKernel'``.
+
+        Returns:
+            dpctl.compiler.SyclKernel:
+                The kernel object.
+
+        Raises:
+            ValueError:
+                If the kernel could not be created for any reason.
+        """
+        cdef bytes name = kernel_name.encode("utf8")
+        cdef DPCTLSyclKernelRef KRef = NULL
+
         if self._is_sycl_source:
-            return SyclKernel._create(
-                    DPCTLKernelBundle_GetSyclKernel(
-                        self._kernel_bundle_ref, name
-                    ),
-                    kernel_name
-                )
-        return SyclKernel._create(
-            DPCTLKernelBundle_GetKernel(self._kernel_bundle_ref, name),
-            kernel_name
-        )
+            KRef = DPCTLKernelBundle_GetSyclKernel(
+                self._kernel_bundle_ref, name
+            )
+        else:
+            KRef = DPCTLKernelBundle_GetKernel(self._kernel_bundle_ref, name)
+        if KRef is NULL:
+            # returning a SyclKernel with a NULL reference would defer the
+            # failure to the submission of the kernel
+            raise ValueError(
+                f"Kernel '{kernel_name}' could not be created from the kernel "
+                "bundle. The kernel bundle may not define a kernel with that "
+                "name, or the module it was built from may require device "
+                "libraries that were not linked into it."
+            )
+        return SyclKernel._create(KRef, kernel_name)
 
     def has_sycl_kernel(self, str kernel_name):
         name = kernel_name.encode("utf8")

dpctl/tests/test_sycl_compiler.py

@@ -233,6 +233,21 @@ def test_create_kernel_bundle_from_source_ocl():
     _check_multi_kernel_bundle(kb)
 
 
+def test_get_sycl_kernel_undefined_name():
+    try:
+        q = dpctl.SyclQueue()
+    except dpctl.SyclQueueCreationError:
+        pytest.skip("Could not create default queue")
+    spirv_file = get_spirv_abspath("multi_kernel.spv")
+    with open(spirv_file, "rb") as fin:
+        spirv = fin.read()
+    kb = dpc.create_kernel_bundle_from_spirv(q, spirv)
+
+    assert not kb.has_sycl_kernel("no_such_kernel")
+    with pytest.raises(ValueError):
+        kb.get_sycl_kernel("no_such_kernel")
+
+
 def test_create_kernel_bundle_from_spirv_ocl():
     q = _get_opencl_queue_or_skip()
     spirv_file = get_spirv_abspath("multi_kernel.spv")

libsyclinterface/source/dpctl_sycl_queue_interface.cpp

@@ -498,6 +498,17 @@ DPCTLQueue_SubmitRange(__dpctl_keep const DPCTLSyclKernelRef KRef,
                        __dpctl_keep const DPCTLSyclEventRef *DepEvents,
                        size_t NDepEvents)
 {
+    if (!KRef) {
+        error_handler("Cannot submit the kernel as the kernel is a nullptr.",
+                      __FILE__, __func__, __LINE__, error_level::error);
+        return nullptr;
+    }
+    if (!QRef) {
+        error_handler("Cannot submit the kernel as the queue is a nullptr.",
+                      __FILE__, __func__, __LINE__, error_level::error);
+        return nullptr;
+    }
+
     auto Kernel = unwrap<kernel>(KRef);
     auto Queue = unwrap<queue>(QRef);
     event e;
@@ -563,6 +574,17 @@ DPCTLQueue_SubmitNDRange(__dpctl_keep const DPCTLSyclKernelRef KRef,
                          __dpctl_keep const DPCTLSyclEventRef *DepEvents,
                          size_t NDepEvents)
 {
+    if (!KRef) {
+        error_handler("Cannot submit the kernel as the kernel is a nullptr.",
+                      __FILE__, __func__, __LINE__, error_level::error);
+        return nullptr;
+    }
+    if (!QRef) {
+        error_handler("Cannot submit the kernel as the queue is a nullptr.",
+                      __FILE__, __func__, __LINE__, error_level::error);
+        return nullptr;
+    }
+
     auto Kernel = unwrap<kernel>(KRef);
     auto Queue = unwrap<queue>(QRef);
     event e;

libsyclinterface/tests/test_sycl_queue_submit.cpp

@@ -389,6 +389,47 @@ TEST_F(TestQueueSubmit, CheckForUnsupportedArgTy)
     ASSERT_TRUE(ERef == nullptr);
 }
 
+TEST_F(TestQueueSubmit, CheckForNullKernel)
+{
+    std::size_t Range[] = {SIZE};
+    std::size_t lRange[] = {1};
+    std::size_t RANGE_NDIMS = 1;
+    DPCTLSyclEventRef ERef = nullptr;
+
+    EXPECT_NO_FATAL_FAILURE(
+        ERef = DPCTLQueue_SubmitRange(nullptr, QRef, nullptr, nullptr, 0, Range,
+                                      RANGE_NDIMS, nullptr, 0));
+    ASSERT_TRUE(ERef == nullptr);
+
+    EXPECT_NO_FATAL_FAILURE(ERef = DPCTLQueue_SubmitNDRange(
+                                nullptr, QRef, nullptr, nullptr, 0, Range,
+                                lRange, RANGE_NDIMS, nullptr, 0));
+    ASSERT_TRUE(ERef == nullptr);
+}
+
+TEST_F(TestQueueSubmit, CheckForNullQueue)
+{
+    std::size_t Range[] = {SIZE};
+    std::size_t lRange[] = {1};
+    std::size_t RANGE_NDIMS = 1;
+    DPCTLSyclEventRef ERef = nullptr;
+
+    auto kernel = DPCTLKernelBundle_GetKernel(KBRef, "_ZTS11RangeKernelIiE");
+    ASSERT_TRUE(kernel != nullptr);
+
+    EXPECT_NO_FATAL_FAILURE(
+        ERef = DPCTLQueue_SubmitRange(kernel, nullptr, nullptr, nullptr, 0,
+                                      Range, RANGE_NDIMS, nullptr, 0));
+    ASSERT_TRUE(ERef == nullptr);
+
+    EXPECT_NO_FATAL_FAILURE(ERef = DPCTLQueue_SubmitNDRange(
+                                kernel, nullptr, nullptr, nullptr, 0, Range,
+                                lRange, RANGE_NDIMS, nullptr, 0));
+    ASSERT_TRUE(ERef == nullptr);
+
+    EXPECT_NO_FATAL_FAILURE(DPCTLKernel_Delete(kernel));
+}
+
 struct TestQueueSubmitBarrier : public ::testing::Test
 {
     DPCTLSyclQueueRef QRef = nullptr;