spirv_std/
arch.rs

1//! SPIR-V Intrinsics
2//!
3//! This module is intended as a low level abstraction over SPIR-V instructions.
4//! These functions will typically map to a single instruction, and will perform
5//! no additional safety checks beyond type-checking.
6#[cfg(target_arch = "spirv")]
7use crate::Integer;
8use crate::{Scalar, SignedInteger, UnsignedInteger, Vector};
9#[cfg(target_arch = "spirv")]
10use core::arch::asm;
11use glam::UVec2;
12
13mod atomics;
14mod barrier;
15mod demote_to_helper_invocation_ext;
16mod derivative;
17mod mesh_shading;
18mod primitive;
19mod ray_tracing;
20mod subgroup;
21
22pub use atomics::*;
23pub use barrier::*;
24pub use demote_to_helper_invocation_ext::*;
25pub use derivative::*;
26pub use mesh_shading::*;
27pub use primitive::*;
28pub use ray_tracing::*;
29pub use subgroup::*;
30
31/// Result is true if any component of `vector` is true, otherwise result is
32/// false.
33#[spirv_std_macros::gpu_only]
34#[doc(alias = "OpAny")]
35#[inline]
36pub fn any<V: Vector<bool, N>, const N: usize>(vector: V) -> bool {
37    let mut result = false;
38
39    unsafe {
40        asm! {
41            "%bool = OpTypeBool",
42            "%vector = OpLoad _ {vector}",
43            "%result = OpAny %bool %vector",
44            "OpStore {result} %result",
45            vector = in(reg) &vector,
46            result = in(reg) &mut result
47        }
48    }
49
50    result
51}
52
53/// Result is true if all components of `vector` is true, otherwise result is
54/// false.
55#[spirv_std_macros::gpu_only]
56#[doc(alias = "OpAll")]
57#[inline]
58pub fn all<V: Vector<bool, N>, const N: usize>(vector: V) -> bool {
59    let mut result = false;
60
61    unsafe {
62        asm! {
63            "%bool = OpTypeBool",
64            "%vector = OpLoad _ {vector}",
65            "%result = OpAll %bool %vector",
66            "OpStore {result} %result",
67            vector = in(reg) &vector,
68            result = in(reg) &mut result
69        }
70    }
71
72    result
73}
74
75/// Extract a single, dynamically selected, component of a vector.
76///
77/// # Safety
78/// Behavior is undefined if `index`’s value is greater than or equal to the
79/// number of components in `vector`.
80#[spirv_std_macros::gpu_only]
81#[doc(alias = "OpVectorExtractDynamic")]
82#[inline]
83pub unsafe fn vector_extract_dynamic<T: Scalar, const N: usize>(
84    vector: impl Vector<T, N>,
85    index: usize,
86) -> T {
87    unsafe {
88        let mut result = T::default();
89
90        asm! {
91            "%vector = OpLoad _ {vector}",
92            "%element = OpVectorExtractDynamic _ %vector {index}",
93            "OpStore {element} %element",
94            vector = in(reg) &vector,
95            index = in(reg) index,
96            element = in(reg) &mut result
97        }
98
99        result
100    }
101}
102
103/// Make a copy of a vector, with a single, variably selected,
104/// component modified.
105///
106/// # Safety
107/// Behavior is undefined if `index`’s value is greater than or equal to the
108/// number of components in `vector`.
109#[spirv_std_macros::gpu_only]
110#[doc(alias = "OpVectorInsertDynamic")]
111#[inline]
112pub unsafe fn vector_insert_dynamic<T: Scalar, V: Vector<T, N>, const N: usize>(
113    vector: V,
114    index: usize,
115    element: T,
116) -> V {
117    unsafe {
118        let mut result = V::default();
119
120        asm! {
121            "%vector = OpLoad _ {vector}",
122            "%element = OpLoad _ {element}",
123            "%new_vector = OpVectorInsertDynamic _ %vector %element {index}",
124            "OpStore {result} %new_vector",
125            vector = in(reg) &vector,
126            index = in(reg) index,
127            element = in(reg) &element,
128            result = in(reg) &mut result,
129        }
130
131        result
132    }
133}
134
135/// Fragment-shader discard. Equivalvent to `discard()` from GLSL
136///
137/// Ceases all further processing in any invocation that executes it: Only
138/// instructions these invocations executed before [kill] have observable side
139/// effects.
140#[spirv_std_macros::gpu_only]
141#[doc(alias = "OpKill", alias = "discard")]
142#[allow(clippy::empty_loop)]
143pub fn kill() -> ! {
144    unsafe { asm!("OpKill", options(noreturn)) }
145}
146
147/// Read from the shader clock with either the `Subgroup` or `Device` scope.
148///
149/// See:
150/// <https://htmlpreview.github.io/?https://github.com/KhronosGroup/SPIRV-Registry/blob/master/extensions/KHR/SPV_KHR_shader_clock.html>
151#[spirv_std_macros::gpu_only]
152#[doc(alias = "OpReadClockKHR")]
153pub fn read_clock_khr<const SCOPE: u32>() -> u64 {
154    unsafe {
155        let mut result: u64;
156
157        asm! {
158            "%uint = OpTypeInt 32 0",
159            "%scope = OpConstant %uint {scope}",
160            "{result} = OpReadClockKHR typeof*{result} %scope",
161            result = out(reg) result,
162            scope = const SCOPE,
163        };
164
165        result
166    }
167}
168
169/// Like `read_clock_khr` but returns a vector to avoid requiring the `Int64`
170/// capability. It returns a 'vector of two-components of 32-bit unsigned
171/// integer type with the first component containing the 32 least significant
172/// bits and the second component containing the 32 most significant bits.'
173#[spirv_std_macros::gpu_only]
174#[doc(alias = "OpReadClockKHR")]
175pub fn read_clock_uvec2_khr<const SCOPE: u32>() -> UVec2 {
176    unsafe {
177        let mut result = UVec2::default();
178
179        asm! {
180            "%uint = OpTypeInt 32 0",
181            "%scope = OpConstant %uint {scope}",
182            "%result = OpReadClockKHR typeof*{result} %scope",
183            "OpStore {result} %result",
184            result = in(reg) &mut result,
185            scope = const SCOPE,
186        };
187
188        result
189    }
190}
191
192#[cfg(target_arch = "spirv")]
193unsafe fn call_glsl_op_with_ints<T: Integer, const OP: u32>(a: T, b: T) -> T {
194    unsafe {
195        let mut result = T::default();
196        asm!(
197            "%glsl = OpExtInstImport \"GLSL.std.450\"",
198            "%a = OpLoad _ {a}",
199            "%b = OpLoad _ {b}",
200            "%result = OpExtInst typeof*{result} %glsl {op} %a %b",
201            "OpStore {result} %result",
202            a = in(reg) &a,
203            b = in(reg) &b,
204            result = in(reg) &mut result,
205            op = const OP
206        );
207        result
208    }
209}
210
211/// Compute the minimum of two unsigned integers via a GLSL extended instruction.
212#[spirv_std_macros::gpu_only]
213pub fn unsigned_min<T: UnsignedInteger>(a: T, b: T) -> T {
214    unsafe { call_glsl_op_with_ints::<_, 38>(a, b) }
215}
216
217/// Compute the maximum of two unsigned integers via a GLSL extended instruction.
218#[spirv_std_macros::gpu_only]
219pub fn unsigned_max<T: UnsignedInteger>(a: T, b: T) -> T {
220    unsafe { call_glsl_op_with_ints::<_, 41>(a, b) }
221}
222
223/// Compute the minimum of two signed integers via a GLSL extended instruction.
224#[spirv_std_macros::gpu_only]
225pub fn signed_min<T: SignedInteger>(a: T, b: T) -> T {
226    unsafe { call_glsl_op_with_ints::<_, 39>(a, b) }
227}
228
229/// Compute the maximum of two signed integers via a GLSL extended instruction.
230#[spirv_std_macros::gpu_only]
231pub fn signed_max<T: SignedInteger>(a: T, b: T) -> T {
232    unsafe { call_glsl_op_with_ints::<_, 42>(a, b) }
233}
234
235/// Index into an array without bounds checking.
236///
237/// The main purpose of this trait is to work around the fact that the regular `get_unchecked*`
238/// methods do not work in in SPIR-V.
239pub trait IndexUnchecked<T> {
240    /// Returns a reference to the element at `index`. The equivalent of `get_unchecked`.
241    ///
242    /// # Safety
243    /// Behavior is undefined if the `index` value is greater than or equal to the length of the array.
244    unsafe fn index_unchecked(&self, index: usize) -> &T;
245    /// Returns a mutable reference to the element at `index`. The equivalent of `get_unchecked_mut`.
246    ///
247    /// # Safety
248    /// Behavior is undefined if the `index` value is greater than or equal to the length of the array.
249    unsafe fn index_unchecked_mut(&mut self, index: usize) -> &mut T;
250}
251
252impl<T> IndexUnchecked<T> for [T] {
253    #[cfg(target_arch = "spirv")]
254    unsafe fn index_unchecked(&self, index: usize) -> &T {
255        unsafe {
256            // FIXME(eddyb) `let mut result = T::default()` uses (for `asm!`), with this.
257            let mut result_slot = core::mem::MaybeUninit::uninit();
258            asm! {
259                "%slice_ptr = OpLoad _ {slice_ptr_ptr}",
260                "%data_ptr = OpCompositeExtract _ %slice_ptr 0",
261                "%result = OpAccessChain _ %data_ptr {index}",
262                "OpStore {result_slot} %result",
263                slice_ptr_ptr = in(reg) &self,
264                index = in(reg) index,
265                result_slot = in(reg) result_slot.as_mut_ptr(),
266            }
267            result_slot.assume_init()
268        }
269    }
270
271    #[cfg(not(target_arch = "spirv"))]
272    unsafe fn index_unchecked(&self, index: usize) -> &T {
273        unsafe { self.get_unchecked(index) }
274    }
275
276    #[cfg(target_arch = "spirv")]
277    unsafe fn index_unchecked_mut(&mut self, index: usize) -> &mut T {
278        unsafe {
279            // FIXME(eddyb) `let mut result = T::default()` uses (for `asm!`), with this.
280            let mut result_slot = core::mem::MaybeUninit::uninit();
281            asm! {
282                "%slice_ptr = OpLoad _ {slice_ptr_ptr}",
283                "%data_ptr = OpCompositeExtract _ %slice_ptr 0",
284                "%result = OpAccessChain _ %data_ptr {index}",
285                "OpStore {result_slot} %result",
286                slice_ptr_ptr = in(reg) &self,
287                index = in(reg) index,
288                result_slot = in(reg) result_slot.as_mut_ptr(),
289            }
290            result_slot.assume_init()
291        }
292    }
293
294    #[cfg(not(target_arch = "spirv"))]
295    unsafe fn index_unchecked_mut(&mut self, index: usize) -> &mut T {
296        unsafe { self.get_unchecked_mut(index) }
297    }
298}
299
300impl<T, const N: usize> IndexUnchecked<T> for [T; N] {
301    #[cfg(target_arch = "spirv")]
302    unsafe fn index_unchecked(&self, index: usize) -> &T {
303        unsafe {
304            // FIXME(eddyb) `let mut result = T::default()` uses (for `asm!`), with this.
305            let mut result_slot = core::mem::MaybeUninit::uninit();
306            asm! {
307                "%result = OpAccessChain _ {array_ptr} {index}",
308                "OpStore {result_slot} %result",
309                array_ptr = in(reg) self,
310                index = in(reg) index,
311                result_slot = in(reg) result_slot.as_mut_ptr(),
312            }
313            result_slot.assume_init()
314        }
315    }
316
317    #[cfg(not(target_arch = "spirv"))]
318    unsafe fn index_unchecked(&self, index: usize) -> &T {
319        unsafe { self.get_unchecked(index) }
320    }
321
322    #[cfg(target_arch = "spirv")]
323    unsafe fn index_unchecked_mut(&mut self, index: usize) -> &mut T {
324        unsafe {
325            // FIXME(eddyb) `let mut result = T::default()` uses (for `asm!`), with this.
326            let mut result_slot = core::mem::MaybeUninit::uninit();
327            asm! {
328                "%result = OpAccessChain _ {array_ptr} {index}",
329                "OpStore {result_slot} %result",
330                array_ptr = in(reg) self,
331                index = in(reg) index,
332                result_slot = in(reg) result_slot.as_mut_ptr(),
333            }
334            result_slot.assume_init()
335        }
336    }
337
338    #[cfg(not(target_arch = "spirv"))]
339    unsafe fn index_unchecked_mut(&mut self, index: usize) -> &mut T {
340        unsafe { self.get_unchecked_mut(index) }
341    }
342}