spirv-std functions to query compute builtins

crates/rustc_codegen_spirv/src/linker/duplicates.rs

@@ -3,7 +3,6 @@ use rspirv::binary::Assemble;
 use rspirv::dr::{Instruction, Module, Operand};
 use rspirv::spirv::{Op, Word};
 use rustc_data_structures::fx::{FxHashMap, FxHashSet};
-use rustc_middle::bug;
 use smallvec::SmallVec;
 use std::collections::hash_map;
 use std::mem;
@@ -104,24 +103,10 @@ fn gather_annotations(annotations: &[Instruction]) -> FxHashMap<Word, Vec<u32>>
         .collect()
 }
 
-fn gather_names(debug_names: &[Instruction]) -> FxHashMap<Word, String> {
-    debug_names
-        .iter()
-        .filter(|inst| inst.class.opcode == Op::Name)
-        .map(|inst| {
-            (
-                inst.operands[0].unwrap_id_ref(),
-                inst.operands[1].unwrap_literal_string().to_owned(),
-            )
-        })
-        .collect()
-}
-
 fn make_dedupe_key(
     inst: &Instruction,
     unresolved_forward_pointers: &FxHashSet<Word>,
     annotations: &FxHashMap<Word, Vec<u32>>,
-    names: &FxHashMap<Word, String>,
 ) -> Vec<u32> {
     let mut data = vec![inst.class.opcode as u32];
 
@@ -144,29 +129,10 @@ fn make_dedupe_key(
             op.assemble_into(&mut data);
         }
     }
