Potentially optimize dot4{I,U}8Packed on Metal

This might allow the Metal compiler to emit faster code, but that's not
confirmed. See
<gpuweb/gpuweb#2677 (comment)>.

Author: robamler

naga/src/back/msl/writer.rs

@@ -121,6 +121,9 @@ const fn scalar_is_int(scalar: crate::Scalar) -> bool {
 /// Prefix for cached clamped level-of-detail values for `ImageLoad` expressions.
 const CLAMPED_LOD_LOAD_PREFIX: &str = "clamped_lod_e";
 
+/// Prefix for reinterpreted expressions using `as_type<T>(...)`.
+const REINTERPRET_PREFIX: &str = "reinterpreted_";
+
 /// Wrapper for identifier names for clamped level-of-detail values
 ///
 /// Values of this type implement [`core::fmt::Display`], formatting as
@@ -156,6 +159,30 @@ impl Display for ArraySizeMember {
     }
 }
 
+/// Wrapper for reinterpreted variables using `as_type<target_type>(orig)`.
+///
+/// Implements [`core::fmt::Display`], formatting as a name derived from
+/// `target_type` and the variable name of `orig`.
+#[derive(Clone, Copy)]
+struct Reinterpreted<'a> {
+    target_type: &'a str,
+    orig: Handle<crate::Expression>,
+}
+
+impl<'a> Reinterpreted<'a> {
+    const fn new(target_type: &'a str, orig: Handle<crate::Expression>) -> Self {
+        Self { target_type, orig }
+    }
+}
+
+impl Display for Reinterpreted<'_> {
+    fn fmt(&self, f: &mut Formatter<'_>) -> core::fmt::Result {
+        f.write_str(REINTERPRET_PREFIX)?;
+        f.write_str(self.target_type)?;
+        self.orig.write_prefixed(f, "_e")
+    }
+}
+
 struct TypeContext<'a> {
     handle: Handle<crate::Type>,
     gctx: proc::GlobalCtx<'a>,
