Skip to content

Commit c797246

Browse files
committed
builtin: add compute builtins, except workgroup_size
1 parent cc1b4d5 commit c797246

File tree

7 files changed

+284
-0
lines changed

7 files changed

+284
-0
lines changed

crates/spirv-std/src/builtin.rs

Lines changed: 60 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,60 @@
1+
//! Functionality to declare builtins, mostly proc macros
2+
//!
3+
//! # Making built-in functions for `spirv-std`
4+
//!
5+
//! Usually, built-ins are implemented as freestanding functions in `spirv-std`. We like to keep function declaration
6+
//! outside the macro to make it easier for users to browse the source code.
7+
//!
8+
//! Example on how to declare an Input Built-in:
9+
//! ```no_run
10+
//! # use spirv_std_macros::gpu_only;
11+
//! #
12+
//! /// GLSL docs short description in #Name section. Remove the first "Contains " since we're using getters instead
13+
//! /// of globals, capitalize and add a dot to the end.
14+
//! ///
15+
//! /// GLSL docs full #Description section.
16+
//! ///
17+
//! /// We're using GLSL documentation of this built-in, which is usually more descriptive than the SPIR-V or WGSL docs.
18+
//! /// Change all references to link with rust-gpu intrinsics.
19+
//! ///
20+
//! /// Update the links of GLSL and WGSL to reference the correct page, keep SPIR-V as is. GLSL may link to the
21+
//! /// [reference](https://registry.khronos.org/OpenGL-Refpages/gl4/) or to the
22+
//! /// [glsl extensions github repo](https://github.com/KhronosGroup/GLSL/tree/main/extensions).
23+
//! /// * GLSL: [`gl_MyBuiltIn`](https://registry.khronos.org/OpenGL-Refpages/gl4/html/gl_LocalInvocationID.xhtml)
24+
//! /// * WGSL: [`my_built_in`](https://www.w3.org/TR/WGSL/#local-invocation-id-builtin-value)
25+
//! /// * SPIRV: [`MyBuiltIn`](https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_builtin)
26+
//! #[doc(alias = "gl_MyBuiltIn")]
27+
//! #[doc(alias = "MyBuiltIn")]
28+
//! #[inline]
29+
//! #[gpu_only]
30+
//! pub fn my_built_in() -> u32 {
31+
//! crate::load_builtin!(MyBuiltIn)
32+
//! }
33+
//! ```
34+
//!
35+
//! Reference links:
36+
//! * [WGSL specification describing builtins](https://www.w3.org/TR/WGSL/#builtin-inputs-outputs)
37+
//! * [SPIR-V specification for builtins](https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_builtin)
38+
//! * [GLSL reference](https://registry.khronos.org/OpenGL-Refpages/gl4/)
39+
//! * [GLSL reference source code](https://github.com/KhronosGroup/OpenGL-Refpages/tree/main/gl4)
40+
//! * [GLSL extensions](https://github.com/KhronosGroup/GLSL/tree/main/extensions)
41+
42+
/// Query SPIR-V (read-only global) built-in values
43+
///
44+
/// See [module level documentation] on how to use these.
45+
#[macro_export]
46+
macro_rules! load_builtin {
47+
($name:ident $(: $ty:ty)?) => {
48+
unsafe {
49+
let mut result $(: $ty)? = Default::default();
50+
::core::arch::asm! {
51+
"%builtin = OpVariable typeof{result_ref} Input",
52+
concat!("OpDecorate %builtin BuiltIn ", stringify!($name)),
53+
"%result = OpLoad typeof*{result_ref} %builtin",
54+
"OpStore {result_ref} %result",
55+
result_ref = in(reg) &mut result,
56+
}
57+
result
58+
}
59+
};
60+
}

