[hlsl-out] More matCx2 fixes (#1989)

* [hlsl-out] fix matCx2 as global uniform

* [hlsl-out] update comments

* [hlsl-out] fix `row_major` not being written on global arrays of matrices and also write it on nested arrays of matrices

* [hlsl-out] fix matCx2's nested inside global arrays

* [hlsl-out] fix struct members of type array<matCx2>

* [hlsl-out] test mat2x4 to make sure our matCx2 code behaves properly

src/back/hlsl/help.rs

@@ -54,6 +54,11 @@ pub(super) struct WrappedStructMatrixAccess {
     pub(super) index: u32,
 }
 
+#[derive(Clone, Copy, Debug, Hash, Eq, Ord, PartialEq, PartialOrd)]
+pub(super) struct WrappedMatCx2 {
+    pub(super) columns: crate::VectorSize,
+}
+
 /// HLSL backend requires its own `ImageQuery` enum.
 ///
 /// It is used inside `WrappedImageQuery` and should be unique per ImageQuery function.
@@ -461,12 +466,36 @@ impl<'a, W: Write> super::Writer<'a, W> {
                                 )?;
                             }
                         }
-                        _ => {
-                            writeln!(
-                                self.out,
-                                "{}{}.{} = {}{};",
-                                INDENT, RETURN_VARIABLE_NAME, field_name, ARGUMENT_VARIABLE_NAME, i,
-                            )?;
+                        ref other => {
+                            // We cast arrays of native HLSL `floatCx2`s to arrays of `matCx2`s
+                            // (where the inner matrix is represented by a struct with C `float2` members).
+                            // See the module-level block comment in mod.rs for details.
+                            if let Some(super::writer::MatrixType {
+                                columns,
+                                rows: crate::VectorSize::Bi,
+                                width: 4,
+                            }) = super::writer::get_inner_matrix_data(module, member.ty)
+                            {
+                                write!(
+                                    self.out,
+                                    "{}{}.{} = (__mat{}x2",
+                                    INDENT, RETURN_VARIABLE_NAME, field_name, columns as u8
+                                )?;
+                                if let crate::TypeInner::Array { base, size, .. } = *other {
+                                    self.write_array_size(module, base, size)?;
+                                }
+                                writeln!(self.out, "){}{};", ARGUMENT_VARIABLE_NAME, i,)?;
+                            } else {
+                                writeln!(
+                                    self.out,
+                                    "{}{}.{} = {}{};",
+                                    INDENT,
+                                    RETURN_VARIABLE_NAME,
+                                    field_name,
+                                    ARGUMENT_VARIABLE_NAME,
+                                    i,
+                                )?;
+                            }
                         }
                     }
                 }
@@ -1050,4 +1079,117 @@ impl<'a, W: Write> super::Writer<'a, W> {
         }
         Ok(())
     }
