spirv_std/
byte_addressable_buffer.rs

1//! Container for an untyped blob of data.
2
3use core::mem;
4
5#[spirv(buffer_load_intrinsic)]
6// HACK(eddyb) try to prevent MIR inlining from breaking our intrinsics.
7#[inline(never)]
8#[spirv_std_macros::gpu_only]
9unsafe fn buffer_load_intrinsic<T>(
10    buffer: &[u32],
11    // FIXME(eddyb) should be `usize`.
12    offset: u32,
13) -> T {
14    unsafe {
15        // NOTE(eddyb) this doesn't work with `rustc_codegen_spirv` and is only here
16        // for explanatory purposes, and to cause some kind of verbose error if
17        // `#[spirv(buffer_load_intrinsic)]` fails to replace calls to this function.
18        buffer
19            .as_ptr()
20            .cast::<u8>()
21            .add(offset as usize)
22            .cast::<T>()
23            .read()
24    }
25}
26
27#[spirv(buffer_store_intrinsic)]
28// HACK(eddyb) try to prevent MIR inlining from breaking our intrinsics.
29#[inline(never)]
30#[spirv_std_macros::gpu_only]
31unsafe fn buffer_store_intrinsic<T>(
32    buffer: &mut [u32],
33    // FIXME(eddyb) should be `usize`.
34    offset: u32,
35    value: T,
36) {
37    unsafe {
38        // NOTE(eddyb) this doesn't work with `rustc_codegen_spirv` and is only here
39        // for explanatory purposes, and to cause some kind of verbose error if
40        // `#[spirv(buffer_store_intrinsic)]` fails to replace calls to this function.
41        buffer
42            .as_mut_ptr()
43            .cast::<u8>()
44            .add(offset as usize)
45            .cast::<T>()
46            .write(value);
47    }
48}
49
50/// `ByteAddressableBuffer` is a view to an untyped blob of data, allowing
51/// loads and stores of arbitrary basic data types at arbitrary indices.
52///
53/// # Alignment
54/// All data must be aligned to size 4, each element within the data (e.g.
55/// struct fields) must have a size and alignment of a multiple of 4, and the
56/// `byte_index` passed to load and store must be a multiple of 4. Technically
57/// it is not a *byte* addressable buffer, but rather a *word* buffer, but this
58/// naming and behavior was inherited from HLSL (where it's UB to pass in an
59/// index not a multiple of 4).
60///
61/// # Safety
62/// Using these functions allows reading a different type from the buffer than
63/// was originally written (by a previous `store()` or the host API), allowing
64/// all sorts of safety guarantees to be bypassed, making it effectively a
65/// transmute.
66#[repr(transparent)]
67pub struct ByteAddressableBuffer<T> {
68    /// The underlying array of bytes, able to be directly accessed.
69    pub data: T,
70}
71
72fn bounds_check<T>(data: &[u32], byte_index: u32) {
73    let sizeof = mem::size_of::<T>() as u32;
74    if byte_index % 4 != 0 {
75        panic!("`byte_index` should be a multiple of 4");
76    }
77    let last_byte = byte_index + sizeof;
78    let len = data.len() as u32 * 4;
79    if byte_index + sizeof > len {
80        panic!(
81            "index out of bounds: the len is {} but loading {} bytes at `byte_index` {} reads until {} (exclusive)",
82            len, sizeof, byte_index, last_byte,
83        );
84    }
85}
86
87impl<'a> ByteAddressableBuffer<&'a [u32]> {
88    /// Creates a `ByteAddressableBuffer` from the untyped blob of data.
89    #[inline]
90    pub fn from_slice(data: &'a [u32]) -> Self {
91        Self { data }
92    }
93
94    /// Loads an arbitrary type from the buffer. `byte_index` must be a
95    /// multiple of 4.
96    ///
97    /// # Safety
98    /// See [`Self`].
99    pub unsafe fn load<T>(&self, byte_index: u32) -> T {
100        bounds_check::<T>(self.data, byte_index);
101        unsafe { buffer_load_intrinsic(self.data, byte_index) }
102    }
103
104    /// Loads an arbitrary type from the buffer. `byte_index` must be a
105    /// multiple of 4.
106    ///
107    /// # Safety
108    /// See [`Self`]. Additionally, bounds or alignment checking is not performed.
109    pub unsafe fn load_unchecked<T>(&self, byte_index: u32) -> T {
110        unsafe { buffer_load_intrinsic(self.data, byte_index) }
111    }
112}
113
114impl<'a> ByteAddressableBuffer<&'a mut [u32]> {
115    /// Creates a `ByteAddressableBuffer` from the untyped blob of data.
116    #[inline]
117    pub fn from_mut_slice(data: &'a mut [u32]) -> Self {
118        Self { data }
119    }
120
121    /// Create a non-mutable `ByteAddressableBuffer` from this mutable one.
122    #[inline]
123    pub fn as_ref(&self) -> ByteAddressableBuffer<&[u32]> {
124        ByteAddressableBuffer { data: self.data }
125    }
126
127    /// Loads an arbitrary type from the buffer. `byte_index` must be a
128    /// multiple of 4.
129    ///
130    /// # Safety
131    /// See [`Self`].
132    #[inline]
133    pub unsafe fn load<T>(&self, byte_index: u32) -> T {
134        unsafe { self.as_ref().load(byte_index) }
135    }
136
137    /// Loads an arbitrary type from the buffer. `byte_index` must be a
138    /// multiple of 4.
139    ///
140    /// # Safety
141    /// See [`Self`]. Additionally, bounds or alignment checking is not performed.
142    #[inline]
143    pub unsafe fn load_unchecked<T>(&self, byte_index: u32) -> T {
144        unsafe { self.as_ref().load_unchecked(byte_index) }
145    }
146
147    /// Stores an arbitrary type into the buffer. `byte_index` must be a
148    /// multiple of 4.
149    ///
150    /// # Safety
151    /// See [`Self`].
152    pub unsafe fn store<T>(&mut self, byte_index: u32, value: T) {
153        bounds_check::<T>(self.data, byte_index);
154        unsafe {
155            buffer_store_intrinsic(self.data, byte_index, value);
156        }
157    }
158
159    /// Stores an arbitrary type into the buffer. `byte_index` must be a
160    /// multiple of 4.
161    ///
162    /// # Safety
163    /// See [`Self`]. Additionally, bounds or alignment checking is not performed.
164    pub unsafe fn store_unchecked<T>(&mut self, byte_index: u32, value: T) {
165        unsafe {
166            buffer_store_intrinsic(self.data, byte_index, value);
167        }
168    }
169}