diff --git a/oneapi-rs/examples/kernel_launch_derive.rs b/oneapi-rs/examples/kernel_launch_derive.rs index 39dc0d7..d3aee58 100644 --- a/oneapi-rs/examples/kernel_launch_derive.rs +++ b/oneapi-rs/examples/kernel_launch_derive.rs @@ -8,12 +8,7 @@ use oneapi_rs::prelude::*; -use oneapi_rs::{ - buffer::Buffer, - queue::Queue, - range::NdRange, - usm::{SharedAllocator, UsmAllocator}, -}; +use oneapi_rs::{buffer::SharedBuffer, queue::Queue, range::NdRange}; static IOTA_SRC: &str = r#" #include @@ -31,7 +26,7 @@ void iota(float start, float *ptr) { #[derive(KernelArgumentList)] struct IotaArgs<'a> { start: f32, - ptr: &'a mut Buffer>, + ptr: &'a mut SharedBuffer, } fn main() { diff --git a/oneapi-rs/src/buffer.rs b/oneapi-rs/src/buffer.rs index 71718ea..fb4cfae 100644 --- a/oneapi-rs/src/buffer.rs +++ b/oneapi-rs/src/buffer.rs @@ -21,7 +21,9 @@ use pin_project::pin_project; use crate::{ event::{Event, EventFuture}, kernel::KernelArgument, - usm::{HostAccessible, UsmAlloc}, + usm::{ + DeviceAllocator, HostAccessible, HostAllocator, SharedAllocator, UsmAlloc, UsmAllocator, + }, }; /// The Buffer struct defines a shared array of one, two or three dimensions that can be used @@ -101,6 +103,10 @@ impl Drop for Buffer { } } +pub type HostBuffer = Buffer>; +pub type SharedBuffer = Buffer>; +pub type DeviceBuffer = Buffer>; + /// A [`Buffer`] whose initialization has been enqueued. You need to wait/await it. pub struct EnqueuedBuffer { buffer: Buffer, @@ -121,6 +127,10 @@ impl EnqueuedBuffer { } } +pub type EnqueuedHostBuffer = EnqueuedBuffer>; +pub type EnqueuedSharedBuffer = EnqueuedBuffer>; +pub type EnqueuedDeviceBuffer = EnqueuedBuffer>; + #[pin_project] /// A [`Future`] which represents a pending [`Buffer`] allocation. pub struct BufferFuture { @@ -151,6 +161,10 @@ impl IntoFuture for EnqueuedBuffer { } } +pub type HostBufferFuture = BufferFuture>; +pub type SharedBufferFuture = BufferFuture>; +pub type DeviceBufferFuture = BufferFuture>; + unsafe impl KernelArgument for Buffer { unsafe fn as_raw_arg(&self) -> &[u8] { unsafe { self.as_raw_arg_impl() } diff --git a/oneapi-rs/src/queue.rs b/oneapi-rs/src/queue.rs index 635ad3c..859d507 100644 --- a/oneapi-rs/src/queue.rs +++ b/oneapi-rs/src/queue.rs @@ -12,13 +12,16 @@ use bytemuck::Pod; use oneapi_rs_sys::{queue::ffi, types::ffi::EventPtr}; use crate::{ - buffer::{Buffer, EnqueuedBuffer}, + buffer::{ + Buffer, DeviceBuffer, EnqueuedBuffer, EnqueuedDeviceBuffer, EnqueuedHostBuffer, + EnqueuedSharedBuffer, HostBuffer, SharedBuffer, + }, context::Context, device::Device, event::Event, kernel::{Kernel, KernelArgumentList}, range::{NdRange, ValidDimension}, - usm::{DeviceAllocator, HostAllocator, SharedAllocator, UsmAlloc, UsmAllocator}, + usm::{UsmAlloc, UsmAllocator}, }; /// The `Queue` connects a host program to a single device. Programs submit tasks to a device via the @@ -43,10 +46,7 @@ impl Queue { } /// Allocates zeroed memory and creates a host-side [`Buffer`] that can store an array of T. - pub fn alloc_host( - &mut self, - len: usize, - ) -> EnqueuedBuffer> { + pub fn alloc_host(&mut self, len: usize) -> EnqueuedHostBuffer { unsafe { let mut buffer = self.alloc_uninit_host(len); let event = self.memset(&mut buffer, 0); @@ -55,10 +55,7 @@ impl Queue { } /// Allocates zeroed memory and creates a shared [`Buffer`] that can store an array of T. - pub fn alloc_shared( - &mut self, - len: usize, - ) -> EnqueuedBuffer> { + pub fn alloc_shared(&mut self, len: usize) -> EnqueuedSharedBuffer { unsafe { let mut buffer = self.alloc_uninit_shared(len); let event = self.memset(&mut buffer, 0); @@ -67,10 +64,7 @@ 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> { + pub fn alloc_device(&mut self, len: usize) -> EnqueuedDeviceBuffer { unsafe { let mut buffer = self.alloc_uninit_device(len); let event = self.memset(&mut buffer, 0); @@ -80,30 +74,21 @@ impl Queue { /// 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( - &self, - len: usize, - ) -> Buffer> { + pub unsafe fn alloc_uninit_host(&self, len: usize) -> HostBuffer { let allocator = UsmAllocator::from(self); unsafe { Buffer::new(allocator, len) } } /// Allocates memory and creates a shared [`Buffer`] that can store an array of T. /// Safety: the buffer contents are uninitialized. - pub unsafe fn alloc_uninit_shared( - &self, - len: usize, - ) -> Buffer> { + pub unsafe fn alloc_uninit_shared(&self, len: usize) -> SharedBuffer { let allocator = UsmAllocator::from(self); 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> { + pub unsafe fn alloc_uninit_device(&self, len: usize) -> DeviceBuffer { let allocator = UsmAllocator::from(self); unsafe { Buffer::new(allocator, len) } }