+
+    pub(super) fn write_mat_cx2_typedef_and_functions(
+        &mut self,
+        WrappedMatCx2 { columns }: WrappedMatCx2,
+    ) -> BackendResult {
+        use crate::back::INDENT;
+
+        // typedef
+        write!(self.out, "typedef struct {{ ")?;
+        for i in 0..columns as u8 {
+            write!(self.out, "float2 _{}; ", i)?;
+        }
+        writeln!(self.out, "}} __mat{}x2;", columns as u8)?;
+
+        // __get_col_of_mat
+        writeln!(
+            self.out,
+            "float2 __get_col_of_mat{}x2(__mat{}x2 mat, uint idx) {{",
+            columns as u8, columns as u8
+        )?;
+        writeln!(self.out, "{}switch(idx) {{", INDENT)?;
+        for i in 0..columns as u8 {
+            writeln!(self.out, "{}case {}: {{ return mat._{}; }}", INDENT, i, i)?;
+        }
+        writeln!(self.out, "{}default: {{ return (float2)0; }}", INDENT)?;
+        writeln!(self.out, "{}}}", INDENT)?;
+        writeln!(self.out, "}}")?;
+
+        // __set_col_of_mat
+        writeln!(
+            self.out,
+            "void __set_col_of_mat{}x2(__mat{}x2 mat, uint idx, float2 value) {{",
+            columns as u8, columns as u8
+        )?;
+        writeln!(self.out, "{}switch(idx) {{", INDENT)?;
+        for i in 0..columns as u8 {
+            writeln!(
+                self.out,
+                "{}case {}: {{ mat._{} = value; break; }}",
+                INDENT, i, i
+            )?;
+        }
+        writeln!(self.out, "{}}}", INDENT)?;
+        writeln!(self.out, "}}")?;
+
+        // __set_el_of_mat
+        writeln!(
+            self.out,
+            "void __set_el_of_mat{}x2(__mat{}x2 mat, uint idx, uint vec_idx, float value) {{",
+            columns as u8, columns as u8
+        )?;
+        writeln!(self.out, "{}switch(idx) {{", INDENT)?;
+        for i in 0..columns as u8 {
+            writeln!(
+                self.out,
+                "{}case {}: {{ mat._{}[vec_idx] = value; break; }}",
+                INDENT, i, i
+            )?;
+        }
+        writeln!(self.out, "{}}}", INDENT)?;
+        writeln!(self.out, "}}")?;
+
+        writeln!(self.out)?;
+
+        Ok(())
+    }
+
+    pub(super) fn write_all_mat_cx2_typedefs_and_functions(
+        &mut self,
+        module: &crate::Module,
+    ) -> BackendResult {
+        for (handle, _) in module.global_variables.iter() {
+            let global = &module.global_variables[handle];
+
+            if global.space == crate::AddressSpace::Uniform {
+                if let Some(super::writer::MatrixType {
+                    columns,
+                    rows: crate::VectorSize::Bi,
+                    width: 4,
+                }) = super::writer::get_inner_matrix_data(module, global.ty)
+                {
+                    let entry = WrappedMatCx2 { columns };
+                    if !self.wrapped.mat_cx2s.contains(&entry) {
+                        self.write_mat_cx2_typedef_and_functions(entry)?;
+                        self.wrapped.mat_cx2s.insert(entry);
+                    }
+                }
+            }
+        }
+
+        for (_, ty) in module.types.iter() {
+            if let crate::TypeInner::Struct { ref members, .. } = ty.inner {
+                for member in members.iter() {
+                    if let crate::TypeInner::Array { .. } = module.types[member.ty].inner {
+                        if let Some(super::writer::MatrixType {
+                            columns,
+                            rows: crate::VectorSize::Bi,
+                            width: 4,
+                        }) = super::writer::get_inner_matrix_data(module, member.ty)
+                        {
+                            let entry = WrappedMatCx2 { columns };
+                            if !self.wrapped.mat_cx2s.contains(&entry) {
+                                self.write_mat_cx2_typedef_and_functions(entry)?;
+                                self.wrapped.mat_cx2s.insert(entry);
+                            }
+                        }
+                    }
+                }
+            }
+        }
+
+        Ok(())
+    }
 }

src/back/hlsl/mod.rs

@@ -57,8 +57,8 @@ that the columns of a `matKx2<f32>` need only be [aligned as required
 for `vec2<f32>`][ilov], which is [eight-byte alignment][8bb].
 
 To compensate for this, any time a `matKx2<f32>` appears in a WGSL
-`uniform` variable, whether directly as the variable's type or as a
-struct member, we actually emit `K` separate `float2` members, and
+`uniform` variable, whether directly as the variable's type or as part
+of a struct/array, we actually emit `K` separate `float2` members, and
 assemble/disassemble the matrix from its columns (in WGSL; rows in
 HLSL) upon load and store.
 
@@ -92,14 +92,10 @@ float3x2 GetMatmOnBaz(Baz obj) {
 We also emit an analogous `Set` function, as well as functions for
 accessing individual columns by dynamic index.
 
-At present, we do not generate correct HLSL when `matCx2<f32>` us used
-directly as the type of a WGSL `uniform` global ([#1837]).
-
 [hlsl]: https://docs.microsoft.com/en-us/windows/win32/direct3dhlsl/dx-graphics-hlsl
 [ilov]: https://gpuweb.github.io/gpuweb/wgsl/#internal-value-layout
 [16bb]: https://github.com/microsoft/DirectXShaderCompiler/wiki/Buffer-Packing#constant-buffer-packing
 [8bb]: https://gpuweb.github.io/gpuweb/wgsl/#alignment-and-size
-[#1837]: https://github.com/gfx-rs/naga/issues/1837
 */
 
 mod conv;
@@ -253,6 +249,7 @@ struct Wrapped {
     image_queries: crate::FastHashSet<help::WrappedImageQuery>,
     constructors: crate::FastHashSet<help::WrappedConstructor>,
     struct_matrix_access: crate::FastHashSet<help::WrappedStructMatrixAccess>,
+    mat_cx2s: crate::FastHashSet<help::WrappedMatCx2>,
 }
 
 impl Wrapped {
@@ -261,6 +258,7 @@ impl Wrapped {
         self.image_queries.clear();
         self.constructors.clear();
         self.struct_matrix_access.clear();
+        self.mat_cx2s.clear();
     }
 }