spirv_std/
arch.rs

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