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

oneapi_rs/
usm.rs
