[naga] Teach the constant evaluator vector/vector operators.

Allow constant evaluation of binary operators whose left and right
operands are both vectors.

CHANGELOG.md

@@ -68,7 +68,8 @@ This feature allowed you to call `global_id` on any wgpu opaque handle to get a
 
 #### Naga
 
-- Naga's WGSL front and back ends now have experimental support for 64-bit floating-point literals: `1.0lf` denotes an `f64` value. There has been experimental support for an `f64` type for a while, but until now there was no syntax for writing literals with that type. As before, Naga module validation rejects `f64` values unless `naga::valid::Capabilities::FLOAT64` is requested. By @jimblandy in [#4747](https://github.com/gfx-rs/wgpu/pull/4747).
+- Naga'sn WGSL front and back ends now have experimental support for 64-bit floating-point literals: `1.0lf` denotes an `f64` value. There has been experimental support for an `f64` type for a while, but until now there was no syntax for writing literals with that type. As before, Naga module validation rejects `f64` values unless `naga::valid::Capabilities::FLOAT64` is requested. By @jimblandy in [#4747](https://github.com/gfx-rs/wgpu/pull/4747).
+- Naga constant evaluation can now process binary operators whose operands are both vectors. By @jimblandy in [#4861](https://github.com/gfx-rs/wgpu/pull/4861).
 
 ### Changes
 

naga/src/proc/constant_evaluator.rs

@@ -1361,12 +1361,107 @@ impl<'a> ConstantEvaluator<'a> {
                 }
                 Expression::Compose { ty, components }
             }
+            (
+                &Expression::Compose {
+                    components: ref left_components,
+                    ty: left_ty,
+                },
+                &Expression::Compose {
+                    components: ref right_components,
+                    ty: right_ty,
+                },
+            ) => {
+                // We have to make a copy of the component lists, because the
+                // call to `binary_op_vector` needs `&mut self`, but `self` owns
+                // the component lists.
+                let left_flattened = crate::proc::flatten_compose(
+                    left_ty,
+                    left_components,
+                    self.expressions,
+                    self.types,
+                );
+                let right_flattened = crate::proc::flatten_compose(
+                    right_ty,
+                    right_components,
+                    self.expressions,
+                    self.types,
+                );
+
+                // `flatten_compose` doesn't return an `ExactSizeIterator`, so
+                // make a reasonable guess of the capacity we'll need.
+                let mut flattened = Vec::with_capacity(left_components.len());
+                flattened.extend(left_flattened.zip(right_flattened));
+
+                match (&self.types[left_ty].inner, &self.types[right_ty].inner) {
+                    (
+                        &TypeInner::Vector {
+                            size: left_size, ..
+                        },
+                        &TypeInner::Vector {
+                            size: right_size, ..
+                        },
+                    ) if left_size == right_size => {
+                        self.binary_op_vector(op, left_size, &flattened, left_ty, span)?
+                    }
+                    _ => return Err(ConstantEvaluatorError::InvalidBinaryOpArgs),
+                }
+            }
             _ => return Err(ConstantEvaluatorError::InvalidBinaryOpArgs),
         };
 
         self.register_evaluated_expr(expr, span)
     }
 
+    fn binary_op_vector(
+        &mut self,
+        op: BinaryOperator,
+        size: crate::VectorSize,
+        components: &[(Handle<Expression>, Handle<Expression>)],
+        left_ty: Handle<Type>,
+        span: Span,
+    ) -> Result<Expression, ConstantEvaluatorError> {
+        let ty = match op {
+            // Relational operators produce vectors of booleans.
+            BinaryOperator::Equal
+            | BinaryOperator::NotEqual
+            | BinaryOperator::Less
+            | BinaryOperator::LessEqual
+            | BinaryOperator::Greater
+            | BinaryOperator::GreaterEqual => self.types.insert(
+                Type {
+                    name: None,
+                    inner: TypeInner::Vector {
+                        size,
+                        scalar: crate::Scalar::BOOL,
+                    },
+                },
+                span,
+            ),
+
+            // Other operators produce the same type as their left
+            // operand.
+            BinaryOperator::Add
+            | BinaryOperator::Subtract
+            | BinaryOperator::Multiply
+            | BinaryOperator::Divide
+            | BinaryOperator::Modulo
+            | BinaryOperator::And
+            | BinaryOperator::ExclusiveOr
+            | BinaryOperator::InclusiveOr
+            | BinaryOperator::LogicalAnd
+            | BinaryOperator::LogicalOr
+            | BinaryOperator::ShiftLeft
+            | BinaryOperator::ShiftRight => left_ty,
+        };
+
+        let components = components
+            .iter()
+            .map(|&(left, right)| self.binary_op(op, left, right, span))
+            .collect::<Result<Vec<_>, _>>()?;
+
+        Ok(Expression::Compose { ty, components })
+    }
+
     /// Deep copy `expr` from `expressions` into `self.expressions`.
     ///
     /// Return the root of the new copy.

