Add C++ multidimensional launch bindings

Author: szymon-zadworny

oneapi-rs-sys/include/queue.hpp

@@ -15,6 +15,8 @@
 
 namespace sycl_shims {
 struct EventPtr;
+struct Range2;
+struct Range3;
 } // namespace sycl_shims
 
 namespace sycl_shims::queue {
@@ -28,5 +30,16 @@ std::unique_ptr<Event> memset(std::unique_ptr<Queue> &, std::uint8_t *ptr,
                               rust::Vec<EventPtr>);
 std::unique_ptr<Event> barrier(std::unique_ptr<Queue> &, rust::Vec<EventPtr>);
 void wait(std::unique_ptr<Queue> &);
-std::unique_ptr<Event> launch(std::unique_ptr<Queue> &, Kernel const &, rust::Slice<rust::slice<std::uint8_t const> const> args);
+
+std::unique_ptr<Event>
+launch_1d(std::unique_ptr<Queue> &, unsigned long global_size, unsigned long 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<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<rust::slice<std::uint8_t const> const> args);
 } // namespace sycl_shims::queue

oneapi-rs-sys/src/queue-sys.rs

@@ -12,6 +12,8 @@ pub mod ffi {
     extern "C++" {
         include!("oneapi-rs-sys/src/types-sys.rs.h");
         type EventPtr = crate::types::ffi::EventPtr;
+        type Range2 = crate::types::ffi::Range2;
+        type Range3 = crate::types::ffi::Range3;
     }
 
     unsafe extern "C++" {
@@ -33,15 +35,40 @@ pub mod ffi {
         fn new_queue_from_device(device: &Device) -> UniquePtr<Queue>;
         fn get_context(queue: &Queue) -> UniquePtr<Context>;
         fn clone(queue: &Queue) -> UniquePtr<Queue>;
+
         unsafe fn memset(
             queue: &mut UniquePtr<Queue>,
             ptr: *mut u8,
             value: i32,
             num_bytes: usize,
             dep_events: Vec<EventPtr>,
         ) -> UniquePtr<Event>;
+
         fn barrier(queue: &mut UniquePtr<Queue>, dep_events: Vec<EventPtr>) -> UniquePtr<Event>;
         fn wait(queue: &mut UniquePtr<Queue>);
-        unsafe fn launch(queue: &mut UniquePtr<Queue>, kernel: &Kernel, args: &[&[u8]]) -> UniquePtr<Event>;
+
+        unsafe fn launch_1d(
+            queue: &mut UniquePtr<Queue>,
+            global_size: u64,
+            local_size: u64,
+            kernel: &Kernel,
+            args: &[&[u8]],
+        ) -> UniquePtr<Event>;
+
+        unsafe fn launch_2d(
+            queue: &mut UniquePtr<Queue>,
+            global_size: Range2,
+            local_size: Range2,
+            kernel: &Kernel,
+            args: &[&[u8]],
+        ) -> UniquePtr<Event>;
+
+        unsafe fn launch_3d(
+            queue: &mut UniquePtr<Queue>,
+            global_size: Range3,
+            local_size: Range3,
+            kernel: &Kernel,
+            args: &[&[u8]],
+        ) -> UniquePtr<Event>;
     }
 }

oneapi-rs-sys/src/queue.cpp

@@ -55,12 +55,44 @@ std::unique_ptr<Event> barrier(std::unique_ptr<Queue> &queue,
 
 void wait(std::unique_ptr<Queue> &queue) { queue->wait(); }
 
-std::unique_ptr<Event> launch(std::unique_ptr<Queue> &queue, Kernel const &kernel, rust::Slice<rust::slice<std::uint8_t const> const> args) {
+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) {
   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()));
-    
-    cgh.parallel_for(sycl::nd_range{{1024}, {16}}, kernel);
+
+    cgh.parallel_for(nd_range, kernel);
   }));
 }
+
+std::unique_ptr<Event>
+launch_1d(std::unique_ptr<Queue> &queue, unsigned long global_size,
+          unsigned long local_size, Kernel const &kernel,
+          rust::Slice<rust::slice<std::uint8_t const> const> args) {
+  return launch(queue, sycl::nd_range<1>{{global_size}, {local_size}}, 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) {
+  return launch(queue,
+                sycl::nd_range<2>{{global_size.x, global_size.y},
+                                  {local_size.x, local_size.y}},
+                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) {
+  return launch(queue,
+                sycl::nd_range<3>{{global_size.x, global_size.y, global_size.z},
+                                  {local_size.x, local_size.y, local_size.z}},
+                kernel, args);
+}
 } // namespace sycl_shims::queue

oneapi-rs-sys/src/types-sys.rs

@@ -77,6 +77,17 @@ pub mod ffi {
         Complete,
         Unknown,
     }
+    
+    struct Range2 {
+        x: u64,
+        y: u64,
+    }
+
+    struct Range3 {
+        x: u64,
+        y: u64,
+        z: u64,
+    }
 
     impl UniquePtr<Device> {}
     impl UniquePtr<Platform> {}