Add memcpy binding

oneapi-rs-sys/include/queue.hpp

@@ -46,4 +46,8 @@ 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> memcpy(std::unique_ptr<Queue> &, std::uint8_t *dest,
+                              std::uint8_t const *src, std::size_t num_bytes,
+                              rust::Vec<EventPtr>);
 } // namespace sycl_shims::queue

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

@@ -71,5 +71,13 @@ pub mod ffi {
             kernel: &Kernel,
             args: &[&[u8]],
         ) -> UniquePtr<Event>;
+
+        unsafe fn memcpy(
+            queue: &mut UniquePtr<Queue>,
+            dest: *mut u8,
+            src: *const u8,
+            num_bytes: usize,
+            dep_events: Vec<EventPtr>,
+        ) -> UniquePtr<Event>;
     }
 }

oneapi-rs-sys/src/queue.cpp

@@ -98,4 +98,14 @@ launch_3d(std::unique_ptr<Queue> &queue, Range3 global_size, Range3 local_size,
           {local_size.data[0], local_size.data[1], local_size.data[2]}},
       kernel, args);
 }
+
+std::unique_ptr<Event> memcpy(std::unique_ptr<Queue> &queue, std::uint8_t *dest,
+                              std::uint8_t const *src, std::size_t num_bytes,
+                              rust::Vec<EventPtr> dep_events) {
+  std::vector<sycl::event> deps;
+  for (auto &&e : dep_events)
+    deps.push_back(std::move(*e.ptr));
+
+  return std::make_unique<Event>(queue->memcpy(dest, src, num_bytes, deps));
+}
 } // namespace sycl_shims::queue

oneapi-rs/examples/kernel_launch.rs

@@ -21,19 +21,31 @@ void iota(double start, double *ptr) {
 }
 "#;
 
