Add an read_clock_khr function that calls OpReadClockKHR (#757)

* Add support for OpReadClockKHR

* Use a const scope

* Use full scope import path

* Apply some suggestions

* Add cfg flags and scopes as raw u32s

* Run cargo fmt

* Fix comment accuracy

* Update crates/rustc_codegen_spirv/src/spirv_type_constraints.rs

Co-authored-by: Ashley Hauck <953151+khyperia@users.noreply.github.com>

* Run rustfmt

* Add the shader clock feature and capability to the compile test rust flags and bless the changed test errors

Co-authored-by: Ashley Hauck <953151+khyperia@users.noreply.github.com>

crates/rustc_codegen_spirv/src/spirv_type_constraints.rs

@@ -745,7 +745,10 @@ pub fn instruction_signatures(op: Op) -> Option<&'static [InstSig<'static>]> {
         // SPV_AMD_shader_fragment_mask
         Op::FragmentMaskFetchAMD | Op::FragmentFetchAMD => reserved!(SPV_AMD_shader_fragment_mask),
         // SPV_KHR_shader_clock
-        Op::ReadClockKHR => reserved!(SPV_KHR_shader_clock),
+        Op::ReadClockKHR => {
+            // NOTE(eddyb) we actually use these despite not being in the standard yet.
+            // reserved!(SPV_KHR_shader_clock)
+        }
         // SPV_NV_mesh_shader
         Op::WritePackedPrimitiveIndices4x8NV => reserved!(SPV_NV_mesh_shader),
         // SPV_NV_ray_tracing

crates/spirv-std/src/arch.rs

@@ -150,3 +150,53 @@ pub unsafe fn vector_insert_dynamic<T: Scalar, V: Vector<T, N>, const N: usize>(
 pub fn kill() -> ! {
     unsafe { asm!("OpKill", options(noreturn)) }
 }
+
+/// Read from the shader clock with either the `Subgroup` or `Device` scope.
+///
+/// See:
+/// <https://htmlpreview.github.io/?https://github.com/KhronosGroup/SPIRV-Registry/blob/master/extensions/KHR/SPV_KHR_shader_clock.html>
+#[cfg(all(
+    target_feature = "Int64",
+    target_feature = "ShaderClockKHR",
+    target_feature = "ext:SPV_KHR_shader_clock"
+))]
+#[spirv_std_macros::gpu_only]
+#[doc(alias = "OpReadClockKHR")]
+pub unsafe fn read_clock_khr<const SCOPE: u32>() -> u64 {
+    let mut result: u64;
+
+    asm! {
+        "%uint = OpTypeInt 32 0",
+        "%scope = OpConstant %uint {scope}",
+        "{result} = OpReadClockKHR typeof*{result} %scope",
+        result = out(reg) result,
+        scope = const SCOPE,
+    };
+
+    result
+}
+
+/// Like `read_clock_khr` but returns a vector to avoid requiring the `Int64`
+/// capability. It returns a 'vector of two-components of 32-bit unsigned
+/// integer type with the first component containing the 32 least significant
+/// bits and the second component containing the 32 most significant bits.'
+#[cfg(all(
+    target_feature = "ShaderClockKHR",
+    target_feature = "ext:SPV_KHR_shader_clock"
+))]
+#[spirv_std_macros::gpu_only]
+#[doc(alias = "OpReadClockKHR")]
+pub unsafe fn read_clock_uvec2_khr<V: Vector<u32, 2>, const SCOPE: u32>() -> V {
+    let mut result = V::default();
+
+    asm! {
+        "%uint = OpTypeInt 32 0",
+        "%scope = OpConstant %uint {scope}",
+        "%result = OpReadClockKHR typeof*{result} %scope",
+        "OpStore {result} %result",
+        result = in(reg) &mut result,
+        scope = const SCOPE,
+    };
+
+    result
+}

tests/src/main.rs