crates/spirv-std/src/compute.rs

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,93 @@
1+
//! compute shader built-ins
2+
3+
use glam::UVec3;
4+
5+
/// The index of work item currently being operated on by a compute shader.
6+
///
7+
/// In the compute language, [`local_invocation_id`] is an input variable containing the n-dimensional index of the
8+
/// local work invocation within the work group that the current shader is executing in. The possible values for this
9+
/// variable range across the local work group size, i.e., `(0,0,0)` to
10+
/// `(`[`workgroup_size`]`.x - 1, `[`workgroup_size`]`.y - 1, `[`workgroup_size`]`.z - 1)`.
11+
///
12+
/// * GLSL: [`gl_LocalInvocationID`](https://registry.khronos.org/OpenGL-Refpages/gl4/html/gl_LocalInvocationID.xhtml)
13+
/// * WGSL: [`local_invocation_id`](https://www.w3.org/TR/WGSL/#local-invocation-id-builtin-value)
14+
/// * SPIR-V: [`LocalInvocationId`](https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_builtin)
15+
#[doc(alias = "gl_LocalInvocationID")]
16+
#[doc(alias = "LocalInvocationId")]
17+
#[inline]
18+
#[gpu_only]
19+
pub fn local_invocation_id() -> UVec3 {
20+
crate::load_builtin!(LocalInvocationId)
21+
}
22+
23+
/// The local linear index of work item currently being operated on by a compute shader.
24+
///
25+
/// In the compute language, [`local_invocation_index`] is a derived input variable containing the 1-dimensional
26+
/// linearized index of the work invocation within the work group that the current shader is executing on. The value of
27+
/// [`local_invocation_index`] is equal to [`local_invocation_id`]`.z * `[`workgroup_size`]`.x * `[`workgroup_size`]`.y`
28+
/// `+ `[`local_invocation_id`]`.y * `[`workgroup_size`]`.x + `[`local_invocation_id`]`.x`.
29+
///
30+
/// * GLSL: [`gl_LocalInvocationIndex`](https://registry.khronos.org/OpenGL-Refpages/gl4/html/gl_LocalInvocationIndex.xhtml)
31+
/// * WGSL: [`local_invocation_index`](https://www.w3.org/TR/WGSL/#local-invocation-index-builtin-value)
32+
/// * SPIR-V: [`LocalInvocationIndex`](https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_builtin)
33+
#[doc(alias = "gl_LocalInvocationIndex")]
34+
#[doc(alias = "LocalInvocationIndex")]
35+
#[inline]
36+
#[gpu_only]
37+
pub fn local_invocation_index() -> u32 {
38+
crate::load_builtin!(LocalInvocationIndex)
39+
}
40+
41+
/// The global index of work item currently being operated on by a compute shader.
42+
///
43+
/// In the compute language, [`global_invocation_id`] is a derived input variable containing the n-dimensional index of
44+
/// the work invocation within the global work group that the current shader is executing on. The value of
45+
/// [`global_invocation_id`] is equal to [`workgroup_id`] * [`workgroup_size`] + [`local_invocation_id`].
46+
///
47+
/// * GLSL: [`gl_GlobalInvocationID`](https://registry.khronos.org/OpenGL-Refpages/gl4/html/gl_GlobalInvocationID.xhtml)
48+
/// * WGSL: [`global_invocation_id`](https://www.w3.org/TR/WGSL/#global-invocation-index-builtin-value)
49+
/// * SPIR-V: [`GlobalInvocationId`](https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_builtin)
50+
#[doc(alias = "gl_GlobalInvocationID")]
51+
#[doc(alias = "GlobalInvocationId")]
52+
#[inline]
53+
#[gpu_only]
54+
pub fn global_invocation_id() -> UVec3 {
55+
crate::load_builtin!(GlobalInvocationId)
56+
}
57+
58+
// custom: do not mention `glDispatchCompute` directly, be more general across APIs
59+
/// The number of workgroups that have been dispatched to a compute shader.
60+
///
61+
/// In the compute language, [`num_workgroups`] contains the total number of work groups that will execute the compute
62+
/// shader. The components of [`num_workgroups`] are equal to the `x`, `y`, and `z` parameters passed to the dispatch
63+
/// command.
64+
///
65+
/// * GLSL: [`gl_NumWorkGroups`](https://registry.khronos.org/OpenGL-Refpages/gl4/html/gl_NumWorkGroups.xhtml)
66+
/// * WGSL: [`num_workgroups`](https://www.w3.org/TR/WGSL/#num-workgroups-builtin-value)
67+
/// * SPIR-V: [`NumWorkgroups`](https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_builtin)
68+
#[doc(alias = "gl_NumWorkGroups")]
69+
#[doc(alias = "NumWorkgroups")]
70+
#[inline]
71+
#[gpu_only]
72+
pub fn num_workgroups() -> UVec3 {
73+
crate::load_builtin!(NumWorkgroups)
74+
}
75+
76+
// custom: do not mention `glDispatchCompute` directly, be more general across APIs
77+
/// The index of the workgroup currently being operated on by a compute shader.
78+
///
79+
/// In the compute language, [`workgroup_id`] contains the 3-dimensional index of the global work group that the current
80+
/// compute shader invocation is executing within. The possible values range across the parameters passed into the
81+
/// dispatch command, i.e., from `(0, 0, 0)` to
82+
/// `(`[`num_workgroups`]`.x - 1, `[`num_workgroups`]`.y - 1, `[`num_workgroups`]`.z - 1)`.
83+
///
84+
/// * GLSL: [`gl_WorkGroupID`](https://registry.khronos.org/OpenGL-Refpages/gl4/html/gl_WorkGroupID.xhtml)
85+
/// * WGSL: [`workgroup_id`](https://www.w3.org/TR/WGSL/#workgroup-id-builtin-value)
86+
/// * SPIR-V: [`WorkgroupId`](https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_builtin)
87+
#[doc(alias = "gl_WorkGroupID")]
88+
#[doc(alias = "WorkgroupId")]
89+
#[inline]
90+
#[gpu_only]
91+
pub fn workgroup_id() -> UVec3 {
92+
crate::load_builtin!(WorkgroupId)
93+
}

