diff --git a/oneapi-rs-sys/include/queue.hpp b/oneapi-rs-sys/include/queue.hpp index 8ac204a..635aaf6 100644 --- a/oneapi-rs-sys/include/queue.hpp +++ b/oneapi-rs-sys/include/queue.hpp @@ -46,4 +46,8 @@ std::unique_ptr launch_3d(std::unique_ptr &, Range3 global_size, Range3 local_size, Kernel const &, rust::Slice const> args); + +std::unique_ptr memcpy(std::unique_ptr &, std::uint8_t *dest, + std::uint8_t const *src, std::size_t num_bytes, + rust::Vec); } // namespace sycl_shims::queue diff --git a/oneapi-rs-sys/src/queue-sys.rs b/oneapi-rs-sys/src/queue-sys.rs index d66e620..5e77ef0 100644 --- a/oneapi-rs-sys/src/queue-sys.rs +++ b/oneapi-rs-sys/src/queue-sys.rs @@ -71,5 +71,13 @@ pub mod ffi { kernel: &Kernel, args: &[&[u8]], ) -> UniquePtr; + + unsafe fn memcpy( + queue: &mut UniquePtr, + dest: *mut u8, + src: *const u8, + num_bytes: usize, + dep_events: Vec, + ) -> UniquePtr; } } diff --git a/oneapi-rs-sys/src/queue.cpp b/oneapi-rs-sys/src/queue.cpp index afef17b..8fce66d 100644 --- a/oneapi-rs-sys/src/queue.cpp +++ b/oneapi-rs-sys/src/queue.cpp @@ -98,4 +98,14 @@ launch_3d(std::unique_ptr &queue, Range3 global_size, Range3 local_size, {local_size.data[0], local_size.data[1], local_size.data[2]}}, kernel, args); } + +std::unique_ptr memcpy(std::unique_ptr &queue, std::uint8_t *dest, + std::uint8_t const *src, std::size_t num_bytes, + rust::Vec dep_events) { + std::vector deps; + for (auto &&e : dep_events) + deps.push_back(std::move(*e.ptr)); + + return std::make_unique(queue->memcpy(dest, src, num_bytes, deps)); +} } // namespace sycl_shims::queue diff --git a/oneapi-rs/examples/kernel_launch.rs b/oneapi-rs/examples/kernel_launch.rs index c7fe823..10bf4f8 100644 --- a/oneapi-rs/examples/kernel_launch.rs +++ b/oneapi-rs/examples/kernel_launch.rs @@ -21,9 +21,10 @@ void iota(double start, double *ptr) { } "#; -fn main() { +#[tokio::main] +async fn main() { let mut queue = Queue::new(); - let mut buffer = queue.alloc_shared::(1024).wait(); + let mut device_buffer = queue.alloc_device::(1024).await; let kernel = queue .get_context() @@ -31,9 +32,20 @@ fn main() { .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::(1024).await; + + queue.copy(&device_buffer, &mut host_buffer).await; - for e in buffer.iter() { + for e in host_buffer.iter() { print!("{e} "); } println!(); diff --git a/oneapi-rs/src/buffer.rs b/oneapi-rs/src/buffer.rs index 6d67adf..71718ea 100644 --- a/oneapi-rs/src/buffer.rs +++ b/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,7 +61,7 @@ impl Buffer { } } - pub(crate) fn get_byte_ptr(&mut self) -> *mut u8 { + pub(crate) fn get_byte_ptr(&self) -> *mut u8 { self.data.as_ptr().cast() } @@ -69,6 +69,10 @@ impl Buffer { 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; @@ -76,14 +80,14 @@ impl Buffer { } } -impl Deref for Buffer { +impl Deref for Buffer { type Target = [T]; fn deref(&self) -> &Self::Target { unsafe { slice::from_raw_parts(self.data.as_ptr(), self.len) } } } -impl DerefMut for Buffer { +impl DerefMut for Buffer { fn deref_mut(&mut self) -> &mut Self::Target { unsafe { slice::from_raw_parts_mut(self.data.as_ptr(), self.len) } } diff --git a/oneapi-rs/src/queue.rs b/oneapi-rs/src/queue.rs index c55b765..635ad3c 100644 --- a/oneapi-rs/src/queue.rs +++ b/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( + &mut self, + len: usize, + ) -> EnqueuedBuffer> { + 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( @@ -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( + &self, + len: usize, + ) -> Buffer> { + 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( @@ -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(&mut self, src: &Buffer, dst: &mut Buffer) -> 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( + &mut self, + src: &Buffer, + dst: &mut Buffer, + 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::>(); + + let amount = min(src.get_len(), dst.get_len()); + let num_bytes = amount * size_of::(); + unsafe { + ffi::memcpy( + &mut self.0, + dst.get_byte_ptr(), + src.get_byte_ptr(), + num_bytes, + dep_events, + ) + } + .into() + } } impl From<&Device> for Queue { diff --git a/oneapi-rs/src/usm.rs b/oneapi-rs/src/usm.rs index 4e27c5e..518c5ab 100644 --- a/oneapi-rs/src/usm.rs +++ b/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 From<&Queue> for UsmAllocator { fn from(queue: &Queue) -> Self { Self { @@ -60,7 +63,7 @@ unsafe impl Allocator for UsmAllocator { /// 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 {} + /// 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 {}