[naga] Implement constant evaluation for the cross builtin. (#7404)

Add support for `naga::ir::MathFunction::Cross` to
`naga::proc::constant_evaluator`.

In the tests:

- Change `naga/tests/in/wgsl/cross.wgsl` to use more interesting
  argument values. Rather than passing the same vector twice, which
  yields a cross product of zero, pass in the x and y unit vectors,
  whose cross product is the z unit vector. Update snapshot output.

- Replace `validation::bad_cross_builtin_args` with a new test,
  `builtin_cross_product_args`, that is written more in the style of
  the other tests in this module, and does not depend on the WGSL
  front end. Because this PR changes the behavior of the constant
  evaluator, this test stopped behaving as expected.

- In `wgsl_errors::check`, move a `panic!` out of a closure so that
  the `#[track_caller]` attribute works properly.

Author: jimblandy

naga/src/proc/constant_evaluator.rs

@@ -474,7 +474,6 @@ impl ExpressionKindTracker {
                         fun,
                         Mf::Dot
                             | Mf::Outer
-                            | Mf::Cross
                             | Mf::Distance
                             | Mf::Length
                             | Mf::Normalize
@@ -1347,12 +1346,116 @@ impl<'a> ConstantEvaluator<'a> {
                 component_wise_concrete_int(self, span, [arg], |ci| Ok(first_leading_bit(ci)))
             }
 
+            // vector
+            crate::MathFunction::Cross => self.cross_product(arg, arg1.unwrap(), span),
+
             fun => Err(ConstantEvaluatorError::NotImplemented(format!(
                 "{fun:?} built-in function"
             ))),
         }
     }
 
+    /// Vector cross product.
+    fn cross_product(
+        &mut self,
+        a: Handle<Expression>,
+        b: Handle<Expression>,
+        span: Span,
+    ) -> Result<Handle<Expression>, ConstantEvaluatorError> {
+        use Literal as Li;
+
+        let (a, ty) = self.extract_vec::<3>(a)?;
+        let (b, _) = self.extract_vec::<3>(b)?;
+
+        let product = match (a, b) {
+            (
+                [Li::AbstractInt(a0), Li::AbstractInt(a1), Li::AbstractInt(a2)],
+                [Li::AbstractInt(b0), Li::AbstractInt(b1), Li::AbstractInt(b2)],
+            ) => {
+                // `cross` has no overload for AbstractInt, so AbstractInt
+                // arguments are automatically converted to AbstractFloat. Since
+                // `f64` has a much wider range than `i64`, there's no danger of
+                // overflow here.
+                let p = cross_product(
+                    [a0 as f64, a1 as f64, a2 as f64],
+                    [b0 as f64, b1 as f64, b2 as f64],
+                );
+                [
+                    Li::AbstractFloat(p[0]),
+                    Li::AbstractFloat(p[1]),
+                    Li::AbstractFloat(p[2]),
+                ]
+            }
+            (
+                [Li::AbstractFloat(a0), Li::AbstractFloat(a1), Li::AbstractFloat(a2)],
+                [Li::AbstractFloat(b0), Li::AbstractFloat(b1), Li::AbstractFloat(b2)],
+            ) => {
+                let p = cross_product([a0, a1, a2], [b0, b1, b2]);
+                [
+                    Li::AbstractFloat(p[0]),
+                    Li::AbstractFloat(p[1]),
+                    Li::AbstractFloat(p[2]),
+                ]
+            }
+            ([Li::F16(a0), Li::F16(a1), Li::F16(a2)], [Li::F16(b0), Li::F16(b1), Li::F16(b2)]) => {
+                let p = cross_product([a0, a1, a2], [b0, b1, b2]);
+                [Li::F16(p[0]), Li::F16(p[1]), Li::F16(p[2])]
+            }
+            ([Li::F32(a0), Li::F32(a1), Li::F32(a2)], [Li::F32(b0), Li::F32(b1), Li::F32(b2)]) => {
+                let p = cross_product([a0, a1, a2], [b0, b1, b2]);
+                [Li::F32(p[0]), Li::F32(p[1]), Li::F32(p[2])]
+            }
+            ([Li::F64(a0), Li::F64(a1), Li::F64(a2)], [Li::F64(b0), Li::F64(b1), Li::F64(b2)]) => {
+                let p = cross_product([a0, a1, a2], [b0, b1, b2]);
+                [Li::F64(p[0]), Li::F64(p[1]), Li::F64(p[2])]
+            }
+            _ => return Err(ConstantEvaluatorError::InvalidMathArg),
+        };
+
+        let p0 = self.register_evaluated_expr(Expression::Literal(product[0]), span)?;
+        let p1 = self.register_evaluated_expr(Expression::Literal(product[1]), span)?;
+        let p2 = self.register_evaluated_expr(Expression::Literal(product[2]), span)?;
+
+        self.register_evaluated_expr(
+            Expression::Compose {
+                ty,
+                components: vec![p0, p1, p2],
+            },
+            span,
+        )
+    }
+
+    /// Extract the values of a `vecN` from `expr`.
+    ///
+    /// Return the value of `expr`, whose type is `vecN<S>` for some
+    /// vector size `N` and scalar `S`, as an array of `N` [`Literal`]
+    /// values.
+    ///
+    /// Also return the type handle from the `Compose` expression.
+    fn extract_vec<const N: usize>(
+        &mut self,
+        expr: Handle<Expression>,
+    ) -> Result<([Literal; N], Handle<Type>), ConstantEvaluatorError> {
+        let span = self.expressions.get_span(expr);
+        let expr = self.eval_zero_value_and_splat(expr, span)?;
+        let Expression::Compose { ty, ref components } = self.expressions[expr] else {
+            return Err(ConstantEvaluatorError::InvalidMathArg);
+        };
+
+        let mut value = [Literal::Bool(false); N];
+        for (component, elt) in
+            crate::proc::flatten_compose(ty, components, self.expressions, self.types)
+                .zip(value.iter_mut())
+        {
+            let Expression::Literal(literal) = self.expressions[component] else {
+                return Err(ConstantEvaluatorError::InvalidMathArg);
+            };
+            *elt = literal;
+        }
+
+        Ok((value, ty))
+    }
+
     fn array_length(
         &mut self,
         array: Handle<Expression>,
@@ -2689,6 +2792,19 @@ impl TryFromAbstract<i64> for f16 {
     }
 }
 