-fn main() {
+#[tokio::main]
+async fn main() {
     let mut queue = Queue::new();
-    let mut buffer = queue.alloc_shared::<f64>(1024).wait();
+    let mut device_buffer = queue.alloc_device::<f64>(1024).await;
 
     let kernel = queue
         .get_context()
         .create_kernel_bundle_from_source(IOTA_SRC)
         .build()
         .get_kernel("iota");
 
-    unsafe { queue.launch(NdRange::new([1024], [16]), &kernel, (3.14, &mut buffer)) }.wait();
+    unsafe {
+        queue.launch(
+            NdRange::new([1024], [16]),
+            &kernel,
+            (3.14, &mut device_buffer),
+        )
+    }
+    .await;
+
+    let mut host_buffer = queue.alloc_host::<f64>(1024).await;
+
+    queue.copy(&device_buffer, &mut host_buffer).await;
 
-    for e in buffer.iter() {
+    for e in host_buffer.iter() {
         print!("{e} ");
     }
     println!();

oneapi-rs/src/buffer.rs

@@ -21,7 +21,7 @@ use pin_project::pin_project;
 use crate::{
     event::{Event, EventFuture},
     kernel::KernelArgument,
-    usm::UsmAlloc,
+    usm::{HostAccessible, UsmAlloc},
 };
 
 /// The Buffer struct defines a shared array of one, two or three dimensions that can be used
@@ -61,29 +61,33 @@ impl<T, A: UsmAlloc> Buffer<T, A> {
         }
     }
 
-    pub(crate) fn get_byte_ptr(&mut self) -> *mut u8 {
+    pub(crate) fn get_byte_ptr(&self) -> *mut u8 {
         self.data.as_ptr().cast()
     }
 
     pub(crate) fn get_byte_size(&self) -> usize {
         self.layout.size()
     }
 
+    pub(crate) fn get_len(&self) -> usize {
+        self.len
+    }
+
     unsafe fn as_raw_arg_impl(&self) -> &[u8] {
         let data_ptr: *const NonNull<_> = &self.data;
         let cast_ptr = data_ptr as *const u8;
         unsafe { slice::from_raw_parts(cast_ptr, std::mem::size_of_val(&cast_ptr)) }
     }
 }
 
-impl<T, A: UsmAlloc> Deref for Buffer<T, A> {
+impl<T, A: UsmAlloc + HostAccessible> Deref for Buffer<T, A> {
     type Target = [T];
     fn deref(&self) -> &Self::Target {
         unsafe { slice::from_raw_parts(self.data.as_ptr(), self.len) }
     }
 }
 
-impl<T, A: UsmAlloc> DerefMut for Buffer<T, A> {
+impl<T, A: UsmAlloc + HostAccessible> DerefMut for Buffer<T, A> {
     fn deref_mut(&mut self) -> &mut Self::Target {
         unsafe { slice::from_raw_parts_mut(self.data.as_ptr(), self.len) }
     }

oneapi-rs/src/queue.rs

@@ -6,6 +6,8 @@
 // SPDX-License-Identifier: MIT OR Apache-2.0
 //
 
+use std::cmp::min;
+
 use bytemuck::Pod;
 use oneapi_rs_sys::{queue::ffi, types::ffi::EventPtr};
 
@@ -16,7 +18,7 @@ use crate::{
     event::Event,
     kernel::{Kernel, KernelArgumentList},
     range::{NdRange, ValidDimension},
-    usm::{HostAllocator, SharedAllocator, UsmAlloc, UsmAllocator},
+    usm::{DeviceAllocator, HostAllocator, SharedAllocator, UsmAlloc, UsmAllocator},
 };
 
 /// The `Queue` connects a host program to a single device. Programs submit tasks to a device via the
@@ -64,6 +66,18 @@ impl Queue {
         }
     }
 
+    /// Allocates zeroed memory and creates a device [`Buffer`] that can store an array of T.
+    pub fn alloc_device<T: Pod>(
+        &mut self,
+        len: usize,
+    ) -> EnqueuedBuffer<T, UsmAllocator<DeviceAllocator>> {
+        unsafe {
+            let mut buffer = self.alloc_uninit_device(len);
+            let event = self.memset(&mut buffer, 0);
+            EnqueuedBuffer::new(buffer, event)
+        }
+    }
+
     /// Allocates memory and creates a host-side [`Buffer`] that can store an array of T.
     /// Safety: the buffer contents are uninitialized.
     pub unsafe fn alloc_uninit_host<T>(
@@ -84,6 +98,16 @@ impl Queue {
         unsafe { Buffer::new(allocator, len) }
     }
 
+    /// Allocates memory and creates a device-side [`Buffer`] that can store an array of T.
+    /// Safety: the buffer contents are uninitialized.
+    pub unsafe fn alloc_uninit_device<T>(
+        &self,
+        len: usize,
+    ) -> Buffer<T, UsmAllocator<DeviceAllocator>> {
+        let allocator = UsmAllocator::from(self);
+        unsafe { Buffer::new(allocator, len) }
+    }
+
     /// Sets memory allocated with USM allocations.
     /// Safety: the caller must make sure the underlying memory isn't being aliased somewhere else.
     pub unsafe fn memset<T, A: UsmAlloc>(
@@ -147,6 +171,57 @@ impl Queue {
     {
         unsafe { nd_range.launch(self, kernel, args) }
     }
+
+    /// Copies the contents of the source buffer to the destination buffer.
+    ///
+    /// If the buffer sizes don't match the destination buffer is filled with as many elements from
+    /// source buffer as possible.
+    pub fn copy<T, A1, A2>(&mut self, src: &Buffer<T, A1>, dst: &mut Buffer<T, A2>) -> Event
+    where
+        T: Pod,
+        A1: UsmAlloc,
+        A2: UsmAlloc,
+    {
+        self.copy_with_deps(src, dst, &[])
+    }
+
+    /// Copies the contents of the source buffer to the destination buffer after all specified
+    /// events finish.
+    ///
+    /// If the buffer sizes don't match the destination buffer is filled with as many elements from
+    /// source buffer as possible.
+    pub fn copy_with_deps<T, A1, A2>(
+        &mut self,
+        src: &Buffer<T, A1>,
+        dst: &mut Buffer<T, A2>,
+        dep_events: &[&Event],
+    ) -> Event
+    where
+        T: Pod,
+        A1: UsmAlloc,
+        A2: UsmAlloc,
+    {
+        // TODO: Resolve the C++ lifetime elision issue
+        let dep_events = dep_events
+            .iter()
+            .map(|e| EventPtr {
+                ptr: (*e).clone().0,
+            })
+            .collect::<Vec<_>>();
+
+        let amount = min(src.get_len(), dst.get_len());
+        let num_bytes = amount * size_of::<T>();
+        unsafe {
+            ffi::memcpy(
+                &mut self.0,
+                dst.get_byte_ptr(),
+                src.get_byte_ptr(),
+                num_bytes,
+                dep_events,
+            )
+        }
+        .into()
+    }
 }
 
 impl From<&Device> for Queue {

oneapi-rs/src/usm.rs

@@ -30,6 +30,9 @@ pub trait UsmAllocatorKind {
     unsafe fn alloc(alignment: usize, num_bytes: usize, queue: &Queue) -> CxxResult<*mut u8>;
 }
 
+/// A marker trait for host-accessible USM allocators.
+pub unsafe trait HostAccessible {}
+
 impl<T: UsmAllocatorKind> From<&Queue> for UsmAllocator<T> {
     fn from(queue: &Queue) -> Self {
         Self {
@@ -60,7 +63,7 @@ unsafe impl<T: UsmAllocatorKind> Allocator for UsmAllocator<T> {
 /// An allocator for Device-side buffers
 /// Safety: memory allocated by this allocator cannot be accessed on the host side
 #[allow(dead_code)]
-pub(crate) struct DeviceAllocator;
+pub struct DeviceAllocator;
 
 impl UsmAllocatorKind for DeviceAllocator {
     unsafe fn alloc(alignment: usize, num_bytes: usize, queue: &Queue) -> CxxResult<*mut u8> {
@@ -77,6 +80,8 @@ impl UsmAllocatorKind for HostAllocator {
     }
 }
 
+unsafe impl HostAccessible for UsmAllocator<HostAllocator> {}
+
 /// An allocator for shared memory buffers
 pub struct SharedAllocator;
 
@@ -85,3 +90,5 @@ impl UsmAllocatorKind for SharedAllocator {
         unsafe { ffi::aligned_alloc_shared(alignment, num_bytes, &queue.0) }
     }
 }
+
+unsafe impl HostAccessible for UsmAllocator<SharedAllocator> {}