[naga]: Let TypeInner::Matrix hold a Scalar, not just a width.

Let `naga::TypeInner::Matrix` hold a full `Scalar`, with a kind and byte width, not merely a byte width, to make it possible to represent matrices of AbstractFloats for WGSL.
gfx-rs · Nov 21, 2023 · 8b3518b · 8b3518b
1 parent 42058cf
commit 8b3518b
Show file tree

Hide file tree

Showing 37 changed files with 225 additions and 246 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -95,6 +95,8 @@ Passing an owned value `window` to `Surface` will return a `Surface<'static>`. S
 
 - When reading GLSL, fix the argument types of the double-precision floating-point overloads of the `dot`, `reflect`, `distance`, and `ldexp` builtin functions. Correct the WGSL generated for constructing 64-bit floating-point matrices. Add tests for all the above. By @jimblandy in [#4684](https://github.com/gfx-rs/wgpu/pull/4684).
 
+- Allow Naga's IR types to represent matrices with elements elements of any scalar kind. This makes it possible for Naga IR types to represent WGSL abstract matrices. By @jimblandy in [#4735](https://github.com/gfx-rs/wgpu/pull/4735).
+
 - When evaluating const-expressions and generating SPIR-V, properly handle `Compose` expressions whose operands are `Splat` expressions. Such expressions are created and marked as constant by the constant evaluator. By @jimblandy in [#4695](https://github.com/gfx-rs/wgpu/pull/4695).
 
 - Preserve the source spans for constants and expressions correctly across module compaction. By @jimblandy in [#4696](https://github.com/gfx-rs/wgpu/pull/4696).
@@ -2353,4 +2355,4 @@ DeviceDescriptor {
 - concept of the storage hub
 - basic recording of passes and command buffers
 - submission-based lifetime tracking and command buffer recycling
-- automatic resource transitions
+- automatic resource transitions
diff --git a/naga/src/back/glsl/features.rs b/naga/src/back/glsl/features.rs
@@ -275,11 +275,9 @@ impl<'a, W> Writer<'a, W> {
 
         for (ty_handle, ty) in self.module.types.iter() {
             match ty.inner {
-                TypeInner::Scalar(scalar) => self.scalar_required_features(scalar),
-                TypeInner::Vector { scalar, .. } => self.scalar_required_features(scalar),
-                TypeInner::Matrix { width, .. } => {
-                    self.scalar_required_features(Scalar::float(width))
-                }
+                TypeInner::Scalar(scalar)
+                | TypeInner::Vector { scalar, .. }
+                | TypeInner::Matrix { scalar, .. } => self.scalar_required_features(scalar),
                 TypeInner::Array { base, size, .. } => {
                     if let TypeInner::Array { .. } = self.module.types[base].inner {
                         self.features.request(Features::ARRAY_OF_ARRAYS)

diff --git a/naga/src/back/glsl/mod.rs b/naga/src/back/glsl/mod.rs
@@ -985,11 +985,11 @@ impl<'a, W: Write> Writer<'a, W> {
             TypeInner::Matrix {
                 columns,
                 rows,
-                width,
+                scalar,
             } => write!(
                 self.out,
                 "{}mat{}x{}",
-                glsl_scalar(crate::Scalar::float(width))?.prefix,
+                glsl_scalar(scalar)?.prefix,
                 columns as u8,
                 rows as u8
             )?,

diff --git a/naga/src/back/hlsl/conv.rs b/naga/src/back/hlsl/conv.rs
@@ -47,10 +47,10 @@ impl crate::TypeInner {
             Self::Matrix {
                 columns,
                 rows,
-                width,
+                scalar,
             } => {
-                let stride = Alignment::from(rows) * width as u32;
-                let last_row_size = rows as u32 * width as u32;
+                let stride = Alignment::from(rows) * scalar.width as u32;
+                let last_row_size = rows as u32 * scalar.width as u32;
                 ((columns as u32 - 1) * stride) + last_row_size
             }
             Self::Array { base, size, stride } => {
@@ -82,10 +82,10 @@ impl crate::TypeInner {
             crate::TypeInner::Matrix {
                 columns,
                 rows,
-                width,
+                scalar,
             } => Cow::Owned(format!(
                 "{}{}x{}",
-                crate::Scalar::float(width).to_hlsl_str()?,
+                scalar.to_hlsl_str()?,
                 crate::back::vector_size_str(columns),
                 crate::back::vector_size_str(rows),
             )),

diff --git a/naga/src/back/hlsl/help.rs b/naga/src/back/hlsl/help.rs
@@ -656,10 +656,9 @@ impl<'a, W: Write> super::Writer<'a, W> {
             _ => unreachable!(),
         };
         let vec_ty = match module.types[member.ty].inner {
-            crate::TypeInner::Matrix { rows, width, .. } => crate::TypeInner::Vector {
-                size: rows,
-                scalar: crate::Scalar::float(width),
-            },
+            crate::TypeInner::Matrix { rows, scalar, .. } => {
+                crate::TypeInner::Vector { size: rows, scalar }
+            }
             _ => unreachable!(),
         };
         self.write_value_type(module, &vec_ty)?;
@@ -736,9 +735,7 @@ impl<'a, W: Write> super::Writer<'a, W> {
             _ => unreachable!(),
         };
         let scalar_ty = match module.types[member.ty].inner {
-            crate::TypeInner::Matrix { width, .. } => {
-                crate::TypeInner::Scalar(crate::Scalar::float(width))
-            }
+            crate::TypeInner::Matrix { scalar, .. } => crate::TypeInner::Scalar(scalar),
             _ => unreachable!(),
         };
         self.write_value_type(module, &scalar_ty)?;

diff --git a/naga/src/back/hlsl/storage.rs b/naga/src/back/hlsl/storage.rs
@@ -180,23 +180,20 @@ impl<W: fmt::Write> super::Writer<'_, W> {
             crate::TypeInner::Matrix {
                 columns,
                 rows,
-                width,
+                scalar,
             } => {
                 write!(
                     self.out,
                     "{}{}x{}(",
-                    crate::Scalar::float(width).to_hlsl_str()?,
+                    scalar.to_hlsl_str()?,
                     columns as u8,
                     rows as u8,
                 )?;
 
                 // Note: Matrices containing vec3s, due to padding, act like they contain vec4s.
-                let row_stride = Alignment::from(rows) * width as u32;
+                let row_stride = Alignment::from(rows) * scalar.width as u32;
                 let iter = (0..columns as u32).map(|i| {
-                    let ty_inner = crate::TypeInner::Vector {
-                        size: rows,
-                        scalar: crate::Scalar::float(width),
-                    };
+                    let ty_inner = crate::TypeInner::Vector { size: rows, scalar };
                     (TypeResolution::Value(ty_inner), i * row_stride)
                 });
                 self.write_storage_load_sequence(module, var_handle, iter, func_ctx)?;
@@ -316,7 +313,7 @@ impl<W: fmt::Write> super::Writer<'_, W> {
             crate::TypeInner::Matrix {
                 columns,
                 rows,
-                width,
+                scalar,
             } => {
                 // first, assign the value to a temporary
                 writeln!(self.out, "{level}{{")?;
@@ -325,7 +322,7 @@ impl<W: fmt::Write> super::Writer<'_, W> {
                     self.out,
                     "{}{}{}x{} {}{} = ",
                     level.next(),
-                    crate::Scalar::float(width).to_hlsl_str()?,
+                    scalar.to_hlsl_str()?,
                     columns as u8,
                     rows as u8,
                     STORE_TEMP_NAME,
@@ -335,16 +332,13 @@ impl<W: fmt::Write> super::Writer<'_, W> {
                 writeln!(self.out, ";")?;
 
                 // Note: Matrices containing vec3s, due to padding, act like they contain vec4s.
-                let row_stride = Alignment::from(rows) * width as u32;
+                let row_stride = Alignment::from(rows) * scalar.width as u32;
 
                 // then iterate the stores
                 for i in 0..columns as u32 {
                     self.temp_access_chain
                         .push(SubAccess::Offset(i * row_stride));
-                    let ty_inner = crate::TypeInner::Vector {
-                        size: rows,
-                        scalar: crate::Scalar::float(width),
-                    };
+                    let ty_inner = crate::TypeInner::Vector { size: rows, scalar };
                     let sv = StoreValue::TempIndex {
                         depth,
                         index: i,
@@ -467,10 +461,10 @@ impl<W: fmt::Write> super::Writer<'_, W> {
                     crate::TypeInner::Vector { scalar, .. } => Parent::Array {
                         stride: scalar.width as u32,
                     },
-                    crate::TypeInner::Matrix { rows, width, .. } => Parent::Array {
+                    crate::TypeInner::Matrix { rows, scalar, .. } => Parent::Array {
                         // The stride between matrices is the count of rows as this is how
                         // long each column is.
-                        stride: Alignment::from(rows) * width as u32,
+                        stride: Alignment::from(rows) * scalar.width as u32,
                     },
                     _ => unreachable!(),
                 },

diff --git a/naga/src/back/hlsl/writer.rs b/naga/src/back/hlsl/writer.rs
@@ -908,12 +908,9 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 TypeInner::Matrix {
                     rows,
                     columns,
-                    width,
+                    scalar,
                 } if member.binding.is_none() && rows == crate::VectorSize::Bi => {
-                    let vec_ty = crate::TypeInner::Vector {
-                        size: rows,
-                        scalar: crate::Scalar::float(width),
-                    };
+                    let vec_ty = crate::TypeInner::Vector { size: rows, scalar };
                     let field_name_key = NameKey::StructMember(handle, index as u32);
 
                     for i in 0..columns as u8 {
@@ -1037,7 +1034,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
             TypeInner::Matrix {
                 columns,
                 rows,
-                width,
+                scalar,
             } => {
                 // The IR supports only float matrix
                 // https://docs.microsoft.com/en-us/windows/win32/direct3dhlsl/dx-graphics-hlsl-matrix
@@ -1046,7 +1043,7 @@ impl<'a, W: fmt::Write> super::Writer<'a, W> {
                 write!(
                     self.out,
                     "{}{}x{}",
-                    crate::Scalar::float(width).to_hlsl_str()?,
+                    scalar.to_hlsl_str()?,
                     back::vector_size_str(columns),
                     back::vector_size_str(rows),
                 )?;
@@ -3241,11 +3238,11 @@ pub(super) fn get_inner_matrix_data(
         TypeInner::Matrix {
             columns,
             rows,
-            width,
+            scalar,
         } => Some(MatrixType {
             columns,
             rows,
-            width,
+            width: scalar.width,
         }),
         TypeInner::Array { base, .. } => get_inner_matrix_data(module, base),
         _ => None,
@@ -3276,12 +3273,12 @@ pub(super) fn get_inner_matrix_of_struct_array_member(
             TypeInner::Matrix {
                 columns,
                 rows,
-                width,
+                scalar,
             } => {
                 mat_data = Some(MatrixType {
                     columns,
                     rows,
-                    width,
+                    width: scalar.width,
                 })
             }
             TypeInner::Array { base, .. } => {
@@ -3333,12 +3330,12 @@ fn get_inner_matrix_of_global_uniform(
             TypeInner::Matrix {
                 columns,
                 rows,
-                width,
+                scalar,
             } => {
                 mat_data = Some(MatrixType {
                     columns,
                     rows,
-                    width,
+                    width: scalar.width,
                 })
             }
             TypeInner::Array { base, .. } => {

diff --git a/naga/src/back/msl/writer.rs b/naga/src/back/msl/writer.rs
@@ -1942,11 +1942,11 @@ impl<W: Write> Writer<W> {
                 crate::TypeInner::Matrix {
                     columns,
                     rows,
-                    width,
+                    scalar,
                 } => {
                     let target_scalar = crate::Scalar {
                         kind,
-                        width: convert.unwrap_or(width),
+                        width: convert.unwrap_or(scalar.width),
                     };
                     put_numeric_type(&mut self.out, target_scalar, &[rows, columns])?;
                     write!(self.out, "(")?;
@@ -2555,10 +2555,9 @@ impl<W: Write> Writer<W> {
             TypeResolution::Value(crate::TypeInner::Matrix {
                 columns,
                 rows,
-                width,
+                scalar,
             }) => {
-                let element = crate::Scalar::float(width);
-                put_numeric_type(&mut self.out, element, &[rows, columns])?;
+                put_numeric_type(&mut self.out, scalar, &[rows, columns])?;
             }
             TypeResolution::Value(ref other) => {
                 log::warn!("Type {:?} isn't a known local", other); //TEMP!

diff --git a/naga/src/back/spv/block.rs b/naga/src/back/spv/block.rs
@@ -494,7 +494,7 @@ impl<'w> BlockContext<'w> {
                         crate::TypeInner::Matrix {
                             columns,
                             rows,
-                            width,
+                            scalar,
                         } => {
                             self.write_matrix_matrix_column_op(
                                 block,
@@ -504,7 +504,7 @@ impl<'w> BlockContext<'w> {
                                 right_id,
                                 columns,
                                 rows,
-                                width,
+                                scalar.width,
                                 spirv::Op::FAdd,
                             );
 
@@ -522,7 +522,7 @@ impl<'w> BlockContext<'w> {
                         crate::TypeInner::Matrix {
                             columns,
                             rows,
-                            width,
+                            scalar,
                         } => {
                             self.write_matrix_matrix_column_op(
                                 block,
@@ -532,7 +532,7 @@ impl<'w> BlockContext<'w> {
                                 right_id,
                                 columns,
                                 rows,
-                                width,
+                                scalar.width,
                                 spirv::Op::FSub,
                             );
 
@@ -1141,9 +1141,7 @@ impl<'w> BlockContext<'w> {
                     match *self.fun_info[expr].ty.inner_with(&self.ir_module.types) {
                         crate::TypeInner::Scalar(scalar) => (scalar, None, false),
                         crate::TypeInner::Vector { scalar, size } => (scalar, Some(size), false),
-                        crate::TypeInner::Matrix { width, .. } => {
-                            (crate::Scalar::float(width), None, true)
-                        }
+                        crate::TypeInner::Matrix { scalar, .. } => (scalar, None, true),
                         ref other => {
                             log::error!("As source {:?}", other);
                             return Err(Error::Validation("Unexpected Expression::As source"));

diff --git a/naga/src/back/spv/mod.rs b/naga/src/back/spv/mod.rs
@@ -367,11 +367,11 @@ fn make_local(inner: &crate::TypeInner) -> Option<LocalType> {
         crate::TypeInner::Matrix {
             columns,
             rows,
-            width,
+            scalar,
         } => LocalType::Matrix {
             columns,
             rows,
-            width,
+            width: scalar.width,
         },
         crate::TypeInner::Pointer { base, space } => LocalType::Pointer {
             base,

diff --git a/naga/src/back/spv/writer.rs b/naga/src/back/spv/writer.rs
@@ -1766,10 +1766,10 @@ impl Writer {
         if let crate::TypeInner::Matrix {
             columns: _,
             rows,
-            width,
+            scalar,
         } = *member_array_subty_inner
         {
-            let byte_stride = Alignment::from(rows) * width as u32;
+            let byte_stride = Alignment::from(rows) * scalar.width as u32;
             self.annotations.push(Instruction::member_decorate(
                 struct_id,
                 index as u32,

diff --git a/naga/src/back/wgsl/writer.rs b/naga/src/back/wgsl/writer.rs
@@ -524,14 +524,14 @@ impl<W: Write> Writer<W> {
             TypeInner::Matrix {
                 columns,
                 rows,
-                width,
+                scalar,
             } => {
                 write!(
                     self.out,
                     "mat{}x{}<{}>",
                     back::vector_size_str(columns),
                     back::vector_size_str(rows),
-                    scalar_kind_str(crate::Scalar::float(width))
+                    scalar_kind_str(scalar)
                 )?;
             }
             TypeInner::Pointer { base, space } => {
@@ -1412,12 +1412,11 @@ impl<W: Write> Writer<W> {
                     TypeInner::Matrix {
                         columns,
                         rows,
-                        width,
-                        ..
+                        scalar,
                     } => {
                         let scalar = crate::Scalar {
                             kind,
-                            width: convert.unwrap_or(width),
+                            width: convert.unwrap_or(scalar.width),
                         };
                         let scalar_kind_str = scalar_kind_str(scalar);
                         write!(