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