compiletest: shared memory reductions, using the same type in buffers and as shared memory

Author: Firestar99

tests/compiletests/ui/arch/shared/reduction_array.rs

@@ -0,0 +1,76 @@
+// build-pass
+
+use core::ops::{Add, AddAssign, Deref, DerefMut};
+use spirv_std::arch::workgroup_memory_barrier_with_group_sync;
+use spirv_std::glam::*;
+use spirv_std::spirv;
+
+#[derive(Copy, Clone, Debug)]
+pub struct Value(pub [f32; 4]);
+
+impl Deref for Value {
+    type Target = [f32; 4];
+
+    fn deref(&self) -> &Self::Target {
+        &self.0
+    }
+}
+
+impl DerefMut for Value {
+    fn deref_mut(&mut self) -> &mut Self::Target {
+        &mut self.0
+    }
+}
+
+impl Add for Value {
+    type Output = Self;
+
+    fn add(self, rhs: Self) -> Self::Output {
+        Self([
+            self[0] + rhs[0],
+            self[1] + rhs[1],
+            self[2] + rhs[2],
+            self[3] + rhs[3],
+        ])
+    }
+}
+
+impl AddAssign for Value {
+    fn add_assign(&mut self, rhs: Self) {
+        *self = *self + rhs;
+    }
+}
+
+pub const WG_SIZE_SHIFT: usize = 5;
+pub const WG_SIZE: usize = 1 << WG_SIZE_SHIFT;
+
+// threads must be a literal, constants don't work
+const _: () = {
+    assert!(WG_SIZE == 32);
+};
+#[spirv(compute(threads(32)))]
+pub fn main(
+    #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] input: &[Value],
+    #[spirv(descriptor_set = 0, binding = 1, storage_buffer)] output: &mut Value,
+    #[spirv(workgroup)] shared: &mut [Value; WG_SIZE],
+    #[spirv(local_invocation_index)] inv_id: UVec3,
+) {
+    unsafe {
+        let inv_id = inv_id.x as usize;
+        shared[inv_id] = input[inv_id];
+        workgroup_memory_barrier_with_group_sync();
+
+        let mut mask = WG_SIZE << 1;
+        while mask != 0 {
+            if inv_id < mask {
+                shared[inv_id] += shared[inv_id + mask];
+            }
+            workgroup_memory_barrier_with_group_sync();
+            mask <<= 1;
+        }
+
+        if inv_id == 0 {
+            *output = shared[0];
+        }
+    }
+}

tests/compiletests/ui/arch/shared/reduction_big_struct.rs

@@ -0,0 +1,84 @@
+// build-pass
+
+use core::ops::{Add, AddAssign};
+use spirv_std::arch::workgroup_memory_barrier_with_group_sync;
+use spirv_std::glam::*;
+use spirv_std::spirv;
+
+#[derive(Copy, Clone, Debug)]
+pub struct Nested(pub i32);
+
+impl Add for Nested {
+    type Output = Self;
+
+    fn add(self, rhs: Self) -> Self::Output {
+        Self(self.0 + rhs.0)
+    }
+}
+
+impl AddAssign for Nested {
+    fn add_assign(&mut self, rhs: Self) {
+        *self = *self + rhs;
+    }
+}
+
+#[derive(Copy, Clone, Debug)]
+pub struct Value {
+    pub a: f32,
+    pub b: UVec4,
+    pub c: Nested,
+    pub d: Mat4,
+}
+
+impl Add for Value {
+    type Output = Self;
+
+    fn add(self, rhs: Self) -> Self::Output {
+        Self {
+            a: self.a + rhs.a,
+            b: self.b + rhs.b,
+            c: self.c + rhs.c,
+            d: self.d + rhs.d,
+        }
+    }
+}
+
+impl AddAssign for Value {
+    fn add_assign(&mut self, rhs: Self) {
+        *self = *self + rhs;
+    }
+}
+
+pub const WG_SIZE_SHIFT: usize = 5;
+pub const WG_SIZE: usize = 1 << WG_SIZE_SHIFT;
+
+// threads must be a literal, constants don't work
+const _: () = {
+    assert!(WG_SIZE == 32);
+};
+#[spirv(compute(threads(32)))]
+pub fn main(
+    #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] input: &[Value],
+    #[spirv(descriptor_set = 0, binding = 1, storage_buffer)] output: &mut Value,
+    #[spirv(workgroup)] shared: &mut [Value; WG_SIZE],
+    #[spirv(local_invocation_index)] inv_id: UVec3,
+) {
+    unsafe {
+        let inv_id = inv_id.x as usize;
+        shared[inv_id] = input[inv_id];
+        workgroup_memory_barrier_with_group_sync();
+
+        let mut mask = WG_SIZE << 1;
+        while mask != 0 {
+            if inv_id < mask {
+                shared[inv_id] += shared[inv_id + mask];
+            }
+            workgroup_memory_barrier_with_group_sync();
+            mask <<= 1;
+        }
+
+        if inv_id == 0 {
+            *output = shared[0];
+        }
+    }
+}

