[naga spv-out] Handle 2-row matrices in uniform buffers

Two-row matrices in uniform buffers in Vulkan/SPIR-V's default
"extended" (std140) layout do not adhere to WGSL/Naga IR's layout
rules - each column is aligned to a minimum of 16 bytes rather than
just the vector size.

To work around this, we emit additional std140 compatible type
declarations for each type used in a uniform buffer that requires
one. Any two-row matrix struct member is decomposed into the
containing struct as separate vectors for each column. Two-row
matrices used directly as uniform buffers are wrapped in a struct with
a vector member for each column, and arrays (of arrays, etc) of
two-row matrices are declared as arrays (of arrays) of structs
containing a vector member for each column.

When loading such a value from a uniform buffer, we convert from the
std140 compatible type to the regular type immediately after loading.
Accesses of a uniform two-row matrix's columns using constant indices
are rewritten to access the containing struct's vector members
directly. Accesses using dynamic indices are implemented by loading
and converting the matrix, extracting each column then `switch`ing on
the index to select the correct column.

We can now remove the expected failure annotation for the Vulkan
backend on the uniform_input struct_layout test, as it now passes. We
additionally make the hlsl_mat_cx*.wgsl snapshot tests run for the
SPIR-V backend too, and remove the "hlsl" prefix from the name, as
they are no longer just relevant for HLSL.

Author: jamienicol

naga/src/back/spv/block.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/helpers.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/index.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