-    if let Some(id) = inst.result_id {
-        if let Some(annos) = annotations.get(&id) {
-            data.extend_from_slice(annos);
-        }
-        if inst.class.opcode == Op::Variable {
-            // Names only matter for OpVariable.
-            if let Some(name) = names.get(&id) {
-                // Jump through some hoops to shove a String into a Vec<u32>.
-                //
-                // FIXME(eddyb) this should `.assemble_into(&mut data)` the
-                // `Operand::LiteralString(...)` from the original `Op::Name`.
-                for chunk in name.as_bytes().chunks(4) {
-                    let slice = match *chunk {
-                        [a] => [a, 0, 0, 0],
-                        [a, b] => [a, b, 0, 0],
-                        [a, b, c] => [a, b, c, 0],
-                        [a, b, c, d] => [a, b, c, d],
-                        _ => bug!(),
-                    };
-                    data.push(u32::from_le_bytes(slice));
-                }
-            }
-        }
+    if let Some(id) = inst.result_id
+        && let Some(annos) = annotations.get(&id)
+    {
+        data.extend_from_slice(annos);
     }
 
     data
@@ -198,7 +164,6 @@ pub fn remove_duplicate_types(module: &mut Module) {
 
     // Collect a map from type ID to an annotation "key blob" (to append to the type key)
     let annotations = gather_annotations(&module.annotations);
-    let names = gather_names(&module.debug_names);
 
     for inst in &mut module.types_global_values {
         if inst.class.opcode == Op::TypeForwardPointer
@@ -222,7 +187,7 @@ pub fn remove_duplicate_types(module: &mut Module) {
         // all_inst_iter_mut pass below. However, the code is a lil bit cleaner this way I guess.
         rewrite_inst_with_rules(inst, &rewrite_rules);
 
-        let key = make_dedupe_key(inst, &unresolved_forward_pointers, &annotations, &names);
+        let key = make_dedupe_key(inst, &unresolved_forward_pointers, &annotations);
 
         match key_to_result_id.entry(key) {
             hash_map::Entry::Vacant(entry) => {

crates/spirv-std/src/builtin.rs

@@ -0,0 +1,60 @@
+//! Functionality to declare builtins, mostly proc macros
+//!
+//! # Making built-in functions for `spirv-std`
+//!
+//! Usually, built-ins are implemented as freestanding functions in `spirv-std`. We like to keep function declaration
+//! outside the macro to make it easier for users to browse the source code.
+//!
+//! Example on how to declare an Input Built-in:
+//! ```no_run
+//! # use spirv_std_macros::gpu_only;
+//! #
+//! /// GLSL docs short description in #Name section. Remove the first "Contains " since we're using getters instead
+//! /// of globals, capitalize and add a dot to the end.
+//! ///
+//! /// GLSL docs full #Description section.
+//! ///
+//! /// We're using GLSL documentation of this built-in, which is usually more descriptive than the SPIR-V or WGSL docs.
+//! /// Change all references to link with rust-gpu intrinsics.
+//! ///
+//! /// Update the links of GLSL and WGSL to reference the correct page, keep SPIR-V as is. GLSL may link to the
+//! /// [reference](https://registry.khronos.org/OpenGL-Refpages/gl4/) or to the
+//! /// [glsl extensions github repo](https://github.com/KhronosGroup/GLSL/tree/main/extensions).
+//! /// * GLSL: [`gl_MyBuiltIn`](https://registry.khronos.org/OpenGL-Refpages/gl4/html/gl_LocalInvocationID.xhtml)
+//! /// * WGSL: [`my_built_in`](https://www.w3.org/TR/WGSL/#local-invocation-id-builtin-value)
+//! /// * SPIRV: [`MyBuiltIn`](https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_builtin)
+//! #[doc(alias = "gl_MyBuiltIn")]
+//! #[doc(alias = "MyBuiltIn")]
+//! #[inline]
+//! #[gpu_only]
+//! pub fn my_built_in() -> u32 {
+//!     crate::load_builtin!(MyBuiltIn)
+//! }
+//! ```
+//!
+//! Reference links:
+//! * [WGSL specification describing builtins](https://www.w3.org/TR/WGSL/#builtin-inputs-outputs)
+//! * [SPIR-V specification for builtins](https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_builtin)
+//! * [GLSL reference](https://registry.khronos.org/OpenGL-Refpages/gl4/)
+//!   * [GLSL reference source code](https://github.com/KhronosGroup/OpenGL-Refpages/tree/main/gl4)
+//!   * [GLSL extensions](https://github.com/KhronosGroup/GLSL/tree/main/extensions)
+
+/// Query SPIR-V (read-only global) built-in values
+///
+/// See [module level documentation] on how to use these.
+#[macro_export]
+macro_rules! load_builtin {
+    ($name:ident $(: $ty:ty)?) => {
+        unsafe {
+            let mut result $(: $ty)? = Default::default();
+            ::core::arch::asm! {
+                "%builtin = OpVariable typeof{result_ref} Input",
+                concat!("OpDecorate %builtin BuiltIn ", stringify!($name)),
+                "%result = OpLoad typeof*{result_ref} %builtin",
+                "OpStore {result_ref} %result",
+                result_ref = in(reg) &mut result,
+            }
+            result
+        }
+    };
+}

crates/spirv-std/src/compute.rs

@@ -0,0 +1,93 @@
+//! compute shader built-ins
+
+use glam::UVec3;
+
+/// The index of work item currently being operated on by a compute shader.
+///
+/// In the compute language, [`local_invocation_id`] is an input variable containing the n-dimensional index of the
+/// local work invocation within the work group that the current shader is executing in. The possible values for this
+/// variable range across the local work group size, i.e., `(0,0,0)` to
+/// `(workgroup_size.x - 1, workgroup_size.y - 1, workgroup_size.z - 1)`.
+///
+/// * GLSL: [`gl_LocalInvocationID`](https://registry.khronos.org/OpenGL-Refpages/gl4/html/gl_LocalInvocationID.xhtml)
+/// * WGSL: [`local_invocation_id`](https://www.w3.org/TR/WGSL/#local-invocation-id-builtin-value)
+/// * SPIR-V: [`LocalInvocationId`](https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_builtin)
+#[doc(alias = "gl_LocalInvocationID")]
+#[doc(alias = "LocalInvocationId")]
+#[inline]
+#[gpu_only]
+pub fn local_invocation_id() -> UVec3 {
+    crate::load_builtin!(LocalInvocationId)
+}
+
+/// The local linear index of work item currently being operated on by a compute shader.
+///
+/// In the compute language, [`local_invocation_index`] is a derived input variable containing the 1-dimensional
+/// linearized index of the work invocation within the work group that the current shader is executing on. The value of
+/// [`local_invocation_index`] is equal to [`local_invocation_id`]`.z * workgroup_size.x * workgroup_size.y`
+/// `+ `[`local_invocation_id`]`.y * workgroup_size.x + `[`local_invocation_id`]`.x`.
+///
+/// * GLSL: [`gl_LocalInvocationIndex`](https://registry.khronos.org/OpenGL-Refpages/gl4/html/gl_LocalInvocationIndex.xhtml)
+/// * WGSL: [`local_invocation_index`](https://www.w3.org/TR/WGSL/#local-invocation-index-builtin-value)
+/// * SPIR-V: [`LocalInvocationIndex`](https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_builtin)
+#[doc(alias = "gl_LocalInvocationIndex")]
+#[doc(alias = "LocalInvocationIndex")]
+#[inline]
+#[gpu_only]
+pub fn local_invocation_index() -> u32 {
+    crate::load_builtin!(LocalInvocationIndex)
+}
+
+/// The global index of work item currently being operated on by a compute shader.
+///
+/// In the compute language, [`global_invocation_id`] is a derived input variable containing the n-dimensional index of
+/// the work invocation within the global work group that the current shader is executing on. The value of
+/// [`global_invocation_id`] is equal to [`workgroup_id`]` * workgroup_size + `[`local_invocation_id`].
+///
+/// * GLSL: [`gl_GlobalInvocationID`](https://registry.khronos.org/OpenGL-Refpages/gl4/html/gl_GlobalInvocationID.xhtml)
+/// * WGSL: [`global_invocation_id`](https://www.w3.org/TR/WGSL/#global-invocation-index-builtin-value)
+/// * SPIR-V: [`GlobalInvocationId`](https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_builtin)
+#[doc(alias = "gl_GlobalInvocationID")]
+#[doc(alias = "GlobalInvocationId")]
+#[inline]
+#[gpu_only]
+pub fn global_invocation_id() -> UVec3 {
+    crate::load_builtin!(GlobalInvocationId)
+}
+
+// custom: do not mention `glDispatchCompute` directly, be more general across APIs
+/// The number of workgroups that have been dispatched to a compute shader.
+///
+/// In the compute language, [`num_workgroups`] contains the total number of work groups that will execute the compute
+/// shader. The components of [`num_workgroups`] are equal to the `x`, `y`, and `z` parameters passed to the dispatch
+/// command.
+///
+/// * GLSL: [`gl_NumWorkGroups`](https://registry.khronos.org/OpenGL-Refpages/gl4/html/gl_NumWorkGroups.xhtml)
+/// * WGSL: [`num_workgroups`](https://www.w3.org/TR/WGSL/#num-workgroups-builtin-value)
+/// * SPIR-V: [`NumWorkgroups`](https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_builtin)
+#[doc(alias = "gl_NumWorkGroups")]
+#[doc(alias = "NumWorkgroups")]
+#[inline]
+#[gpu_only]
+pub fn num_workgroups() -> UVec3 {
+    crate::load_builtin!(NumWorkgroups)
+}
+
+// custom: do not mention `glDispatchCompute` directly, be more general across APIs
+/// The index of the workgroup currently being operated on by a compute shader.
+///
+/// In the compute language, [`workgroup_id`] contains the 3-dimensional index of the global work group that the current
+/// compute shader invocation is executing within. The possible values range across the parameters passed into the
+/// dispatch command, i.e., from `(0, 0, 0)` to
+/// `(`[`num_workgroups`]`.x - 1, `[`num_workgroups`]`.y - 1, `[`num_workgroups`]`.z - 1)`.
+///
+/// * GLSL: [`gl_WorkGroupID`](https://registry.khronos.org/OpenGL-Refpages/gl4/html/gl_WorkGroupID.xhtml)
+/// * WGSL: [`workgroup_id`](https://www.w3.org/TR/WGSL/#workgroup-id-builtin-value)
+/// * SPIR-V: [`WorkgroupId`](https://registry.khronos.org/SPIR-V/specs/unified1/SPIRV.html#_builtin)
+#[doc(alias = "gl_WorkGroupID")]
+#[doc(alias = "WorkgroupId")]
+#[inline]
+#[gpu_only]
+pub fn workgroup_id() -> UVec3 {
+    crate::load_builtin!(WorkgroupId)
+}

crates/spirv-std/src/lib.rs

@@ -93,7 +93,9 @@ pub use macros::{debug_printf, debug_printfln};
 pub mod arch;
 pub mod atomic;
 pub mod barrier;
+pub mod builtin;
 pub mod byte_addressable_buffer;
+pub mod compute;
 pub mod debug_printf;
 pub mod float;
 pub mod fragment;