@@ -1470,14 +1497,14 @@ impl<W: Write> Writer<W> {
 
     /// Emit code for the arithmetic expression of the dot product.
     ///
-    /// The argument `extractor` is a function that accepts a `Writer`, a handle to a vector,
-    /// and an index. writes out the expression for the component at that index.
-    fn put_dot_product(
+    /// The argument `extractor` is a function that accepts a `Writer`, a vector, and
+    /// an index. It writes out the expression for the vector component at that index.
+    fn put_dot_product<T: Copy>(
         &mut self,
-        arg: Handle<crate::Expression>,
-        arg1: Handle<crate::Expression>,
+        arg: T,
+        arg1: T,
         size: usize,
-        extractor: impl Fn(&mut Self, Handle<crate::Expression>, usize) -> BackendResult,
+        extractor: impl Fn(&mut Self, T, usize) -> BackendResult,
     ) -> BackendResult {
         // Write parentheses around the dot product expression to prevent operators
         // with different precedences from applying earlier.
@@ -2206,24 +2233,22 @@ impl<W: Write> Writer<W> {
                         ),
                     },
                     fun @ (Mf::Dot4I8Packed | Mf::Dot4U8Packed) => {
-                        let conversion = match fun {
-                            Mf::Dot4I8Packed => "int",
-                            Mf::Dot4U8Packed => "",
+                        // The two function arguments were already reinterpreted as packed (signed
+                        // or unsigned) chars in `Self::put_block`.
+                        let packed_type = match fun {
+                            Mf::Dot4I8Packed => "packed_char4",
+                            Mf::Dot4U8Packed => "packed_uchar4",
                             _ => unreachable!(),
                         };
 
                         return self.put_dot_product(
-                            arg,
-                            arg1.unwrap(),
+                            Reinterpreted::new(packed_type, arg),
+                            Reinterpreted::new(packed_type, arg1.unwrap()),
                             4,
                             |writer, arg, index| {
-                                write!(writer.out, "({}(", conversion)?;
-                                writer.put_expression(arg, context, true)?;
-                                if index == 3 {
-                                    write!(writer.out, ") >> 24)")?;
-                                } else {
-                                    write!(writer.out, ") << {} >> 24)", (3 - index) * 8)?;
-                                }
+                                // MSL implicitly promotes these (signed or unsigned) chars to
+                                // `int` or `uint` in the multiplication, so no overflow can occur.
+                                write!(writer.out, "{arg}[{index}]")?;
                                 Ok(())
                             },
                         );
@@ -3362,17 +3387,57 @@ impl<W: Write> Writer<W> {
             match *statement {
                 crate::Statement::Emit(ref range) => {
                     for handle in range.clone() {
-                        // `ImageLoad` expressions covered by the `Restrict` bounds check policy
-                        // may need to cache a clamped version of their level-of-detail argument.
-                        if let crate::Expression::ImageLoad {
-                            image,
-                            level: mip_level,
-                            ..
-                        } = context.expression.function.expressions[handle]
-                        {
-                            self.put_cache_restricted_level(
-                                handle, image, mip_level, level, context,
-                            )?;
+                        use crate::MathFunction as Mf;
+
+                        match context.expression.function.expressions[handle] {
+                            // `ImageLoad` expressions covered by the `Restrict` bounds check policy
+                            // may need to cache a clamped version of their level-of-detail argument.
+                            crate::Expression::ImageLoad {
+                                image,
+                                level: mip_level,
+                                ..
+                            } => {
+                                self.put_cache_restricted_level(
+                                    handle, image, mip_level, level, context,
+                                )?;
+                            }
+
+                            // If we are going to write a `Dot4I8Packed` or `Dot4U8Packed` then we
+                            // introduce two intermediate variables that recast the two arguments as
+                            // packed (signed or unsigned) chars. The actual dot product is
+                            // implemented in `Self::put_expression`, and it uses both of these
+                            // intermediate variables multiple times.
+                            crate::Expression::Math {
+                                fun: fun @ (Mf::Dot4I8Packed | Mf::Dot4U8Packed),
+                                arg,
+                                arg1,
+                                ..
+                            } => {
+                                let arg1 = arg1.unwrap();
+                                let packed_type = match fun {
+                                    Mf::Dot4I8Packed => "packed_char4",
+                                    Mf::Dot4U8Packed => "packed_uchar4",
+                                    _ => unreachable!(),
+                                };
+
+                                write!(
+                                    self.out,
+                                    "{level}{packed_type} {0} = as_type<{packed_type}>(",
+                                    Reinterpreted::new(packed_type, arg)
+                                )?;
+                                self.put_expression(arg, &context.expression, true)?;
+                                writeln!(self.out, ");")?;
+
+                                write!(
+                                    self.out,
+                                    "{level}{packed_type} {0} = as_type<{packed_type}>(",
+                                    Reinterpreted::new(packed_type, arg1)
+                                )?;
+                                self.put_expression(arg1, &context.expression, true)?;
+                                writeln!(self.out, ");")?;
+                            }
+
+                            _ => (),
                         }
 
                         let ptr_class = context.expression.resolve_type(handle).pointer_space();

naga/tests/out/msl/wgsl-functions.msl

@@ -29,14 +29,22 @@ int test_integer_dot_product(
 
 uint test_packed_integer_dot_product(
 ) {
-    int c_5_ = ( + (int(1u) << 24 >> 24) * (int(2u) << 24 >> 24) + (int(1u) << 16 >> 24) * (int(2u) << 16 >> 24) + (int(1u) << 8 >> 24) * (int(2u) << 8 >> 24) + (int(1u) >> 24) * (int(2u) >> 24));
-    uint c_6_ = ( + ((3u) << 24 >> 24) * ((4u) << 24 >> 24) + ((3u) << 16 >> 24) * ((4u) << 16 >> 24) + ((3u) << 8 >> 24) * ((4u) << 8 >> 24) + ((3u) >> 24) * ((4u) >> 24));
+    packed_char4 reinterpreted_packed_char4_e0 = as_type<packed_char4>(1u);
+    packed_char4 reinterpreted_packed_char4_e1 = as_type<packed_char4>(2u);
+    int c_5_ = ( + reinterpreted_packed_char4_e0[0] * reinterpreted_packed_char4_e1[0] + reinterpreted_packed_char4_e0[1] * reinterpreted_packed_char4_e1[1] + reinterpreted_packed_char4_e0[2] * reinterpreted_packed_char4_e1[2] + reinterpreted_packed_char4_e0[3] * reinterpreted_packed_char4_e1[3]);
+    packed_uchar4 reinterpreted_packed_uchar4_e3 = as_type<packed_uchar4>(3u);
+    packed_uchar4 reinterpreted_packed_uchar4_e4 = as_type<packed_uchar4>(4u);
+    uint c_6_ = ( + reinterpreted_packed_uchar4_e3[0] * reinterpreted_packed_uchar4_e4[0] + reinterpreted_packed_uchar4_e3[1] * reinterpreted_packed_uchar4_e4[1] + reinterpreted_packed_uchar4_e3[2] * reinterpreted_packed_uchar4_e4[2] + reinterpreted_packed_uchar4_e3[3] * reinterpreted_packed_uchar4_e4[3]);
     uint _e7 = 5u + c_6_;
     uint _e9 = 6u + c_6_;
-    int c_7_ = ( + (int(_e7) << 24 >> 24) * (int(_e9) << 24 >> 24) + (int(_e7) << 16 >> 24) * (int(_e9) << 16 >> 24) + (int(_e7) << 8 >> 24) * (int(_e9) << 8 >> 24) + (int(_e7) >> 24) * (int(_e9) >> 24));
+    packed_char4 reinterpreted_packed_char4_e7 = as_type<packed_char4>(_e7);
+    packed_char4 reinterpreted_packed_char4_e9 = as_type<packed_char4>(_e9);
+    int c_7_ = ( + reinterpreted_packed_char4_e7[0] * reinterpreted_packed_char4_e9[0] + reinterpreted_packed_char4_e7[1] * reinterpreted_packed_char4_e9[1] + reinterpreted_packed_char4_e7[2] * reinterpreted_packed_char4_e9[2] + reinterpreted_packed_char4_e7[3] * reinterpreted_packed_char4_e9[3]);
     uint _e12 = 7u + c_6_;
     uint _e14 = 8u + c_6_;
-    uint c_8_ = ( + ((_e12) << 24 >> 24) * ((_e14) << 24 >> 24) + ((_e12) << 16 >> 24) * ((_e14) << 16 >> 24) + ((_e12) << 8 >> 24) * ((_e14) << 8 >> 24) + ((_e12) >> 24) * ((_e14) >> 24));
+    packed_uchar4 reinterpreted_packed_uchar4_e12 = as_type<packed_uchar4>(_e12);
+    packed_uchar4 reinterpreted_packed_uchar4_e14 = as_type<packed_uchar4>(_e14);
+    uint c_8_ = ( + reinterpreted_packed_uchar4_e12[0] * reinterpreted_packed_uchar4_e14[0] + reinterpreted_packed_uchar4_e12[1] * reinterpreted_packed_uchar4_e14[1] + reinterpreted_packed_uchar4_e12[2] * reinterpreted_packed_uchar4_e14[2] + reinterpreted_packed_uchar4_e12[3] * reinterpreted_packed_uchar4_e14[3]);
     return c_8_;
 }