Make byte addressable buffer take &self, add support for matrix (#749)

crates/rustc_codegen_spirv/src/builder/byte_addressable_buffer.rs

@@ -51,7 +51,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
     }
 
     #[allow(clippy::too_many_arguments)]
-    fn load_vec_or_arr(
+    fn load_vec_mat_arr(
         &mut self,
         original_type: Word,
         result_type: Word,
@@ -104,21 +104,22 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 let val = self.load_u32(array, dynamic_word_index, constant_word_offset);
                 self.bitcast(val, result_type)
             }
-            SpirvType::Vector { element, count } => self.load_vec_or_arr(
-                original_type,
-                result_type,
-                array,
-                dynamic_word_index,
-                constant_word_offset,
-                element,
-                count,
-            ),
+            SpirvType::Vector { element, count } | SpirvType::Matrix { element, count } => self
+                .load_vec_mat_arr(
+                    original_type,
+                    result_type,
+                    array,
+                    dynamic_word_index,
+                    constant_word_offset,
+                    element,
+                    count,
+                ),
             SpirvType::Array { element, count } => {
                 let count = match self.builder.lookup_const_u64(count) {
                     Some(count) => count as u32,
                     None => return self.load_err(original_type, result_type),
                 };
-                self.load_vec_or_arr(
+                self.load_vec_mat_arr(
                     original_type,
                     result_type,
                     array,
@@ -229,7 +230,7 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
     }
 
     #[allow(clippy::too_many_arguments)]
-    fn store_vec_or_arr(
+    fn store_vec_mat_arr(
         &mut self,
         original_type: Word,
         value: SpirvValue,
@@ -278,21 +279,22 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
                 let value_u32 = self.bitcast(value, u32_ty);
                 self.store_u32(array, dynamic_word_index, constant_word_offset, value_u32);
             }
-            SpirvType::Vector { element, count } => self.store_vec_or_arr(
-                original_type,
-                value,
-                array,
-                dynamic_word_index,
-                constant_word_offset,
-                element,
-                count,
-            ),
+            SpirvType::Vector { element, count } | SpirvType::Matrix { element, count } => self
+                .store_vec_mat_arr(
+                    original_type,
+                    value,
+                    array,
+                    dynamic_word_index,
+                    constant_word_offset,
+                    element,
+                    count,
+                ),
             SpirvType::Array { element, count } => {
                 let count = match self.builder.lookup_const_u64(count) {
                     Some(count) => count as u32,
                     None => return self.store_err(original_type, value),
                 };
-                self.store_vec_or_arr(
+                self.store_vec_mat_arr(
                     original_type,
                     value,
                     array,

crates/spirv-std/src/byte_addressable_buffer.rs

@@ -43,7 +43,7 @@ impl<'a> ByteAddressableBuffer<'a> {
     /// This function allows writing a type to an untyped buffer, then reading a different type
     /// from the same buffer, allowing all sorts of safety guarantees to be bypassed (effectively a
     /// transmute)
-    pub unsafe fn load<T>(self, byte_index: u32) -> T {
+    pub unsafe fn load<T>(&self, byte_index: u32) -> T {
         if byte_index + mem::size_of::<T>() as u32 > self.data.len() as u32 {
             panic!("Index out of range")
         }
@@ -58,7 +58,7 @@ impl<'a> ByteAddressableBuffer<'a> {
     /// This function allows writing a type to an untyped buffer, then reading a different type
     /// from the same buffer, allowing all sorts of safety guarantees to be bypassed (effectively a
     /// transmute). Additionally, bounds checking is not performed.
-    pub unsafe fn load_unchecked<T>(self, byte_index: u32) -> T {
+    pub unsafe fn load_unchecked<T>(&self, byte_index: u32) -> T {
         buffer_load_intrinsic(self.data, byte_index)
     }
 
@@ -69,7 +69,7 @@ impl<'a> ByteAddressableBuffer<'a> {
     /// This function allows writing a type to an untyped buffer, then reading a different type
     /// from the same buffer, allowing all sorts of safety guarantees to be bypassed (effectively a
     /// transmute)
-    pub unsafe fn store<T>(self, byte_index: u32, value: T) {
+    pub unsafe fn store<T>(&mut self, byte_index: u32, value: T) {
         if byte_index + mem::size_of::<T>() as u32 > self.data.len() as u32 {
             panic!("Index out of range")
         }
@@ -84,7 +84,7 @@ impl<'a> ByteAddressableBuffer<'a> {
     /// This function allows writing a type to an untyped buffer, then reading a different type
     /// from the same buffer, allowing all sorts of safety guarantees to be bypassed (effectively a
     /// transmute). Additionally, bounds checking is not performed.
-    pub unsafe fn store_unchecked<T>(self, byte_index: u32, value: T) {
+    pub unsafe fn store_unchecked<T>(&mut self, byte_index: u32, value: T) {
         buffer_store_intrinsic(self.data, byte_index, value);
     }
 }

tests/ui/byte_addressable_buffer/arr.rs

@@ -19,7 +19,7 @@ pub fn store(
     val: [i32; 4],
 ) {
     unsafe {
-        let buf = ByteAddressableBuffer::new(buf);
+        let mut buf = ByteAddressableBuffer::new(buf);
         buf.store(5, val);
     }
 }

tests/ui/byte_addressable_buffer/big_struct.rs

@@ -28,7 +28,7 @@ pub fn store(
     val: BigStruct,
 ) {
     unsafe {
-        let buf = ByteAddressableBuffer::new(buf);
+        let mut buf = ByteAddressableBuffer::new(buf);
         buf.store(5, val);
     }
 }

tests/ui/byte_addressable_buffer/complex.rs

@@ -34,7 +34,7 @@ pub fn store(
     val: Nesty,
 ) {
     unsafe {
-        let buf = ByteAddressableBuffer::new(buf);
+        let mut buf = ByteAddressableBuffer::new(buf);
         buf.store(5, val);
     }
 }

tests/ui/byte_addressable_buffer/f32.rs

@@ -16,7 +16,7 @@ pub fn load(
 #[spirv(fragment)]
 pub fn store(#[spirv(descriptor_set = 0, binding = 0, storage_buffer)] buf: &mut [u32], val: f32) {
     unsafe {
-        let buf = ByteAddressableBuffer::new(buf);
+        let mut buf = ByteAddressableBuffer::new(buf);
         buf.store(5, val);
     }
 }

tests/ui/byte_addressable_buffer/u32.rs

@@ -16,7 +16,7 @@ pub fn load(
 #[spirv(fragment)]
 pub fn store(#[spirv(descriptor_set = 0, binding = 0, storage_buffer)] buf: &mut [u32], val: u32) {
     unsafe {
-        let buf = ByteAddressableBuffer::new(buf);
+        let mut buf = ByteAddressableBuffer::new(buf);
         buf.store(5, val);
     }
 }

tests/ui/byte_addressable_buffer/vec.rs

@@ -2,21 +2,36 @@
 
 use spirv_std::{glam::Vec4, ByteAddressableBuffer};
 
+#[spirv(matrix)]
+pub struct Mat4 {
+    x: Vec4,
+    y: Vec4,
+    z: Vec4,
+    w: Vec4,
+}
+
 #[spirv(fragment)]
 pub fn load(
     #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] buf: &mut [u32],
     out: &mut Vec4,
+    outmat: &mut Mat4,
 ) {
     unsafe {
         let buf = ByteAddressableBuffer::new(buf);
         *out = buf.load(5);
+        *outmat = buf.load(5);
     }
 }
 
 #[spirv(fragment)]
-pub fn store(#[spirv(descriptor_set = 0, binding = 0, storage_buffer)] buf: &mut [u32], val: Vec4) {
+pub fn store(
+    #[spirv(descriptor_set = 0, binding = 0, storage_buffer)] buf: &mut [u32],
+    val: Vec4,
+    valmat: Mat4,
+) {
     unsafe {
-        let buf = ByteAddressableBuffer::new(buf);
+        let mut buf = ByteAddressableBuffer::new(buf);
         buf.store(5, val);
+        buf.store(5, valmat);
     }
 }