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

oneapi_rs/
usm.rs