naga/tests/in/const-exprs.wgsl

@@ -84,3 +84,6 @@ fn map_texture_kind(texture_kind: i32) -> u32 {
 fn compose_of_splat() {
     var x = vec4f(vec3f(1.0), 2.0).wzyx;
 }
+
+const add_vec = vec2(1.0f) + vec2(3.0f, 4.0f);
+const compare_vec = vec2(3.0f) == vec2(3.0f, 4.0f);

naga/tests/out/glsl/const-exprs.main.Compute.glsl

@@ -17,6 +17,8 @@ const vec4 DIV = vec4(0.44444445, 0.0, 0.0, 0.0);
 const int TEXTURE_KIND_REGULAR = 0;
 const int TEXTURE_KIND_WARP = 1;
 const int TEXTURE_KIND_SKY = 2;
+const vec2 add_vec = vec2(4.0, 5.0);
+const bvec2 compare_vec = bvec2(true, false);
 
 
 void swizzle_of_compose() {

naga/tests/out/hlsl/const-exprs.hlsl

@@ -10,6 +10,8 @@ static const float4 DIV = float4(0.44444445, 0.0, 0.0, 0.0);
 static const int TEXTURE_KIND_REGULAR = 0;
 static const int TEXTURE_KIND_WARP = 1;
 static const int TEXTURE_KIND_SKY = 2;
+static const float2 add_vec = float2(4.0, 5.0);
+static const bool2 compare_vec = bool2(true, false);
 
 void swizzle_of_compose()
 {

naga/tests/out/msl/const-exprs.msl

@@ -16,6 +16,8 @@ constant metal::float4 DIV = metal::float4(0.44444445, 0.0, 0.0, 0.0);
 constant int TEXTURE_KIND_REGULAR = 0;
 constant int TEXTURE_KIND_WARP = 1;
 constant int TEXTURE_KIND_SKY = 2;
+constant metal::float2 add_vec = metal::float2(4.0, 5.0);
+constant metal::bool2 compare_vec = metal::bool2(true, false);
 
 void swizzle_of_compose(
 ) {

naga/tests/out/spv/const-exprs.spvasm

@@ -1,143 +1,152 @@
 ; SPIR-V
 ; Version: 1.1
 ; Generator: rspirv
-; Bound: 100
+; Bound: 109
 OpCapability Shader
 %1 = OpExtInstImport "GLSL.std.450"
 OpMemoryModel Logical GLSL450
-OpEntryPoint GLCompute %91 "main"
+OpEntryPoint GLCompute %100 "main"
-OpExecutionMode %91 LocalSize 2 3 1
+OpExecutionMode %100 LocalSize 2 3 1
 %2 = OpTypeVoid
 %3 = OpTypeInt 32 0
 %4 = OpTypeInt 32 1
 %5 = OpTypeVector %4 4
 %7 = OpTypeFloat 32
 %6 = OpTypeVector %7 4
-%8 = OpConstant  %3  2
+%8 = OpTypeVector %7 2
-%9 = OpConstant  %4  3
+%10 = OpTypeBool
-%10 = OpConstant  %4  4
+%9 = OpTypeVector %10 2
-%11 = OpConstant  %4  8
+%11 = OpConstant  %3  2
-%12 = OpConstant  %7  3.141
+%12 = OpConstant  %4  3
-%13 = OpConstant  %7  6.282
+%13 = OpConstant  %4  4
-%14 = OpConstant  %7  0.44444445
+%14 = OpConstant  %4  8
-%15 = OpConstant  %7  0.0
+%15 = OpConstant  %7  3.141
-%16 = OpConstantComposite  %6  %14 %15 %15 %15
+%16 = OpConstant  %7  6.282
-%17 = OpConstant  %4  0
+%17 = OpConstant  %7  0.44444445
-%18 = OpConstant  %4  1
+%18 = OpConstant  %7  0.0
-%19 = OpConstant  %4  2
+%19 = OpConstantComposite  %6  %17 %18 %18 %18
-%22 = OpTypeFunction %2
+%20 = OpConstant  %4  0
-%23 = OpConstantComposite  %5  %10 %9 %19 %18
+%21 = OpConstant  %4  1
-%25 = OpTypePointer Function %5
+%22 = OpConstant  %4  2
-%30 = OpTypePointer Function %4
+%23 = OpConstant  %7  4.0
-%34 = OpConstant  %4  6
+%24 = OpConstant  %7  5.0
-%39 = OpConstant  %4  30
+%25 = OpConstantComposite  %8  %23 %24
-%40 = OpConstant  %4  70
+%26 = OpConstantTrue  %10
-%43 = OpConstantNull  %4
+%27 = OpConstantFalse  %10
-%45 = OpConstantNull  %4
+%28 = OpConstantComposite  %9  %26 %27
-%48 = OpConstantNull  %5
+%31 = OpTypeFunction %2
-%59 = OpConstant  %4  -4
+%32 = OpConstantComposite  %5  %13 %12 %22 %21
-%60 = OpConstantComposite  %5  %59 %59 %59 %59
+%34 = OpTypePointer Function %5
-%69 = OpConstant  %7  1.0
+%39 = OpTypePointer Function %4
-%70 = OpConstant  %7  2.0
+%43 = OpConstant  %4  6
-%71 = OpConstantComposite  %6  %70 %69 %69 %69
+%48 = OpConstant  %4  30
-%73 = OpTypePointer Function %6
+%49 = OpConstant  %4  70
-%78 = OpTypeFunction %3 %4
+%52 = OpConstantNull  %4
-%79 = OpConstant  %3  10
+%54 = OpConstantNull  %4
-%80 = OpConstant  %3  20
+%57 = OpConstantNull  %5
-%81 = OpConstant  %3  30
+%68 = OpConstant  %4  -4
-%82 = OpConstant  %3  0
+%69 = OpConstantComposite  %5  %68 %68 %68 %68
-%89 = OpConstantNull  %3
+%78 = OpConstant  %7  1.0
-%21 = OpFunction  %2  None %22
+%79 = OpConstant  %7  2.0
-%20 = OpLabel
+%80 = OpConstantComposite  %6  %79 %78 %78 %78
-%24 = OpVariable  %25  Function %23
+%82 = OpTypePointer Function %6
-OpBranch %26
+%87 = OpTypeFunction %3 %4
-%26 = OpLabel
+%88 = OpConstant  %3  10
+%89 = OpConstant  %3  20
+%90 = OpConstant  %3  30
+%91 = OpConstant  %3  0
+%98 = OpConstantNull  %3
+%30 = OpFunction  %2  None %31
+%29 = OpLabel
+%33 = OpVariable  %34  Function %32
+OpBranch %35
+%35 = OpLabel
 OpReturn
 OpFunctionEnd
-%28 = OpFunction  %2  None %22
+%37 = OpFunction  %2  None %31
-%27 = OpLabel
-%29 = OpVariable  %30  Function %19
-OpBranch %31
-%31 = OpLabel
-OpReturn
-OpFunctionEnd
-%33 = OpFunction  %2  None %22
-%32 = OpLabel
-%35 = OpVariable  %30  Function %34
-OpBranch %36
 %36 = OpLabel
+%38 = OpVariable  %39  Function %22
+OpBranch %40
+%40 = OpLabel
 OpReturn
 OpFunctionEnd
-%38 = OpFunction  %2  None %22
+%42 = OpFunction  %2  None %31
-%37 = OpLabel
+%41 = OpLabel
-%47 = OpVariable  %25  Function %48
+%44 = OpVariable  %39  Function %43
-%42 = OpVariable  %30  Function %43
+OpBranch %45
-%46 = OpVariable  %30  Function %40
+%45 = OpLabel
-%41 = OpVariable  %30  Function %39
-%44 = OpVariable  %30  Function %45
-OpBranch %49
-%49 = OpLabel
-%50 = OpLoad  %4  %41
-OpStore %42 %50
-%51 = OpLoad  %4  %42
-OpStore %44 %51
-%52 = OpLoad  %4  %41
-%53 = OpLoad  %4  %42
-%54 = OpLoad  %4  %44
-%55 = OpLoad  %4  %46
-%56 = OpCompositeConstruct  %5  %52 %53 %54 %55
-OpStore %47 %56
 OpReturn
 OpFunctionEnd
-%58 = OpFunction  %2  None %22
+%47 = OpFunction  %2  None %31
-%57 = OpLabel
+%46 = OpLabel
-%61 = OpVariable  %25  Function %60
+%56 = OpVariable  %34  Function %57
-OpBranch %62
+%51 = OpVariable  %39  Function %52
-%62 = OpLabel
+%55 = OpVariable  %39  Function %49
+%50 = OpVariable  %39  Function %48
+%53 = OpVariable  %39  Function %54
+OpBranch %58
+%58 = OpLabel
+%59 = OpLoad  %4  %50
+OpStore %51 %59
+%60 = OpLoad  %4  %51
+OpStore %53 %60
+%61 = OpLoad  %4  %50
+%62 = OpLoad  %4  %51
+%63 = OpLoad  %4  %53
+%64 = OpLoad  %4  %55
+%65 = OpCompositeConstruct  %5  %61 %62 %63 %64
+OpStore %56 %65
 OpReturn
 OpFunctionEnd
-%64 = OpFunction  %2  None %22
+%67 = OpFunction  %2  None %31
-%63 = OpLabel
-%65 = OpVariable  %25  Function %60
-OpBranch %66
 %66 = OpLabel
+%70 = OpVariable  %34  Function %69
+OpBranch %71
+%71 = OpLabel
 OpReturn
 OpFunctionEnd
-%68 = OpFunction  %2  None %22
+%73 = OpFunction  %2  None %31
-%67 = OpLabel
+%72 = OpLabel
-%72 = OpVariable  %73  Function %71
+%74 = OpVariable  %34  Function %69
-OpBranch %74
+OpBranch %75
-%74 = OpLabel
-OpReturn
-OpFunctionEnd
-%77 = OpFunction  %3  None %78
-%76 = OpFunctionParameter  %4
 %75 = OpLabel
+OpReturn
+OpFunctionEnd
+%77 = OpFunction  %2  None %31
+%76 = OpLabel
+%81 = OpVariable  %82  Function %80
 OpBranch %83
 %83 = OpLabel
-OpSelectionMerge %84 None
+OpReturn
-OpSwitch %76 %88 0 %85 1 %86 2 %87
-%85 = OpLabel
-OpReturnValue %79
-%86 = OpLabel
-OpReturnValue %80
-%87 = OpLabel
-OpReturnValue %81
-%88 = OpLabel
-OpReturnValue %82
-%84 = OpLabel
-OpReturnValue %89
 OpFunctionEnd
-%91 = OpFunction  %2  None %22
+%86 = OpFunction  %3  None %87
-%90 = OpLabel
+%85 = OpFunctionParameter  %4
+%84 = OpLabel
 OpBranch %92
 %92 = OpLabel
-%93 = OpFunctionCall  %2  %21
+OpSelectionMerge %93 None
-%94 = OpFunctionCall  %2  %28
+OpSwitch %85 %97 0 %94 1 %95 2 %96
-%95 = OpFunctionCall  %2  %33
+%94 = OpLabel
-%96 = OpFunctionCall  %2  %38
+OpReturnValue %88
-%97 = OpFunctionCall  %2  %58
+%95 = OpLabel
-%98 = OpFunctionCall  %2  %64
+OpReturnValue %89
-%99 = OpFunctionCall  %2  %68
+%96 = OpLabel
+OpReturnValue %90
+%97 = OpLabel
+OpReturnValue %91
+%93 = OpLabel
+OpReturnValue %98
+OpFunctionEnd
+%100 = OpFunction  %2  None %31
+%99 = OpLabel
+OpBranch %101
+%101 = OpLabel
+%102 = OpFunctionCall  %2  %30
+%103 = OpFunctionCall  %2  %37
+%104 = OpFunctionCall  %2  %42
+%105 = OpFunctionCall  %2  %47
+%106 = OpFunctionCall  %2  %67
+%107 = OpFunctionCall  %2  %73
+%108 = OpFunctionCall  %2  %77
 OpReturn
 OpFunctionEnd

naga/tests/out/wgsl/const-exprs.wgsl

@@ -10,6 +10,8 @@ const DIV: vec4<f32> = vec4<f32>(0.44444445, 0.0, 0.0, 0.0);
 const TEXTURE_KIND_REGULAR: i32 = 0;
 const TEXTURE_KIND_WARP: i32 = 1;
 const TEXTURE_KIND_SKY: i32 = 2;
+const add_vec: vec2<f32> = vec2<f32>(4.0, 5.0);
+const compare_vec: vec2<bool> = vec2<bool>(true, false);
 
 fn swizzle_of_compose() {
     var out: vec4<i32> = vec4<i32>(4, 3, 2, 1);