+fn cross_product<T>(a: [T; 3], b: [T; 3]) -> [T; 3]
+where
+    T: Copy,
+    T: core::ops::Mul<T, Output = T>,
+    T: core::ops::Sub<T, Output = T>,
+{
+    [
+        a[1] * b[2] - a[2] * b[1],
+        a[2] * b[0] - a[0] * b[2],
+        a[0] * b[1] - a[1] * b[0],
+    ]
+}
+
 #[cfg(test)]
 mod tests {
     use alloc::{vec, vec::Vec};

naga/tests/in/wgsl/cross.wgsl

@@ -1,4 +1,4 @@
 // NOTE: invalid combinations are tested in the `validation::bad_cross_builtin_args` test.
 @compute @workgroup_size(1) fn main() {
-    let a = cross(vec3(0., 1., 2.), vec3(0., 1., 2.));
+    let a = cross(vec3(1., 0., 0.), vec3(0., 1., 0.));
 }

naga/tests/out/glsl/cross.main.Compute.glsl

@@ -7,7 +7,7 @@ layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
 
 
 void main() {
-    vec3 a = cross(vec3(0.0, 1.0, 2.0), vec3(0.0, 1.0, 2.0));
+    vec3 a = vec3(0.0, 0.0, 1.0);
     return;
 }
 

naga/tests/out/hlsl/cross.hlsl

@@ -1,6 +1,6 @@
 [numthreads(1, 1, 1)]
 void main()
 {
-    float3 a = cross(float3(0.0, 1.0, 2.0), float3(0.0, 1.0, 2.0));
+    float3 a = float3(0.0, 0.0, 1.0);
     return;
 }

naga/tests/out/msl/cross.msl

@@ -7,6 +7,6 @@ using metal::uint;
 
 kernel void main_(
 ) {
-    metal::float3 a = metal::cross(metal::float3(0.0, 1.0, 2.0), metal::float3(0.0, 1.0, 2.0));
+    metal::float3 a = metal::float3(0.0, 0.0, 1.0);
     return;
 }

naga/tests/out/spv/cross.spvasm

@@ -1,7 +1,7 @@
 ; SPIR-V
 ; Version: 1.1
 ; Generator: rspirv
-; Bound: 14
+; Bound: 12
 OpCapability Shader
 %1 = OpExtInstImport "GLSL.std.450"
 OpMemoryModel Logical GLSL450
@@ -13,12 +13,10 @@ OpExecutionMode %6 LocalSize 1 1 1
 %7 = OpTypeFunction %2
 %8 = OpConstant  %4  0.0
 %9 = OpConstant  %4  1.0
-%10 = OpConstant  %4  2.0
-%11 = OpConstantComposite  %3  %8 %9 %10
+%10 = OpConstantComposite  %3  %8 %8 %9
 %6 = OpFunction  %2  None %7
 %5 = OpLabel
-OpBranch %12
-%12 = OpLabel
-%13 = OpExtInst  %3  %1 Cross %11 %11
+OpBranch %11
+%11 = OpLabel
 OpReturn
 OpFunctionEnd

naga/tests/out/wgsl/cross.wgsl

@@ -1,5 +1,5 @@
 @compute @workgroup_size(1, 1, 1) 
 fn main() {
-    let a = cross(vec3<f32>(0f, 1f, 2f), vec3<f32>(0f, 1f, 2f));
+    const a = vec3<f32>(0f, 0f, 1f);
     return;
 }

naga/tests/validation.rs

@@ -261,56 +261,86 @@ fn emit_workgroup_uniform_load_result() {
     assert!(variant(false).is_err());
 }
 
-#[cfg(feature = "wgsl-in")]
 #[test]
-fn bad_cross_builtin_args() {
-    // NOTE: Things we expect to actually compile are in the `cross` snapshot test.
-    let cases = [
-        (
-            "vec2(0., 1.)",
-            "\
-error: Entry point main at Compute is invalid
-  ┌─ wgsl:3:13
-  │
-3 │     let a = cross(vec2(0., 1.), vec2(0., 1.));
-  │             ^^^^^ naga::ir::Expression [6]
-  │
-  = Expression [6] is invalid
-  = Argument [0] to Cross as expression [2] has an invalid type.
+fn builtin_cross_product_args() {
+    use naga::{MathFunction, Module, Type, TypeInner, VectorSize};
 
-",
-        ),
-        (
-            "vec4(0., 1., 2., 3.)",
-            "\
-error: Entry point main at Compute is invalid
-  ┌─ wgsl:3:13
-  │
-3 │     let a = cross(vec4(0., 1., 2., 3.), vec4(0., 1., 2., 3.));
-  │             ^^^^^ naga::ir::Expression [10]
-  │
-  = Expression [10] is invalid
-  = Argument [0] to Cross as expression [4] has an invalid type.
+    // We want to ensure that the *only* problem with the code is the
+    // arity of the vectors passed to `cross`. So validate two
+    // versions of the module varying only in that aspect.
+    //
+    // Looking at uses of the `wg_load` makes it easy to identify the
+    // differences between the two variants.
+    fn variant(
+        size: VectorSize,
+    ) -> Result<naga::valid::ModuleInfo, naga::WithSpan<naga::valid::ValidationError>> {
+        let span = naga::Span::default();
+        let mut module = Module::default();
+        let ty_vec3f = module.types.insert(
+            Type {
+                name: Some("vecnf".into()),
+                inner: TypeInner::Vector {
+                    size: VectorSize::Tri,
+                    scalar: Scalar::F32,
+                },
+            },
+            span,
+        );
+        let ty_vecnf = module.types.insert(
+            Type {
+                name: Some("vecnf".into()),
+                inner: TypeInner::Vector {
+                    size,
+                    scalar: Scalar::F32,
+                },
+            },
+            span,
+        );
 
-",
-        ),
-    ];
+        let mut fun = Function {
+            result: Some(naga::ir::FunctionResult {
+                ty: ty_vec3f,
+                binding: None,
+            }),
+            ..Function::default()
+        };
+        let ex_zero = fun
+            .expressions
+            .append(Expression::ZeroValue(ty_vecnf), span);
+        let ex_cross = fun.expressions.append(
+            Expression::Math {
+                fun: MathFunction::Cross,
+                arg: ex_zero,
+                arg1: Some(ex_zero),
+                arg2: None,
+                arg3: None,
+            },
+            span,
+        );
 
-    for (invalid_arg, expected_err) in cases {
-        let source = format!(
-            "\
-@compute @workgroup_size(1)
-fn main() {{
-    let a = cross({invalid_arg}, {invalid_arg});
-}}
-"
+        fun.body.push(
+            naga::Statement::Emit(naga::Range::new_from_bounds(ex_cross, ex_cross)),
+            span,
         );
-        let module = naga::front::wgsl::parse_str(&source).unwrap();
-        let err = valid::Validator::new(Default::default(), valid::Capabilities::all())
-            .validate(&module)
-            .expect_err("module should be invalid");
-        assert_eq!(err.emit_to_string(&source), expected_err);
+        fun.body.push(
+            naga::Statement::Return {
+                value: Some(ex_cross),
+            },
+            span,
+        );
+
+        module.functions.append(fun, span);
+
+        valid::Validator::new(
+            valid::ValidationFlags::default(),
+            valid::Capabilities::all(),
+        )
+        .validate(&module)
     }
+
+    assert!(variant(VectorSize::Bi).is_err());
+    variant(VectorSize::Tri).expect("module should validate");
+    assert!(variant(VectorSize::Quad).is_err());
 }
 
 #[cfg(feature = "wgsl-in")]

naga/tests/wgsl_errors.rs

@@ -7,10 +7,10 @@ use naga::valid::Capabilities;
 
 #[track_caller]
 fn check(input: &str, snapshot: &str) {
-    let output = naga::front::wgsl::parse_str(input)
-        .map(|_| panic!("expected parser error, but parsing succeeded!"))
-        .unwrap_err()
-        .emit_to_string(input);
+    let output = match naga::front::wgsl::parse_str(input) {
+        Ok(_) => panic!("expected parser error, but parsing succeeded!"),
+        Err(err) => err.emit_to_string(input),
+    };
     if output != snapshot {
         for diff in diff::lines(snapshot, &output) {
             match diff {