Add Rust NdRange support

Author: szymon-zadworny

oneapi-rs/examples/kernel_launch.rs

@@ -6,7 +6,7 @@
 // SPDX-License-Identifier: MIT OR Apache-2.0
 //
 
-use oneapi_rs::{buffer::Buffer, kernel_bundle::{KernelArgument, KernelArgumentList}, queue::Queue, usm::{SharedAllocator, UsmAllocator}};
+use oneapi_rs::{buffer::Buffer, kernel_bundle::{KernelArgument, KernelArgumentList, NdRange, Range}, queue::Queue, usm::{SharedAllocator, UsmAllocator}};
 
 static IOTA_SRC: &str = r#"
 #include <sycl/sycl.hpp>
@@ -44,7 +44,7 @@ fn main() {
         .build()
         .get_kernel("iota");
 
-    unsafe { queue.launch(&kernel, IotaArgs { start: 3.14, buffer: &mut buffer }) }.wait();
+    unsafe { queue.launch(NdRange::new([1024], [16]), &kernel, IotaArgs { start: 3.14, buffer: &mut buffer }) }.wait();
 
     for e in buffer.iter() {
         print!("{e} ");

oneapi-rs/src/kernel_bundle.rs

@@ -6,9 +6,13 @@
 // SPDX-License-Identifier: MIT OR Apache-2.0
 //
 
+use std::marker::PhantomData;
+
 use bytemuck::Pod;
 use oneapi_rs_sys::{kernel_bundle::ffi, types};
 
+use crate::{event::Event, queue::Queue};
+
 pub struct SourceKernelBundle(pub(crate) cxx::UniquePtr<types::ffi::SourceKernelBundle>);
 
 impl From<cxx::UniquePtr<types::ffi::SourceKernelBundle>> for SourceKernelBundle {
@@ -58,3 +62,80 @@ unsafe impl<T: Pod> KernelArgument for T {
 pub unsafe trait KernelArgumentList<const ARGC: usize> {
     unsafe fn as_raw_arg_list(&self) -> [&[u8]; ARGC];
 }
+
+
+pub type Range<const DIMENSIONS: usize = 1> = [u64; DIMENSIONS];
+
+pub struct NdRange<const DIMENSIONS: usize = 1> {
+    pub group_size: Range<DIMENSIONS>,
+    pub local_size: Range<DIMENSIONS>
+}
+
+impl<const DIMENSIONS: usize> NdRange<DIMENSIONS> {
+    pub fn new(group_size: Range<DIMENSIONS>, local_size: Range<DIMENSIONS>) -> Self {
+        Self {
+            group_size,
+            local_size
+        }
+    }
+}
+
+pub trait ValidDimension {
+    unsafe fn launch<const ARGC: usize>(&self, queue: &mut Queue, kernel: &Kernel, args: impl KernelArgumentList<ARGC>) -> Event;
+}
+
+impl ValidDimension for NdRange<1> {
+    unsafe fn launch<const ARGC: usize>(&self, queue: &mut Queue, kernel: &Kernel, args: impl KernelArgumentList<ARGC>) -> Event {
+        unsafe {
+            oneapi_rs_sys::queue::ffi::launch_1d(
+                &mut queue.0,
+                self.group_size[0],
+                self.local_size[0],
+                &kernel.0,
+                &args.as_raw_arg_list()
+            )
+        }.into()
+    }
+}
+
+impl ValidDimension for NdRange<2> {
+    unsafe fn launch<const ARGC: usize>(&self, queue: &mut Queue, kernel: &Kernel, args: impl KernelArgumentList<ARGC>) -> Event {
+        unsafe {
+            oneapi_rs_sys::queue::ffi::launch_2d(
+                &mut queue.0,
+                types::ffi::Range2 {
+                    x: self.group_size[0],
+                    y: self.group_size[1],
+                },
+                types::ffi::Range2 {
+                    x: self.local_size[0],
+                    y: self.local_size[1],
+                },
+                &kernel.0,
+                &args.as_raw_arg_list()
+            )
+        }.into()
+    }
+}
+
+impl ValidDimension for NdRange<3> {
+    unsafe fn launch<const ARGC: usize>(&self, queue: &mut Queue, kernel: &Kernel, args: impl KernelArgumentList<ARGC>) -> Event {
+        unsafe {
+            oneapi_rs_sys::queue::ffi::launch_3d(
+                &mut queue.0,
+                types::ffi::Range3 {
+                    x: self.group_size[0],
+                    y: self.group_size[1],
+                    z: self.group_size[2],
+                },
+                types::ffi::Range3 {
+                    x: self.local_size[0],
+                    y: self.local_size[1],
+                    z: self.local_size[2],
+                },
+                &kernel.0,
+                &args.as_raw_arg_list()
+            )
+        }.into()
+    }
+}

oneapi-rs/src/queue.rs

@@ -10,7 +10,12 @@ use bytemuck::Pod;
 use oneapi_rs_sys::{queue::ffi, types::ffi::EventPtr};
 
 use crate::{
-    buffer::{Buffer, EnqueuedBuffer}, context::Context, device::Device, event::Event, kernel_bundle::{Kernel, KernelArgumentList}, usm::{HostAllocator, SharedAllocator, UsmAlloc, UsmAllocator},
+    buffer::{Buffer, EnqueuedBuffer},
+    context::Context,
+    device::Device,
+    event::Event,
+    kernel_bundle::{Kernel, KernelArgumentList, NdRange, ValidDimension},
+    usm::{HostAllocator, SharedAllocator, UsmAlloc, UsmAllocator},
 };
 
 /// The `Queue` connects a host program to a single device. Programs submit tasks to a device via the
@@ -128,8 +133,16 @@ impl Queue {
         ffi::wait(&mut self.0);
     }
 
-    pub unsafe fn launch<const ARGC: usize>(&mut self, kernel: &Kernel, args: impl KernelArgumentList<ARGC>) -> Event {
-        unsafe { ffi::launch(&mut self.0, &kernel.0, &args.as_raw_arg_list()) }.into()
+    pub unsafe fn launch<const ARGC: usize, const DIMENSIONS: usize>(
+        &mut self,
+        nd_range: NdRange<DIMENSIONS>,
+        kernel: &Kernel,
+        args: impl KernelArgumentList<ARGC>,
+    ) -> Event
+    where
+        NdRange<DIMENSIONS>: ValidDimension,
+    {
+        unsafe { nd_range.launch(self, kernel, args) }
     }
 }