[naga spv-out] Handle matCx2 in uniform buffers

naga/src/back/spv/helpers.rs

@@ -122,6 +122,44 @@ pub fn global_needs_wrapper(ir_module: &crate::Module, var: &crate::GlobalVariab
     }
 }
 
+/// Returns true if `pointer` refers to two-row matrix which is a member of a
+/// struct in the [`crate::AddressSpace::Uniform`] address space.
+pub fn is_uniform_matcx2_struct_member_access(
+    ir_function: &crate::Function,
+    fun_info: &crate::valid::FunctionInfo,
+    ir_module: &crate::Module,
+    pointer: Handle<crate::Expression>,
+) -> bool {
+    if let crate::TypeInner::Pointer {
+        base: pointer_base_type,
+        space: crate::AddressSpace::Uniform,
+    } = *fun_info[pointer].ty.inner_with(&ir_module.types)
+    {
+        if let crate::TypeInner::Matrix {
+            rows: crate::VectorSize::Bi,
+            ..
+        } = ir_module.types[pointer_base_type].inner
+        {
+            if let crate::Expression::AccessIndex {
+                base: parent_pointer,
+                ..
+            } = ir_function.expressions[pointer]
+            {
+                if let crate::TypeInner::Pointer {
+                    base: parent_type, ..
+                } = *fun_info[parent_pointer].ty.inner_with(&ir_module.types)
+                {
+                    if let crate::TypeInner::Struct { .. } = ir_module.types[parent_type].inner {
+                        return true;
+                    }
+                }
+            }
+        }
+    }
+
+    false
+}
+
 ///HACK: this is taken from std unstable, remove it when std's floor_char_boundary is stable
 trait U8Internal {
     fn is_utf8_char_boundary(&self) -> bool;

naga/src/back/spv/index.rs

@@ -536,17 +536,18 @@ impl BlockContext<'_> {
     /// Emit code to subscript a vector by value with a computed index.
     ///
     /// Return the id of the element value.
+    ///
+    /// If `base_id_override` is provided, it is used as the vector expression
+    /// to be subscripted into, rather than the cached value of `base`.
     pub(super) fn write_vector_access(
         &mut self,
-        expr_handle: Handle<crate::Expression>,
+        result_type_id: Word,
         base: Handle<crate::Expression>,
-        index: Handle<crate::Expression>,
+        base_id_override: Option<Word>,
+        index: GuardedIndex,
         block: &mut Block,
     ) -> Result<Word, Error> {
-        let result_type_id = self.get_expression_type_id(&self.fun_info[expr_handle].ty);
-
-        let base_id = self.cached[base];
-        let index = GuardedIndex::Expression(index);
+        let base_id = base_id_override.unwrap_or_else(|| self.cached[base]);
 
         let result_id = match self.write_bounds_check(base, index, block)? {
             BoundsCheckResult::KnownInBounds(known_index) => {

naga/src/back/spv/mod.rs

@@ -1,7 +1,91 @@
 /*!
 Backend for [SPIR-V][spv] (Standard Portable Intermediate Representation).
 
+# Layout of values in `uniform` buffers
+
+WGSL's ["Internal Layout of Values"][ilov] rules specify how each WGSL type
+should be stored in `uniform` and `storage` buffers, and Naga IR adheres to
+these rules. The SPIR-V we generate must access values in that form, even when
+it is not what Vulkan would use normally. Fortunately the rules for `storage`
+buffers match Vulkan's, but some adjustments must be made when emitting SPIR-V
+for `uniform` buffers.
+
+## Padding in two-row matrices
+
+In Vulkan's ["extended layout"][extended-layout] (also known as std140) used
+for `uniform` buffers, matrices are defined in terms of arrays of their vector
+type, and arrays are defined to have an alignment equal to the alignment of
+their element type rounded up to a multiple of 16. This means that each column
+of the vector has a minimum alignment of 16. WGSL, and consequently Naga IR, on
+the other hand defines each column to have an alignment equal to the alignment
+of the vector type, without being rounded up to 16.
+
+To compensate for this, for any `struct` used as a `uniform` buffer which
+contains a two-row matrix, we declare an additional "std140 compatible" type
+in which each column of the matrix has been decomposed into the containing
+struct. For example, the following WGSL struct type:
+
+```ignore
+struct Baz {
+    m: mat3x2<f32>,
+}
+```
+
+is rendered as the SPIR-V struct type:
+
+```ignore
+OpTypeStruct %v2float %v2float %v2float
+```
+
+This has the effect that struct indices in Naga IR for such types do not
+correspond to the struct indices used in SPIR-V. A mapping of struct indices
+for these types is maintained in [`Std140CompatTypeInfo`].
+
+Additionally, any two-row matrices that are declared directly as uniform
+buffers without being wrapped in a struct are declared as a struct containing a
+vector member for each column. Any array of a two-row matrix in a uniform
+buffer is declared as an array of a struct containing a vector member for each
+column. Any struct or array within a uniform buffer which contains a member or
+whose base type requires requires a std140 compatible type declaration, itself
+requires a std140 compatible type declaration.
+
+Whenever a value of such a type is [`loaded`] we insert code to convert the
+loaded value from the std140 compatible type to the regular type. This occurs
+in `BlockContext::write_checked_load`, making use of the wrapper function
+defined by `Writer::write_wrapped_convert_from_std140_compat_type`. For matrices
+that have been decomposed as separate columns in the containing struct, we load
+each column separately then composite the matrix type in
+`BlockContext::maybe_write_load_uniform_matcx2_struct_member`.
+
+Whenever a column of a matrix that has been decomposed into its containing
+struct is [`accessed`] with a constant index we adjust the emitted access chain
+to access from the containing struct instead, in `BlockContext::write_access_chain`.
+
+Whenever a column of a uniform buffer two-row matrix is [`dynamically accessed`]
+we must first load the matrix type, converting it from its std140 compatible
+type as described above, then access the column using the wrapper function
+defined by `Writer::write_wrapped_matcx2_get_column`. This is handled by
+`BlockContext::maybe_write_uniform_matcx2_dynamic_access`.
+
+Note that this approach differs somewhat from the equivalent code in the HLSL
+backend. For HLSL all structs containing two-row matrices (or arrays of such)
+have their declarations modified, not just those used as uniform buffers.
+Two-row matrices and arrays of such only use modified type declarations when
+used as uniform buffers, or additionally when used as struct member in any
+context. This avoids the need to convert struct values when loading from uniform
+buffers, but when loading arrays and matrices from uniform buffers or from any
+struct the conversion is still required. In contrast, the approach used here
+always requires converting *any* affected type when loading from a uniform
+buffer, but consistently *only* when loading from a uniform buffer. As a result
+this also means we only have to handle loads and not stores, as uniform buffers
+are read-only.
+
 [spv]: https://www.khronos.org/registry/SPIR-V/
+[ilov]: https://gpuweb.github.io/gpuweb/wgsl/#internal-value-layout
+[extended-layout]: https://docs.vulkan.org/spec/latest/chapters/interfaces.html#interfaces-resources-layout
+[`loaded`]: crate::Expression::Load
+[`accessed`]: crate::Expression::AccessIndex
+[`dynamically accessed`]: crate::Expression::Access
 */
 
 mod block;
@@ -461,6 +545,12 @@ enum WrappedFunction {
         left_type_id: Word,
         right_type_id: Word,
     },
+    ConvertFromStd140CompatType {
+        r#type: Handle<crate::Type>,
+    },
+    MatCx2GetColumn {
+        r#type: Handle<crate::Type>,
+    },
 }
 
 /// A map from evaluated [`Expression`](crate::Expression)s to their SPIR-V ids.
@@ -721,6 +811,20 @@ impl BlockContext<'_> {
     }
 }
 
+/// Information about a type for which we have declared a std140 layout
+/// compatible variant, because the type is used in a uniform but does not
+/// adhere to std140 requirements. The uniform will be declared using the
+/// type `type_id`, and the result of any `Load` will be immediately converted
+/// to the base type. This is used for matrices with 2 rows, as well as any
+/// arrays or structs containing such matrices.
+pub struct Std140CompatTypeInfo {
+    /// ID of the std140 compatible type declaration.
+    type_id: Word,
+    /// For structs, a mapping of Naga IR struct member indices to the indices
+    /// used in the generated SPIR-V. For non-struct types this will be empty.
+    member_indices: Vec<u32>,
+}
+
 pub struct Writer {
     physical_layout: PhysicalLayout,
     logical_layout: LogicalLayout,
@@ -760,6 +864,7 @@ pub struct Writer {
     constant_ids: HandleVec<crate::Expression, Word>,
     cached_constants: crate::FastHashMap<CachedConstant, Word>,
     global_variables: HandleVec<crate::GlobalVariable, GlobalVariable>,
+    std140_compat_uniform_types: crate::FastHashMap<Handle<crate::Type>, Std140CompatTypeInfo>,
     binding_map: BindingMap,
 
     // Cached expressions are only meaningful within a BlockContext, but we