nonuniform: ByteAddressableBuffer decorates NonUniform if ShaderNonUniform capability is enabled

Author: Firestar99

crates/rustc_codegen_spirv/src/builder/byte_addressable_buffer.rs

@@ -4,7 +4,8 @@ use crate::maybe_pqp_cg_ssa as rustc_codegen_ssa;
 use super::Builder;
 use crate::builder_spirv::{SpirvValue, SpirvValueExt, SpirvValueKind};
 use crate::spirv_type::SpirvType;
-use rspirv::spirv::Word;
+use rspirv::spirv::{Decoration, Word};
+use rustc_codegen_spirv_types::Capability;
 use rustc_codegen_ssa::traits::BuilderMethods;
 use rustc_errors::ErrorGuaranteed;
 use rustc_span::DUMMY_SP;
@@ -41,11 +42,20 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         };
         let u32_ty = SpirvType::Integer(32, false).def(DUMMY_SP, self);
         let u32_ptr = self.type_ptr_to(u32_ty);
+        let array = array.def(self);
+        let actual_index = actual_index.def(self);
         let ptr = self
             .emit()
-            .in_bounds_access_chain(u32_ptr, None, array.def(self), [actual_index.def(self)])
+            .in_bounds_access_chain(u32_ptr, None, array, [actual_index])
             .unwrap()
             .with_type(u32_ptr);
+        if self.builder.has_capability(Capability::ShaderNonUniform) {
+            // apply NonUniform to the operation and the index
+            self.emit()
+                .decorate(ptr.def(self), Decoration::NonUniform, []);
+            self.emit()
+                .decorate(actual_index, Decoration::NonUniform, []);
+        }
         self.load(u32_ty, ptr, Align::ONE)
     }
 
@@ -233,11 +243,20 @@ impl<'a, 'tcx> Builder<'a, 'tcx> {
         };
         let u32_ty = SpirvType::Integer(32, false).def(DUMMY_SP, self);
         let u32_ptr = self.type_ptr_to(u32_ty);
+        let array = array.def(self);
+        let actual_index = actual_index.def(self);
         let ptr = self
             .emit()
-            .in_bounds_access_chain(u32_ptr, None, array.def(self), [actual_index.def(self)])
+            .in_bounds_access_chain(u32_ptr, None, array, [actual_index])
             .unwrap()
             .with_type(u32_ptr);
+        if self.builder.has_capability(Capability::ShaderNonUniform) {
+            // apply NonUniform to the operation and the index
+            self.emit()
+                .decorate(ptr.def(self), Decoration::NonUniform, []);
+            self.emit()
+                .decorate(actual_index, Decoration::NonUniform, []);
+        }
         self.store(value, ptr, Align::ONE);
         Ok(())
     }