sycl_rs/
usmbox.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    Result,
23    event::{Event, EventFuture},
24    kernel::KernelArgument,
25    usm::{
26        DeviceAllocator, HostAccessible, HostAllocator, SharedAllocator, UsmAlloc, UsmAllocator,
27    },
28};
29
30/// The `UsmBox` struct defines a shared array of one, two or three dimensions that can be used
31/// by the SYCL kernel. `UsmBox`es are generic w.r.t. the type of their data, and the number of
32/// dimensions that the data is stored and accessed through.
33///
34/// A `UsmBox` does not map to only one underlying backend object, and all SYCL backend memory objects
35/// may be temporary for use on a specific device.
36///
37/// `UsmBox`es can be constructed by methods provided by the [`Queue`](`crate::queue::Queue`) class.
38///
39/// The `UsmBox` struct takes a generic parameter [`UsmAlloc`](`crate::usm::UsmAlloc`) for
40/// specifying an allocator which is used by the SYCL runtime when allocating temporary memory on
41/// the host.
42pub struct UsmBox<T, A: UsmAlloc> {
43    data: NonNull<T>,
44    len: usize,
45    layout: Layout,
46    allocator: A,
47}
48
49impl<T, A: UsmAlloc> UsmBox<T, A> {
50    /// Creates a new array given an allocator.
51    /// Safety: returns uninitialized memory.
52    pub(crate) unsafe fn new(allocator: A, len: usize) -> Self {
53        let layout = Layout::array::<T>(len).unwrap();
54        let ptr = match allocator.allocate(layout.clone()) {
55            Ok(ptr) => ptr,
56            _ => handle_alloc_error(layout),
57        };
58
59        Self {
60            data: ptr.cast(),
61            len,
62            layout,
63            allocator,
64        }
65    }
66
67    pub(crate) fn get_byte_ptr(&self) -> *mut u8 {
68        self.data.as_ptr().cast()
69    }
70
71    pub(crate) fn get_byte_size(&self) -> usize {
72        self.layout.size()
73    }
74
75    pub(crate) fn get_len(&self) -> usize {
76        self.len
77    }
78
79    unsafe fn as_raw_arg_impl(&self) -> &[u8] {
80        let data_ptr: *const NonNull<_> = &self.data;
81        let cast_ptr = data_ptr as *const u8;
82        unsafe { slice::from_raw_parts(cast_ptr, std::mem::size_of_val(&cast_ptr)) }
83    }
84}
85
86impl<T, A: UsmAlloc + HostAccessible> Deref for UsmBox<T, A> {
87    type Target = [T];
88    fn deref(&self) -> &Self::Target {
89        unsafe { slice::from_raw_parts(self.data.as_ptr(), self.len) }
90    }
91}
92
93impl<T, A: UsmAlloc + HostAccessible> DerefMut for UsmBox<T, A> {
94    fn deref_mut(&mut self) -> &mut Self::Target {
95        unsafe { slice::from_raw_parts_mut(self.data.as_ptr(), self.len) }
96    }
97}
98
99impl<T, A: UsmAlloc> Drop for UsmBox<T, A> {
100    fn drop(&mut self) {
101        unsafe {
102            self.allocator.deallocate(self.data.cast(), self.layout);
103        }
104    }
105}
106
107pub type HostUsmBox<T> = UsmBox<T, UsmAllocator<HostAllocator>>;
108pub type SharedUsmBox<T> = UsmBox<T, UsmAllocator<SharedAllocator>>;
109pub type DeviceUsmBox<T> = UsmBox<T, UsmAllocator<DeviceAllocator>>;
110
111/// A [`UsmBox`] whose initialization has been enqueued. You need to wait/await it.
112pub struct EnqueuedUsmBox<T, A: UsmAlloc> {
113    array: UsmBox<T, A>,
114    event: Event,
115}
116
117impl<T, A: UsmAlloc> EnqueuedUsmBox<T, A> {
118    pub(crate) fn new(array: UsmBox<T, A>, event: Event) -> Self {
119        Self { array, event }
120    }
121}
122
123impl<T, A: UsmAlloc> EnqueuedUsmBox<T, A> {
124    /// Performs a blocking wait for the [`UsmBox`] initialization to complete. Returns an error if
125    /// a synchronous SYCL exception occurs.
126    ///
127    /// Dropping an enqueued array does not wait for its completion.
128    pub fn wait(mut self) -> Result<UsmBox<T, A>> {
129        self.event.wait().map(|_| self.array)
130    }
131}
132
133pub type EnqueuedHostUsmBox<T> = EnqueuedUsmBox<T, UsmAllocator<HostAllocator>>;
134pub type EnqueuedSharedUsmBox<T> = EnqueuedUsmBox<T, UsmAllocator<SharedAllocator>>;
135pub type EnqueuedDeviceUsmBox<T> = EnqueuedUsmBox<T, UsmAllocator<DeviceAllocator>>;
136
137#[pin_project]
138/// A [`Future`] which represents a pending [`UsmBox`] allocation.
139pub struct UsmBoxFuture<T, A: UsmAlloc> {
140    array: Option<UsmBox<T, A>>,
141    #[pin]
142    event_future: EventFuture,
143}
144
145impl<T, A: UsmAlloc> Future for UsmBoxFuture<T, A> {
146    type Output = Result<UsmBox<T, A>>;
147    fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
148        let this = self.project();
149        this.event_future
150            .poll(cx)
151            .map(|result| result.map(|_| this.array.take().unwrap()))
152    }
153}
154
155impl<T, A: UsmAlloc> IntoFuture for EnqueuedUsmBox<T, A> {
156    type Output = Result<UsmBox<T, A>>;
157    type IntoFuture = UsmBoxFuture<T, A>;
158
159    fn into_future(self) -> Self::IntoFuture {
160        Self::IntoFuture {
161            array: Some(self.array),
162            event_future: self.event.into_future(),
163        }
164    }
165}
166
167pub type HostUsmBoxFuture<T> = UsmBoxFuture<T, UsmAllocator<HostAllocator>>;
168pub type SharedUsmBoxFuture<T> = UsmBoxFuture<T, UsmAllocator<SharedAllocator>>;
169pub type DeviceUsmBoxFuture<T> = UsmBoxFuture<T, UsmAllocator<DeviceAllocator>>;
170
171unsafe impl<T: Pod, A: UsmAlloc> KernelArgument for UsmBox<T, A> {
172    unsafe fn as_raw_arg(&self) -> &[u8] {
173        unsafe { self.as_raw_arg_impl() }
174    }
175}
176
177unsafe impl<T: Pod, A: UsmAlloc> KernelArgument for &UsmBox<T, A> {
178    unsafe fn as_raw_arg(&self) -> &[u8] {
179        unsafe { self.as_raw_arg_impl() }
180    }
181}
182
183unsafe impl<T: Pod, A: UsmAlloc> KernelArgument for &mut UsmBox<T, A> {
184    unsafe fn as_raw_arg(&self) -> &[u8] {
185        unsafe { self.as_raw_arg_impl() }
186    }
187}
sycl_rs/usmbox.rs

sycl_rs/
usmbox.rs
