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

sycl_rs/
usm.rs