crates/spirv-std/src/lib.rs

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,9 @@ pub use macros::{debug_printf, debug_printfln};
9393
pub mod arch;
9494
pub mod atomic;
9595
pub mod barrier;
96+
pub mod builtin;
9697
pub mod byte_addressable_buffer;
98+
pub mod compute;
9799
pub mod debug_printf;
98100
pub mod float;
99101
pub mod fragment;
Lines changed: 31 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,31 @@
1+
// build-pass
2+
// compile-flags: -C llvm-args=--disassemble-globals
3+
// normalize-stderr-test "OpSource .*\n" -> ""
4+
// normalize-stderr-test "OpLine .*\n" -> ""
5+
// normalize-stderr-test "%\d+ = OpString .*\n" -> ""
6+
// normalize-stderr-test "; .*\n" -> ""
7+
// normalize-stderr-test "OpCapability VulkanMemoryModel\n" -> ""
8+
// normalize-stderr-test "OpMemoryModel Logical Vulkan" -> "OpMemoryModel Logical Simple"
9+
// ignore-vulkan1.0
10+
// ignore-spv1.0
11+
// ignore-spv1.1
12+
// ignore-spv1.2
13+
14+
use spirv_std::glam::*;
15+
use spirv_std::spirv;
16+
17+
#[spirv(compute(threads(1)))]
18+
pub fn compute(
19+
#[spirv(storage_buffer, descriptor_set = 0, binding = 0)] buffer: &mut [u32],
20+
#[spirv(local_invocation_id)] local_invocation_id: UVec3,
21+
#[spirv(local_invocation_index)] local_invocation_index: u32,
22+
#[spirv(global_invocation_id)] global_invocation_id: UVec3,
23+
#[spirv(num_workgroups)] num_workgroups: UVec3,
24+
#[spirv(workgroup_id)] workgroup_id: UVec3,
25+
) {
26+
buffer[0] = local_invocation_id.x
27+
+ local_invocation_index
28+
+ global_invocation_id.x
29+
+ num_workgroups.x
30+
+ workgroup_id.x;
31+
}
Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
OpCapability Shader
2+
OpMemoryModel Logical Simple
3+
OpEntryPoint GLCompute %1 "compute" %2 %3 %4 %5 %6 %7
4+
OpExecutionMode %1 LocalSize 1 1 1
5+
OpName %2 "buffer"
6+
OpName %3 "local_invocation_id"
7+
OpName %4 "local_invocation_index"
8+
OpName %5 "global_invocation_id"
9+
OpName %6 "num_workgroups"
10+
OpName %7 "workgroup_id"
11+
OpDecorate %9 ArrayStride 4
12+
OpDecorate %10 Block
13+
OpMemberDecorate %10 0 Offset 0
14+
OpDecorate %2 Binding 0
15+
OpDecorate %2 DescriptorSet 0
16+
OpDecorate %3 BuiltIn LocalInvocationId
17+
OpDecorate %4 BuiltIn LocalInvocationIndex
18+
OpDecorate %5 BuiltIn GlobalInvocationId
19+
OpDecorate %6 BuiltIn NumWorkgroups
20+
OpDecorate %7 BuiltIn WorkgroupId
21+
%11 = OpTypeInt 32 0
22+
%9 = OpTypeRuntimeArray %11
23+
%10 = OpTypeStruct %9
24+
%12 = OpTypePointer StorageBuffer %10
25+
%13 = OpTypeVector %11 3
26+
%14 = OpTypePointer Input %13
27+
%15 = OpTypePointer Input %11
28+
%16 = OpTypeVoid
29+
%17 = OpTypeFunction %16
30+
%18 = OpTypePointer StorageBuffer %9
31+
%2 = OpVariable %12 StorageBuffer
32+
%19 = OpConstant %11 0
33+
%3 = OpVariable %14 Input
34+
%4 = OpVariable %15 Input
35+
%5 = OpVariable %14 Input
36+
%6 = OpVariable %14 Input
37+
%7 = OpVariable %14 Input
38+
%20 = OpTypeBool
39+
%21 = OpTypePointer StorageBuffer %11
Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
// build-pass
2+
// compile-flags: -C llvm-args=--disassemble-globals
3+
// normalize-stderr-test "OpSource .*\n" -> ""
4+
// normalize-stderr-test "OpLine .*\n" -> ""
5+
// normalize-stderr-test "%\d+ = OpString .*\n" -> ""
6+
// normalize-stderr-test "; .*\n" -> ""
7+
// normalize-stderr-test "OpCapability VulkanMemoryModel\n" -> ""
8+
// normalize-stderr-test "OpMemoryModel Logical Vulkan" -> "OpMemoryModel Logical Simple"
9+
// ignore-vulkan1.0
10+
// ignore-spv1.0
11+
// ignore-spv1.1
12+
// ignore-spv1.2
13+
14+
use spirv_std::compute::*;
15+
use spirv_std::glam::*;
16+
use spirv_std::spirv;
17+
18+
#[spirv(compute(threads(1)))]
19+
pub fn compute(#[spirv(storage_buffer, descriptor_set = 0, binding = 0)] buffer: &mut [u32]) {
20+
buffer[0] = local_invocation_id().x
21+
+ local_invocation_index()
22+
+ global_invocation_id().x
23+
+ num_workgroups().x
24+
+ workgroup_id().x;
25+
}
Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
1+
OpCapability Shader
2+
OpMemoryModel Logical Simple
3+
OpEntryPoint GLCompute %1 "compute" %2 %3 %4 %5 %6 %7
4+
OpExecutionMode %1 LocalSize 1 1 1
5+
OpName %2 "buffer"
6+
OpDecorate %10 ArrayStride 4
7+
OpDecorate %11 Block
8+
OpMemberDecorate %11 0 Offset 0
9+
OpDecorate %2 Binding 0
10+
OpDecorate %2 DescriptorSet 0
11+
OpDecorate %3 BuiltIn LocalInvocationId
12+
OpDecorate %4 BuiltIn LocalInvocationIndex
13+
OpDecorate %5 BuiltIn GlobalInvocationId
14+
OpDecorate %6 BuiltIn NumWorkgroups
15+
OpDecorate %7 BuiltIn WorkgroupId
16+
%12 = OpTypeInt 32 0
17+
%10 = OpTypeRuntimeArray %12
18+
%11 = OpTypeStruct %10
19+
%13 = OpTypePointer StorageBuffer %11
20+
%14 = OpTypeVoid
21+
%15 = OpTypeFunction %14
22+
%16 = OpTypePointer StorageBuffer %10
23+
%2 = OpVariable %13 StorageBuffer
24+
%17 = OpConstant %12 0
25+
%18 = OpTypeVector %12 3
26+
%19 = OpTypePointer Input %18
27+
%3 = OpVariable %19 Input
28+
%20 = OpTypePointer Input %12
29+
%4 = OpVariable %20 Input
30+
%5 = OpVariable %19 Input
31+
%6 = OpVariable %19 Input
32+
%7 = OpVariable %19 Input
33+
%21 = OpTypeBool
34+
%22 = OpTypePointer StorageBuffer %12

0 commit comments

Comments
 (0)