tests/compiletests/ui/arch/shared/reduction_u32.rs

@@ -0,0 +1,41 @@
+// build-pass
+
+use spirv_std::arch::workgroup_memory_barrier_with_group_sync;
+use spirv_std::glam::*;
+use spirv_std::spirv;
+
+pub type Value = i32;
+
+pub const WG_SIZE_SHIFT: usize = 5;
+pub const WG_SIZE: usize = 1 << WG_SIZE_SHIFT;
+
+// threads must be a literal, constants don't work
+const _: () = {
+    assert!(WG_SIZE == 32);
+};
+#[spirv(compute(threads(32)))]
+pub fn main(
+    #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] input: &[Value],
+    #[spirv(descriptor_set = 0, binding = 1, storage_buffer)] output: &mut Value,
+    #[spirv(workgroup)] shared: &mut [Value; WG_SIZE],
+    #[spirv(local_invocation_index)] inv_id: UVec3,
+) {
+    unsafe {
+        let inv_id = inv_id.x as usize;
+        shared[inv_id] = input[inv_id];
+        workgroup_memory_barrier_with_group_sync();
+
+        let mut mask = WG_SIZE << 1;
+        while mask != 0 {
+            if inv_id < mask {
+                shared[inv_id] += shared[inv_id + mask];
+            }
+            workgroup_memory_barrier_with_group_sync();
+            mask <<= 1;
+        }
+
+        if inv_id == 0 {
+            *output = shared[0];
+        }
+    }
+}

tests/compiletests/ui/arch/shared/reduction_vec.rs

@@ -0,0 +1,41 @@
+// build-pass
+
+use spirv_std::arch::workgroup_memory_barrier_with_group_sync;
+use spirv_std::glam::*;
+use spirv_std::spirv;
+
+pub type Value = UVec4;
+
+pub const WG_SIZE_SHIFT: usize = 5;
+pub const WG_SIZE: usize = 1 << WG_SIZE_SHIFT;
+
+// threads must be a literal, constants don't work
+const _: () = {
+    assert!(WG_SIZE == 32);
+};
+#[spirv(compute(threads(32)))]
+pub fn main(
+    #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] input: &[Value],
+    #[spirv(descriptor_set = 0, binding = 1, storage_buffer)] output: &mut Value,
+    #[spirv(workgroup)] shared: &mut [Value; WG_SIZE],
+    #[spirv(local_invocation_index)] inv_id: UVec3,
+) {
+    unsafe {
+        let inv_id = inv_id.x as usize;
+        shared[inv_id] = input[inv_id];
+        workgroup_memory_barrier_with_group_sync();
+
+        let mut mask = WG_SIZE << 1;
+        while mask != 0 {
+            if inv_id < mask {
+                shared[inv_id] += shared[inv_id + mask];
+            }
+            workgroup_memory_barrier_with_group_sync();
+            mask <<= 1;
+        }
+
+        if inv_id == 0 {
+            *output = shared[0];
+        }
+    }
+}