Fix kernel launches on CPU devices

CHANGELOG.md

@@ -27,3 +27,6 @@ can be found at: <https://docs.rs/sycl_rs/>
 - Runtime kernel bundle compilation and one-, two-, and three-dimensional
   kernel launch support.
 - Derive support for typed kernel argument lists.
+- `RawKernelArg`, which marks USM pointer kernel arguments apart from plain
+  bytes, so kernels also launch on backends that take pointers through a
+  separate entry point, such as the OpenCL CPU device.

sycl/sycl-rs-derive/src/lib.rs

@@ -43,7 +43,7 @@ pub fn derive_kernel_argument_list(input: TokenStream) -> TokenStream {
     let expanded = quote! {
         unsafe impl #impl_generics #oneapi::kernel::KernelArgumentList<#argc>
         for #ident #ty_generics #where_clause {
-            unsafe fn as_raw_arg_list(&self) -> [&[u8]; #argc] {
+            unsafe fn as_raw_arg_list(&self) -> [#oneapi::kernel::RawKernelArg<'_>; #argc] {
                 [ #(unsafe { self.#members.as_raw_arg() }),* ]
             }
         }
@@ -69,7 +69,7 @@ fn get_single_tuple_impl(argc: usize) -> proc_macro2::TokenStream {
     quote! {
         unsafe impl<#(#types),*> crate::kernel::KernelArgumentList<#argc> for (#(#types),*)
         where #(#types: crate::kernel::KernelArgument),* {
-            unsafe fn as_raw_arg_list(&self) -> [&[u8]; #argc] {
+            unsafe fn as_raw_arg_list(&self) -> [crate::kernel::RawKernelArg<'_>; #argc] {
                 [ #(unsafe { self.#iter.as_raw_arg() }),* ]
             }
         }

sycl/sycl-rs-sys/include/queue.hpp

@@ -18,6 +18,7 @@ struct EventPtr;
 struct Range1;
 struct Range2;
 struct Range3;
+struct KernelArg;
 } // namespace sycl_shims
 
 namespace sycl_shims::queue {
@@ -37,20 +38,17 @@ std::unique_ptr<Event> memset(std::unique_ptr<Queue> &, std::uint8_t *ptr,
 std::unique_ptr<Event> barrier(std::unique_ptr<Queue> &, rust::Vec<EventPtr>);
 void wait(std::unique_ptr<Queue> &);
 
-std::unique_ptr<Event>
-launch_1d(std::unique_ptr<Queue> &, Range1 global_size, Range1 local_size,
-          Kernel const &,
-          rust::Slice<rust::slice<std::uint8_t const> const> args);
+std::unique_ptr<Event> launch_1d(std::unique_ptr<Queue> &, Range1 global_size,
+                                 Range1 local_size, Kernel const &,
+                                 rust::Slice<KernelArg const> args);
 
-std::unique_ptr<Event>
-launch_2d(std::unique_ptr<Queue> &, Range2 global_size, Range2 local_size,
-          Kernel const &,
-          rust::Slice<rust::slice<std::uint8_t const> const> args);
+std::unique_ptr<Event> launch_2d(std::unique_ptr<Queue> &, Range2 global_size,
+                                 Range2 local_size, Kernel const &,
+                                 rust::Slice<KernelArg const> args);
 
-std::unique_ptr<Event>
-launch_3d(std::unique_ptr<Queue> &, Range3 global_size, Range3 local_size,
-          Kernel const &,
-          rust::Slice<rust::slice<std::uint8_t const> const> args);
+std::unique_ptr<Event> launch_3d(std::unique_ptr<Queue> &, Range3 global_size,
+                                 Range3 local_size, Kernel const &,
+                                 rust::Slice<KernelArg const> args);
 
 std::unique_ptr<Event> memcpy(std::unique_ptr<Queue> &, std::uint8_t *dest,
                               std::uint8_t const *src, std::size_t num_bytes,

sycl/sycl-rs-sys/src/queue-sys.rs

@@ -17,6 +17,7 @@ pub mod ffi {
         type Range1 = crate::types::ffi::Range1;
         type Range2 = crate::types::ffi::Range2;
         type Range3 = crate::types::ffi::Range3;
+        type KernelArg<'a> = crate::types::ffi::KernelArg<'a>;
     }
 
     unsafe extern "C++" {
@@ -81,37 +82,37 @@ pub mod ffi {
         /// # Safety
         ///
         /// Each entry in `args` must match the corresponding kernel parameter's size, layout,
-        /// and alignment.
-        unsafe fn launch_1d(
+        /// and alignment, and must be marked as a pointer exactly when it holds a USM pointer.
+        unsafe fn launch_1d<'a>(
             queue: &mut UniquePtr<Queue>,
             global_size: Range1,
             local_size: Range1,
             kernel: &Kernel,
-            args: &[&[u8]],
+            args: &[KernelArg<'a>],
         ) -> Result<UniquePtr<Event>>;
 
         /// # Safety
         ///
         /// Each entry in `args` must match the corresponding kernel parameter's size, layout,
-        /// and alignment.
-        unsafe fn launch_2d(
+        /// and alignment, and must be marked as a pointer exactly when it holds a USM pointer.
+        unsafe fn launch_2d<'a>(
             queue: &mut UniquePtr<Queue>,
             global_size: Range2,
             local_size: Range2,
             kernel: &Kernel,
-            args: &[&[u8]],
+            args: &[KernelArg<'a>],
         ) -> Result<UniquePtr<Event>>;
 
         /// # Safety
         ///
         /// Each entry in `args` must match the corresponding kernel parameter's size, layout,
-        /// and alignment.
-        unsafe fn launch_3d(
+        /// and alignment, and must be marked as a pointer exactly when it holds a USM pointer.
+        unsafe fn launch_3d<'a>(
             queue: &mut UniquePtr<Queue>,
             global_size: Range3,
             local_size: Range3,
             kernel: &Kernel,
-            args: &[&[u8]],
+            args: &[KernelArg<'a>],
         ) -> Result<UniquePtr<Event>>;
     }
 }

sycl/sycl-rs-sys/src/queue.cpp

@@ -10,6 +10,9 @@
 #include "sycl-rs-sys/include/utils.hpp"
 #include "sycl-rs-sys/src/queue-sys.rs.h"
 
+#include <cstring>
+#include <stdexcept>
+
 using sycl::ext::intel::property::queue::immediate_command_list;
 using sycl::property::queue::in_order;
 using sycl_shims::utils::vec_to_vector;
@@ -62,42 +65,65 @@ std::unique_ptr<Event> barrier(std::unique_ptr<Queue> &queue,
 
 void wait(std::unique_ptr<Queue> &queue) { queue->wait(); }
 
+// Pointers go through the typed set_arg overload rather than raw_kernel_arg: a
+// backend may take them through a different entry point than plain bytes (the
+// OpenCL CPU device rejects USM pointers passed as bytes), and raw_kernel_arg
+// only takes non-pointer values.
+static void set_kernel_arg(sycl::handler &cgh, std::size_t index,
+                           KernelArg const &arg) {
+  switch (arg.kind) {
+  case KernelArgKind::Bytes:
+    cgh.set_arg(index,
+                syclexp::raw_kernel_arg(arg.bytes.data(), arg.bytes.size()));
+    return;
+  case KernelArgKind::Pointer: {
+    void *pointer;
+    if (arg.bytes.size() != sizeof(pointer))
+      throw std::invalid_argument(
+          "pointer kernel argument is not pointer-sized");
+    std::memcpy(&pointer, arg.bytes.data(), sizeof(pointer));
+    cgh.set_arg(index, pointer);
+    return;
+  }
+  }
+  throw std::invalid_argument("unknown kernel argument kind");
+}
+
 template <int Dimensions>
 std::unique_ptr<Event>
 launch(std::unique_ptr<Queue> &queue, sycl::nd_range<Dimensions> nd_range,
-       Kernel const &kernel,
-       rust::Slice<rust::slice<std::uint8_t const> const> args) {
+       Kernel const &kernel, rust::Slice<KernelArg const> args) {
   return std::make_unique<Event>(queue->submit([&](sycl::handler &cgh) {
     for (std::size_t i = 0; i < args.size(); ++i)
-      cgh.set_arg(i, syclexp::raw_kernel_arg(args[i].data(), args[i].size()));
+      set_kernel_arg(cgh, i, args[i]);
 
     cgh.parallel_for(nd_range, kernel);
   }));
 }
 
-std::unique_ptr<Event>
-launch_1d(std::unique_ptr<Queue> &queue, Range1 global_size, Range1 local_size,
-          Kernel const &kernel,
-          rust::Slice<rust::slice<std::uint8_t const> const> args) {
+std::unique_ptr<Event> launch_1d(std::unique_ptr<Queue> &queue,
+                                 Range1 global_size, Range1 local_size,
+                                 Kernel const &kernel,
+                                 rust::Slice<KernelArg const> args) {
   return launch(queue,
                 sycl::nd_range<1>{{global_size.data[0]}, {local_size.data[0]}},
                 kernel, args);
 }
 
-std::unique_ptr<Event>
-launch_2d(std::unique_ptr<Queue> &queue, Range2 global_size, Range2 local_size,
-          Kernel const &kernel,
-          rust::Slice<rust::slice<std::uint8_t const> const> args) {
+std::unique_ptr<Event> launch_2d(std::unique_ptr<Queue> &queue,
+                                 Range2 global_size, Range2 local_size,
+                                 Kernel const &kernel,
+                                 rust::Slice<KernelArg const> args) {
   return launch(queue,
                 sycl::nd_range<2>{{global_size.data[0], global_size.data[1]},
                                   {local_size.data[0], local_size.data[1]}},
                 kernel, args);
 }
 
-std::unique_ptr<Event>
-launch_3d(std::unique_ptr<Queue> &queue, Range3 global_size, Range3 local_size,
-          Kernel const &kernel,
-          rust::Slice<rust::slice<std::uint8_t const> const> args) {
+std::unique_ptr<Event> launch_3d(std::unique_ptr<Queue> &queue,
+                                 Range3 global_size, Range3 local_size,
+                                 Kernel const &kernel,
+                                 rust::Slice<KernelArg const> args) {
   return launch(
       queue,
       sycl::nd_range<3>{

sycl/sycl-rs-sys/src/types-sys.rs

@@ -106,6 +106,22 @@ pub mod ffi {
         data: [u64; 3],
     }
 
+    /// How the bytes of a kernel argument are handed to the backend. A backend may take
+    /// pointers through a different entry point than plain bytes, and cannot tell the two
+    /// apart from the bytes alone.
+    #[derive(Debug, Hash)]
+    enum KernelArgKind {
+        /// The bytes of a trivially copyable value that is not a pointer.
+        Bytes,
+        /// The bytes of a pointer to USM memory.
+        Pointer,
+    }
+
+    struct KernelArg<'a> {
+        kind: KernelArgKind,
+        bytes: &'a [u8],
+    }
+
     impl UniquePtr<Device> {}
     impl UniquePtr<Platform> {}
     impl UniquePtr<Queue> {}

sycl/sycl-rs/src/kernel.rs

@@ -49,25 +49,50 @@ impl From<cxx::UniquePtr<types::ffi::Kernel>> for Kernel {
     }
 }
 
+/// The raw representation of a SYCL kernel argument.
+///
+/// A backend may take pointers through a different entry point than plain bytes, and cannot
+/// tell the two apart from the bytes alone, so pointers are marked as such.
+#[derive(Clone, Copy, Debug)]
+#[non_exhaustive]
+pub enum RawKernelArg<'a> {
+    /// The bytes of a trivially copyable value that is not a pointer.
+    Bytes(&'a [u8]),
+    /// The bytes of a pointer to USM memory.
+    Pointer(&'a [u8]),
+}
+
+impl<'a> RawKernelArg<'a> {
+    pub(crate) fn into_ffi(self) -> types::ffi::KernelArg<'a> {
+        let (kind, bytes) = match self {
+            Self::Bytes(bytes) => (types::ffi::KernelArgKind::Bytes, bytes),
+            Self::Pointer(bytes) => (types::ffi::KernelArgKind::Pointer, bytes),
+        };
+        types::ffi::KernelArg { kind, bytes }
+    }
+}
+
 /// Types which can be passed as SYCL kernel arguments.
 ///
 /// # Safety
 ///
 /// A type implementing this trait must mirror the representation and alignment of the
-/// corresponding SYCL kernel argument structure.
+/// corresponding SYCL kernel argument structure, and must return [`RawKernelArg::Pointer`]
+/// exactly when it represents a pointer to USM memory.
 pub unsafe trait KernelArgument {
-    /// Converts self to a raw byte representation.
+    /// Converts self to its raw representation.
     ///
     /// # Safety
     ///
-    /// This function returns a reference to raw bytes. These bytes will be passed to FFI
-    /// functions. The caller must make sure these functions respect Rust's aliasing rules.
-    unsafe fn as_raw_arg(&self) -> &[u8];
+    /// This function returns a reference to raw bytes, marked as a pointer or plain bytes.
+    /// These bytes will be passed to FFI functions. The caller must make sure these functions
+    /// respect Rust's aliasing rules.
+    unsafe fn as_raw_arg(&self) -> RawKernelArg<'_>;
 }
 
 unsafe impl<T: Pod> KernelArgument for T {
-    unsafe fn as_raw_arg(&self) -> &[u8] {
-        bytemuck::bytes_of(self)
+    unsafe fn as_raw_arg(&self) -> RawKernelArg<'_> {
+        RawKernelArg::Bytes(bytemuck::bytes_of(self))
     }
 }
 
@@ -76,25 +101,27 @@ unsafe impl<T: Pod> KernelArgument for T {
 /// # Safety
 ///
 /// A type implementing this trait must mirror the representation and alignment of each
-/// corresponding SYCL kernel argument inside the returned array.
+/// corresponding SYCL kernel argument inside the returned array, and each entry must be
+/// [`RawKernelArg::Pointer`] exactly when it represents a pointer to USM memory.
 pub unsafe trait KernelArgumentList<const ARGC: usize> {
-    /// Converts each struct member to a raw byte representation.
+    /// Converts each struct member to its raw representation.
     ///
     /// # Safety
     ///
-    /// This function returns references to raw bytes. These bytes will be passed to FFI
-    /// functions. The caller must make sure these functions respect Rust's aliasing rules.
-    unsafe fn as_raw_arg_list(&self) -> [&[u8]; ARGC];
+    /// This function returns references to raw bytes, each marked as a pointer or plain bytes.
+    /// These bytes will be passed to FFI functions. The caller must make sure these functions
+    /// respect Rust's aliasing rules.
+    unsafe fn as_raw_arg_list(&self) -> [RawKernelArg<'_>; ARGC];
 }
 
 unsafe impl KernelArgumentList<0> for () {
-    unsafe fn as_raw_arg_list(&self) -> [&[u8]; 0] {
+    unsafe fn as_raw_arg_list(&self) -> [RawKernelArg<'_>; 0] {
         []
     }
 }
 
 unsafe impl<T: KernelArgument> KernelArgumentList<1> for T {
-    unsafe fn as_raw_arg_list(&self) -> [&[u8]; 1] {
+    unsafe fn as_raw_arg_list(&self) -> [RawKernelArg<'_>; 1] {
         [unsafe { self.as_raw_arg() }]
     }
 }

sycl/sycl-rs/src/prelude.rs

@@ -10,7 +10,7 @@ pub use crate::{
     context::Context,
     device::{Aspect, Device},
     info::{self, InfoTarget},
-    kernel::{Kernel, KernelArgument, KernelArgumentList},
+    kernel::{Kernel, KernelArgument, KernelArgumentList, RawKernelArg},
     platform::Platform,
     queue::Queue,
     range::{NdRange, Range},

sycl/sycl-rs/src/range.rs

@@ -11,7 +11,7 @@ use sycl_rs_sys::types;
 
 use crate::{
     event::Event,
-    kernel::{Kernel, KernelArgumentList},
+    kernel::{Kernel, KernelArgumentList, RawKernelArg},
     private::Sealed,
     queue::Queue,
 };
@@ -72,7 +72,7 @@ impl ValidDimension for NdRange<1> {
                     data: self.local_size,
                 },
                 &kernel.0,
-                &args.as_raw_arg_list(),
+                &args.as_raw_arg_list().map(RawKernelArg::into_ffi),
             )
         }
         .map(Into::into)
@@ -97,7 +97,7 @@ impl ValidDimension for NdRange<2> {
                     data: self.local_size,
                 },
                 &kernel.0,
-                &args.as_raw_arg_list(),
+                &args.as_raw_arg_list().map(RawKernelArg::into_ffi),
             )
         }
         .map(Into::into)
@@ -122,7 +122,7 @@ impl ValidDimension for NdRange<3> {
                     data: self.local_size,
                 },
                 &kernel.0,
-                &args.as_raw_arg_list(),
+                &args.as_raw_arg_list().map(RawKernelArg::into_ffi),
             )
         }
         .map(Into::into)