scalar or vector: adjust subgroup_all_equal to accept composites

Author: Firestar99

crates/spirv-std/src/arch/subgroup.rs

@@ -244,24 +244,35 @@ pub fn subgroup_any(predicate: bool) -> bool {
 #[spirv_std_macros::gpu_only]
 #[doc(alias = "OpGroupNonUniformAllEqual")]
 #[inline]
-pub fn subgroup_all_equal<T: ScalarOrVector>(value: T) -> bool {
-    let mut result = false;
+pub fn subgroup_all_equal<T: ScalarOrVectorComposite>(value: T) -> bool {
+    struct Transform(bool);
 
-    unsafe {
-        asm! {
-            "%bool = OpTypeBool",
-            "%u32 = OpTypeInt 32 0",
-            "%subgroup = OpConstant %u32 {subgroup}",
-            "%value = OpLoad _ {value}",
-            "%result = OpGroupNonUniformAllEqual %bool %subgroup %value",
-            "OpStore {result} %result",
-            subgroup = const SUBGROUP,
-            value = in(reg) &value,
-            result = in(reg) &mut result,
+    impl ScalarOrVectorTransform for Transform {
+        #[inline]
+        fn transform<T: ScalarOrVector>(&mut self, value: T) -> T {
+            let mut result = false;
+            unsafe {
+                asm! {
+                    "%bool = OpTypeBool",
+                    "%u32 = OpTypeInt 32 0",
+                    "%subgroup = OpConstant %u32 {subgroup}",
+                    "%value = OpLoad _ {value}",
+                    "%result = OpGroupNonUniformAllEqual %bool %subgroup %value",
+                    "OpStore {result} %result",
+                    subgroup = const SUBGROUP,
+                    value = in(reg) &value,
+                    result = in(reg) &mut result,
+                }
+            }
+            self.0 &= result;
+            value
         }
     }
 
-    result
+    let mut transform = Transform(true);
+    // ignore returned value
+    value.transform(&mut transform);
+    transform.0
 }
 
 /// Result is the `value` of the invocation identified by the id `id` to all active invocations in the group.

tests/compiletests/ui/arch/subgroup/subgroup_composite_all_equals.rs

@@ -0,0 +1,46 @@
+// build-pass
+// compile-flags: -C target-feature=+GroupNonUniform,+GroupNonUniformVote,+ext:SPV_KHR_vulkan_memory_model
+// compile-flags: -C llvm-args=--disassemble-fn=subgroup_composite_all_equals::disassembly
+// normalize-stderr-test "OpLine .*\n" -> ""
+
+use glam::*;
+use spirv_std::ScalarOrVectorComposite;
+use spirv_std::arch::*;
+use spirv_std::spirv;
+
+#[derive(Copy, Clone, ScalarOrVectorComposite)]
+pub struct MyStruct {
+    a: f32,
+    b: UVec3,
+    c: Nested,
+    d: Zst,
+}
+
+#[derive(Copy, Clone, ScalarOrVectorComposite)]
+pub struct Nested(i32);
+
+#[derive(Copy, Clone, ScalarOrVectorComposite)]
+pub struct Zst;
+
+/// this should be 3 `subgroup_all_equal` instructions, with all calls inlined
+fn disassembly(my_struct: MyStruct) -> bool {
+    subgroup_all_equal(my_struct)
+}
+
+#[spirv(compute(threads(32)))]
+pub fn main(
+    #[spirv(local_invocation_index)] inv_id: UVec3,
+    #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] output: &mut u32,
+) {
+    unsafe {
+        let my_struct = MyStruct {
+            a: inv_id.x as f32,
+            b: inv_id,
+            c: Nested(5i32 - inv_id.x as i32),
+            d: Zst,
+        };
+
+        let bool = disassembly(my_struct);
+        *output = u32::from(bool);
+    }
+}

tests/compiletests/ui/arch/subgroup/subgroup_composite_all_equals.stderr

@@ -0,0 +1,15 @@
+%1 = OpFunction  %2  None %3
+%4 = OpFunctionParameter  %5
+%6 = OpLabel
+%8 = OpCompositeExtract  %9  %4 0
+%11 = OpGroupNonUniformAllEqual  %2  %12 %8
+%13 = OpLogicalAnd  %2  %14 %11
+%15 = OpCompositeExtract  %16  %4 1
+%17 = OpGroupNonUniformAllEqual  %2  %12 %15
+%18 = OpLogicalAnd  %2  %13 %17
+%19 = OpCompositeExtract  %20  %4 2
+%21 = OpGroupNonUniformAllEqual  %2  %12 %19
+%22 = OpLogicalAnd  %2  %18 %21
+OpNoLine
+OpReturnValue %22
+OpFunctionEnd