@@ -294,13 +294,23 @@ struct TestDeps {
 /// The RUSTFLAGS passed to all SPIR-V builds.
 // FIXME(eddyb) expose most of these from `spirv-builder`.
 fn rust_flags(codegen_backend_path: &Path) -> String {
+    let target_features = [
+        "Int8",
+        "Int16",
+        "Int64",
+        "Float64",
+        // Only needed for `ui/arch/read_clock_khr.rs`.
+        "ShaderClockKHR",
+        "ext:SPV_KHR_shader_clock",
+    ];
+
     [
         &*format!("-Zcodegen-backend={}", codegen_backend_path.display()),
         "-Coverflow-checks=off",
         "-Cdebug-assertions=off",
         "-Cdebuginfo=2",
         "-Cembed-bitcode=no",
-        "-Ctarget-feature=+Int8,+Int16,+Int64,+Float64",
+        &format!("-Ctarget-feature=+{}", target_features.join(",+")),
         "-Zsymbol-mangling-version=v0",
     ]
     .join(" ")

tests/ui/arch/read_clock_khr.rs

@@ -0,0 +1,16 @@
+// build-pass
+// compile-flags: -Ctarget-feature=+Int64,+ShaderClockKHR,+ext:SPV_KHR_shader_clock
+
+use glam::UVec2;
+use spirv_std::{
+    arch::{read_clock_khr, read_clock_uvec2_khr},
+    memory::Scope,
+};
+
+#[spirv(fragment)]
+pub fn main() {
+    let clock_time = unsafe { read_clock_khr::<{ Scope::Subgroup as u32 }>() };
+
+    let clock_time_uvec2: UVec2 =
+        unsafe { read_clock_uvec2_khr::<_, { Scope::Subgroup as u32 }>() };
+}

tests/ui/dis/asm_op_decorate.stderr

@@ -3,8 +3,10 @@ OpCapability Int16
 OpCapability Int64
 OpCapability Int8
 OpCapability RuntimeDescriptorArray
+OpCapability ShaderClockKHR
 OpCapability Shader
 OpExtension "SPV_EXT_descriptor_indexing"
+OpExtension "SPV_KHR_shader_clock"
 OpMemoryModel Logical Simple
 OpEntryPoint Fragment %1 "main"
 OpExecutionMode %1 OriginUpperLeft

tests/ui/dis/custom_entry_point.stderr

@@ -2,7 +2,9 @@ OpCapability Float64
 OpCapability Int16
 OpCapability Int64
 OpCapability Int8
+OpCapability ShaderClockKHR
 OpCapability Shader
+OpExtension "SPV_KHR_shader_clock"
 OpMemoryModel Logical Simple
 OpEntryPoint Fragment %1 "hello_world"
 OpExecutionMode %1 OriginUpperLeft

tests/ui/dis/generic-fn-op-name.stderr

@@ -2,7 +2,9 @@ OpCapability Float64
 OpCapability Int16
 OpCapability Int64
 OpCapability Int8
+OpCapability ShaderClockKHR
 OpCapability Shader
+OpExtension "SPV_KHR_shader_clock"
 OpMemoryModel Logical Simple
 OpEntryPoint Fragment %1 "main"
 OpExecutionMode %1 OriginUpperLeft

tests/ui/dis/issue-723-indirect-input.stderr

@@ -2,7 +2,9 @@ OpCapability Float64
 OpCapability Int16
 OpCapability Int64
 OpCapability Int8
+OpCapability ShaderClockKHR
 OpCapability Shader
+OpExtension "SPV_KHR_shader_clock"
 OpMemoryModel Logical Simple
 OpEntryPoint Fragment %1 "main" %2
 OpExecutionMode %1 OriginUpperLeft

tests/ui/dis/issue-723-output.stderr

@@ -2,7 +2,9 @@ OpCapability Float64
 OpCapability Int16
 OpCapability Int64
 OpCapability Int8
+OpCapability ShaderClockKHR
 OpCapability Shader
+OpExtension "SPV_KHR_shader_clock"
 OpMemoryModel Logical Simple
 OpEntryPoint Fragment %1 "main" %2
 OpExecutionMode %1 OriginUpperLeft