sycl_rs/
usm.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 crate::queue::Queue;
10
11use allocator_api2::alloc::{AllocError, Allocator};
12use sycl_rs_sys::usm::ffi;
13
14use std::{alloc::Layout, marker::PhantomData, ptr::NonNull};
15
16type CxxResult<T> = cxx::core::result::Result<T, cxx::Exception>;
17
18/// An instance of a USM allocator.
19pub struct UsmAllocator<T: UsmAllocatorKind> {
20    queue: Queue,
21    _kind: PhantomData<T>,
22}
23
24/// A marker trait for USM allocators.
25///
26/// Safety: a type implementing this trait must be a valid USM allocator managed by a SYCL runtime.
27pub unsafe trait UsmAlloc: Allocator {}
28
29unsafe impl<T: UsmAllocatorKind> UsmAlloc for UsmAllocator<T> {}
30
31pub trait UsmAllocatorKind {
32    /// Allocates uninitialized memory.
33    /// Safety: the caller must not read uninitialized memory. The caller must also free this
34    /// memory manually.
35    unsafe fn alloc(alignment: usize, num_bytes: usize, queue: &Queue) -> CxxResult<*mut u8>;
36}
37
38/// A marker trait for host-accessible USM allocators.
39///
40/// Safety: a type implementing this trait must be a valid USM allocator managed by a SYCL runtime,
41/// that allocates memory accessible from the host.
42pub unsafe trait HostAccessible {}
43
44impl<T: UsmAllocatorKind> From<&Queue> for UsmAllocator<T> {
45    fn from(queue: &Queue) -> Self {
46        Self {
47            queue: queue.clone(),
48            _kind: PhantomData,
49        }
50    }
51}
52
53unsafe impl<T: UsmAllocatorKind> Allocator for UsmAllocator<T> {
54    fn allocate(&self, layout: Layout) -> Result<NonNull<[u8]>, AllocError> {
55        let ptr = unsafe { T::alloc(layout.align(), layout.size(), &self.queue) }
56            .map_err(|_e| AllocError)?;
57
58        let ptr = NonNull::new(ptr).ok_or(AllocError)?;
59        let slice = NonNull::slice_from_raw_parts(ptr, layout.size());
60
61        Ok(slice)
62    }
63
64    unsafe fn deallocate(&self, ptr: NonNull<u8>, _layout: Layout) {
65        unsafe {
66            ffi::free(ptr.as_ptr(), &self.queue.0);
67        }
68    }
69}
70
71/// An allocator for Device-side arrays
72///
73/// Safety: memory allocated by this allocator cannot be accessed on the host side
74#[allow(dead_code)]
75pub struct DeviceAllocator;
76
77impl UsmAllocatorKind for DeviceAllocator {
78    unsafe fn alloc(alignment: usize, num_bytes: usize, queue: &Queue) -> CxxResult<*mut u8> {
79        unsafe { ffi::aligned_alloc_device(alignment, num_bytes, &queue.0) }
80    }
81}
82
83/// An allocator for Host-side arrays
84pub struct HostAllocator;
85
86impl UsmAllocatorKind for HostAllocator {
87    unsafe fn alloc(alignment: usize, num_bytes: usize, queue: &Queue) -> CxxResult<*mut u8> {
88        unsafe { ffi::aligned_alloc_host(alignment, num_bytes, &queue.0) }
89    }
90}
91
92unsafe impl HostAccessible for UsmAllocator<HostAllocator> {}
93
94/// An allocator for shared memory arrays
95pub struct SharedAllocator;
96
97impl UsmAllocatorKind for SharedAllocator {
98    unsafe fn alloc(alignment: usize, num_bytes: usize, queue: &Queue) -> CxxResult<*mut u8> {
99        unsafe { ffi::aligned_alloc_shared(alignment, num_bytes, &queue.0) }
100    }
101}
102
103unsafe impl HostAccessible for UsmAllocator<SharedAllocator> {}
sycl_rs/usm.rs

sycl_rs/
usm.rs
