oneapi_rs/
buffer.rs

1//
2// Copyright (C) 2026 Intel Corporation
3//
4// Under the MIT License or the Apache License v2.0.
5// See LICENSE-MIT and LICENSE-APACHE for license information.
6// SPDX-License-Identifier: MIT OR Apache-2.0
7//
8
9use std::{
10    alloc::{Layout, handle_alloc_error},
11    ops::{Deref, DerefMut},
12    pin::Pin,
13    ptr::NonNull,
14    slice,
15    task::{Context, Poll},
16};
17
18use bytemuck::Pod;
19use pin_project::pin_project;
20
21use crate::{
22    event::{Event, EventFuture},
23    kernel::KernelArgument,
24    usm::{
25        DeviceAllocator, HostAccessible, HostAllocator, SharedAllocator, UsmAlloc, UsmAllocator,
26    },
27};
28
29/// The Buffer struct defines a shared array of one, two or three dimensions that can be used
30/// by the SYCL kernel. Buffers are templated on the type of their data, and the number of
31/// dimensions that the data is stored and accessed through.
32///
33/// A Buffer does not map to only one underlying backend object, and all SYCL backend memory objects
34/// may be temporary for use on a specific device.
35///
36/// Buffers can be constructed by methods provided by the [`Queue`](`crate::queue::Queue`) class.
37///
38/// The Buffer struct template takes a template parameter [`UsmAlloc`](`crate::usm::UsmAlloc`) for
39/// specifying an allocator which is used by the SYCL runtime when allocating temporary memory on
40/// the host.
41pub struct Buffer<T, A: UsmAlloc> {
42    data: NonNull<T>,
43    len: usize,
44    layout: Layout,
45    allocator: A,
46}
47
48impl<T, A: UsmAlloc> Buffer<T, A> {
49    /// Creates a new buffer given an allocator.
50    /// Safety: returns uninitialized memory.
51    pub(crate) unsafe fn new(allocator: A, len: usize) -> Self {
52        let layout = Layout::array::<T>(len).unwrap();
53        let ptr = match allocator.allocate(layout.clone()) {
54            Ok(ptr) => ptr,
55            _ => handle_alloc_error(layout),
56        };
57
58        Self {
59            data: ptr.cast(),
60            len,
61            layout,
62            allocator,
63        }
64    }
65
66    pub(crate) fn get_byte_ptr(&self) -> *mut u8 {
67        self.data.as_ptr().cast()
68    }
69
70    pub(crate) fn get_byte_size(&self) -> usize {
71        self.layout.size()
72    }
73
74    pub(crate) fn get_len(&self) -> usize {
75        self.len
76    }
77
78    unsafe fn as_raw_arg_impl(&self) -> &[u8] {
79        let data_ptr: *const NonNull<_> = &self.data;
80        let cast_ptr = data_ptr as *const u8;
81        unsafe { slice::from_raw_parts(cast_ptr, std::mem::size_of_val(&cast_ptr)) }
82    }
83}
84
85impl<T, A: UsmAlloc + HostAccessible> Deref for Buffer<T, A> {
86    type Target = [T];
87    fn deref(&self) -> &Self::Target {
88        unsafe { slice::from_raw_parts(self.data.as_ptr(), self.len) }
89    }
90}
91
92impl<T, A: UsmAlloc + HostAccessible> DerefMut for Buffer<T, A> {
93    fn deref_mut(&mut self) -> &mut Self::Target {
94        unsafe { slice::from_raw_parts_mut(self.data.as_ptr(), self.len) }
95    }
96}
97
98impl<T, A: UsmAlloc> Drop for Buffer<T, A> {
99    fn drop(&mut self) {
100        unsafe {
101            self.allocator.deallocate(self.data.cast(), self.layout);
102        }
103    }
104}
105
106pub type HostBuffer<T> = Buffer<T, UsmAllocator<HostAllocator>>;
107pub type SharedBuffer<T> = Buffer<T, UsmAllocator<SharedAllocator>>;
108pub type DeviceBuffer<T> = Buffer<T, UsmAllocator<DeviceAllocator>>;
109
110/// A [`Buffer`] whose initialization has been enqueued. You need to wait/await it.
111pub struct EnqueuedBuffer<T, A: UsmAlloc> {
112    buffer: Buffer<T, A>,
113    event: Event,
114}
115
116impl<T, A: UsmAlloc> EnqueuedBuffer<T, A> {
117    pub(crate) fn new(buffer: Buffer<T, A>, event: Event) -> Self {
118        Self { buffer, event }
119    }
120}
121
122impl<T, A: UsmAlloc> EnqueuedBuffer<T, A> {
123    /// Waits for [`Buffer`] initialization to finish.
124    pub fn wait(mut self) -> Buffer<T, A> {
125        self.event.wait();
126        self.buffer
127    }
128}
129
130pub type EnqueuedHostBuffer<T> = EnqueuedBuffer<T, UsmAllocator<HostAllocator>>;
131pub type EnqueuedSharedBuffer<T> = EnqueuedBuffer<T, UsmAllocator<SharedAllocator>>;
132pub type EnqueuedDeviceBuffer<T> = EnqueuedBuffer<T, UsmAllocator<DeviceAllocator>>;
133
134#[pin_project]
135/// A [`Future`] which represents a pending [`Buffer`] allocation.
136pub struct BufferFuture<T, A: UsmAlloc> {
137    buffer: Option<Buffer<T, A>>,
138    #[pin]
139    event_future: EventFuture,
140}
141
142impl<T, A: UsmAlloc> Future for BufferFuture<T, A> {
143    type Output = Buffer<T, A>;
144    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
145        let this = self.project();
146        this.event_future
147            .poll(cx)
148            .map(|_| this.buffer.take().unwrap())
149    }
150}
151
152impl<T, A: UsmAlloc> IntoFuture for EnqueuedBuffer<T, A> {
153    type Output = Buffer<T, A>;
154    type IntoFuture = BufferFuture<T, A>;
155
156    fn into_future(self) -> Self::IntoFuture {
157        Self::IntoFuture {
158            buffer: Some(self.buffer),
159            event_future: self.event.into_future(),
160        }
161    }
162}
163
164pub type HostBufferFuture<T> = BufferFuture<T, UsmAllocator<HostAllocator>>;
165pub type SharedBufferFuture<T> = BufferFuture<T, UsmAllocator<SharedAllocator>>;
166pub type DeviceBufferFuture<T> = BufferFuture<T, UsmAllocator<DeviceAllocator>>;
167
168unsafe impl<T: Pod, A: UsmAlloc> KernelArgument for Buffer<T, A> {
169    unsafe fn as_raw_arg(&self) -> &[u8] {
170        unsafe { self.as_raw_arg_impl() }
171    }
172}
173
174unsafe impl<T: Pod, A: UsmAlloc> KernelArgument for &Buffer<T, A> {
175    unsafe fn as_raw_arg(&self) -> &[u8] {
176        unsafe { self.as_raw_arg_impl() }
177    }
178}
179
180unsafe impl<T: Pod, A: UsmAlloc> KernelArgument for &mut Buffer<T, A> {
181    unsafe fn as_raw_arg(&self) -> &[u8] {
182        unsafe { self.as_raw_arg_impl() }
183    }
184}
oneapi_rs/buffer.rs

oneapi_rs/
buffer.rs
