spv-out: Support arrayLength of a dynamically indexed bindings array

naga/src/back/spv/index.rs

@@ -48,25 +48,39 @@ impl<'w> BlockContext<'w> {
         // inside a buffer that is itself an element in a buffer bindings array.
         // SPIR-V requires that runtime-sized arrays are wrapped in structs.
         // See `helpers::global_needs_wrapper` and its uses.
-        let (opt_array_index, global_handle, opt_last_member_index) = match self
+        let (opt_array_index_id, global_handle, opt_last_member_index) = match self
             .ir_function
             .expressions[array]
         {
-            // Note that SPIR-V forbids `OpArrayLength` on a variable pointer,
-            // so we aren't handling `crate::Expression::Access` here.
             crate::Expression::AccessIndex { base, index } => {
                 match self.ir_function.expressions[base] {
                     // The global variable is an array of buffer bindings of structs,
-                    // and we are accessing the last member.
+                    // we are accessing one of them with a static index,
+                    // and the last member of it.
                     crate::Expression::AccessIndex {
                         base: base_outer,
                         index: index_outer,
                     } => match self.ir_function.expressions[base_outer] {
                         crate::Expression::GlobalVariable(handle) => {
-                            (Some(index_outer), handle, Some(index))
+                            let index_id = self.get_index_constant(index_outer);
+                            (Some(index_id), handle, Some(index))
                         }
                         _ => return Err(Error::Validation("array length expression case-1a")),
                     },
+                    // The global variable is an array of buffer bindings of structs,
+                    // we are accessing one of them with a dynamic index,
+                    // and the last member of it.
+                    crate::Expression::Access {
+                        base: base_outer,
+                        index: index_outer,
+                    } => match self.ir_function.expressions[base_outer] {
+                        crate::Expression::GlobalVariable(handle) => {
+                            let index_id = self.cached[index_outer];
+                            (Some(index_id), handle, Some(index))
+                        }
+                        _ => return Err(Error::Validation("array length expression case-1b")),
+                    },
+                    // The global variable is a buffer, and we are accessing the last member.
                     crate::Expression::GlobalVariable(handle) => {
                         let global = &self.ir_module.global_variables[handle];
                         match self.ir_module.types[global.ty].inner {
@@ -79,15 +93,27 @@ impl<'w> BlockContext<'w> {
                     _ => return Err(Error::Validation("array length expression case-1c")),
                 }
             }
+            // The global variable is an array of buffer bindings of arrays.
+            crate::Expression::Access { base, index } => match self.ir_function.expressions[base] {
+                crate::Expression::GlobalVariable(handle) => {
+                    let index_id = self.cached[index];
+                    let global = &self.ir_module.global_variables[handle];
+                    match self.ir_module.types[global.ty].inner {
+                        crate::TypeInner::BindingArray { .. } => (Some(index_id), handle, None),
+                        _ => return Err(Error::Validation("array length expression case-2a")),
+                    }
+                }
+                _ => return Err(Error::Validation("array length expression case-2b")),
+            },
             // The global variable is a run-time array.
             crate::Expression::GlobalVariable(handle) => {
                 let global = &self.ir_module.global_variables[handle];
                 if !global_needs_wrapper(self.ir_module, global) {
-                    return Err(Error::Validation("array length expression case-2"));
+                    return Err(Error::Validation("array length expression case-3"));
                 }
                 (None, handle, None)
             }
-            _ => return Err(Error::Validation("array length expression case-3")),
+            _ => return Err(Error::Validation("array length expression case-4")),
         };
 
         let gvar = self.writer.global_variables[global_handle.index()].clone();
@@ -103,17 +129,16 @@ impl<'w> BlockContext<'w> {
                 (0, gvar.var_id)
             }
         };
-        let structure_id = match opt_array_index {
+        let structure_id = match opt_array_index_id {
             // We are indexing inside a binding array, generate the access op.
-            Some(index) => {
+            Some(index_id) => {
                 let element_type_id = match self.ir_module.types[global.ty].inner {
                     crate::TypeInner::BindingArray { base, size: _ } => {
                         let class = map_storage_class(global.space);
                         self.get_pointer_id(base, class)?
                     }
-                    _ => return Err(Error::Validation("array length expression case-4")),
+                    _ => return Err(Error::Validation("array length expression case-5")),
                 };
-                let index_id = self.get_index_constant(index);
                 let structure_id = self.gen_id();
                 block.body.push(Instruction::access_chain(
                     element_type_id,

naga/src/valid/type.rs

@@ -510,7 +510,6 @@ impl super::Validator {
                 ti.uniform_layout = Ok(Alignment::MIN_UNIFORM);
 
                 let mut min_offset = 0;
-
                 let mut prev_struct_data: Option<(u32, u32)> = None;
 
                 for (i, member) in members.iter().enumerate() {
@@ -662,6 +661,7 @@ impl super::Validator {
                     // Currently Naga only supports binding arrays of structs for non-handle types.
                     match gctx.types[base].inner {
                         crate::TypeInner::Struct { .. } => {}
+                        crate::TypeInner::Array { .. } => {}
                         _ => return Err(TypeError::BindingArrayBaseTypeNotStruct(base)),
                     };
                 }

naga/tests/in/binding-buffer-arrays.wgsl

@@ -24,6 +24,8 @@ fn main(fragment_in: FragmentIn) -> @location(0) u32 {
     u1 += storage_array[non_uniform_index].x;
 
     u1 += arrayLength(&storage_array[0].far);
+    u1 += arrayLength(&storage_array[uniform_index].far);
+    u1 += arrayLength(&storage_array[non_uniform_index].far);
 
     return u1;
 }

naga/tests/out/spv/binding-buffer-arrays.spvasm

@@ -1,7 +1,7 @@
 ; SPIR-V
 ; Version: 1.1
 ; Generator: rspirv
-; Bound: 68
+; Bound: 76
 OpCapability Shader
 OpCapability ShaderNonUniform
 OpExtension "SPV_KHR_storage_buffer_storage_class"
@@ -104,7 +104,17 @@ OpStore %31 %61
 %65 = OpLoad  %3  %31
 %66 = OpIAdd  %3  %65 %64
 OpStore %31 %66
-%67 = OpLoad  %3  %31
+%67 = OpAccessChain  %38  %11 %36
-OpStore %23 %67
+%68 = OpArrayLength  %3  %67 1
+%69 = OpLoad  %3  %31
+%70 = OpIAdd  %3  %69 %68
+OpStore %31 %70
+%71 = OpAccessChain  %38  %11 %37
+%72 = OpArrayLength  %3  %71 1
+%73 = OpLoad  %3  %31
+%74 = OpIAdd  %3  %73 %72
+OpStore %31 %74
+%75 = OpLoad  %3  %31
+OpStore %23 %75
 OpReturn
 OpFunctionEnd

naga/tests/out/wgsl/binding-buffer-arrays.wgsl

@@ -33,6 +33,10 @@ fn main(fragment_in: FragmentIn) -> @location(0) @interpolate(flat) u32 {
     u1_ = (_e23 + _e22);
     let _e29 = u1_;
     u1_ = (_e29 + arrayLength((&storage_array[0].far)));
-    let _e31 = u1_;
+    let _e35 = u1_;
-    return _e31;
+    u1_ = (_e35 + arrayLength((&storage_array[uniform_index].far)));
+    let _e41 = u1_;
+    u1_ = (_e41 + arrayLength((&storage_array[non_uniform_index].far)));
+    let _e43 = u1_;
+    return _e43;
 }