[naga] Let constant evaluation handle Compose of Splat.

CHANGELOG.md

@@ -63,6 +63,8 @@ Bottom level categories:
 
 - 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).
 
+- 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).
+
 ## v0.18.1 (2023-11-15)
 
 (naga version 0.14.1)

naga/src/proc/mod.rs

@@ -695,17 +695,19 @@ pub fn flatten_compose<'arenas>(
     expressions: &'arenas crate::Arena<crate::Expression>,
     types: &'arenas crate::UniqueArena<crate::Type>,
 ) -> impl Iterator<Item = crate::Handle<crate::Expression>> + 'arenas {
-    // Returning `impl Iterator` is a bit tricky. We may or may not want to
-    // flatten the components, but we have to settle on a single concrete
-    // type to return. The below is a single iterator chain that handles
-    // both the flattening and non-flattening cases.
+    // Returning `impl Iterator` is a bit tricky. We may or may not
+    // want to flatten the components, but we have to settle on a
+    // single concrete type to return. This function returns a single
+    // iterator chain that handles both the flattening and
+    // non-flattening cases.
     let (size, is_vector) = if let crate::TypeInner::Vector { size, .. } = types[ty].inner {
         (size as usize, true)
     } else {
         (components.len(), false)
     };
 
-    fn flattener<'c>(
+    /// Flatten `Compose` expressions if `is_vector` is true.
+    fn flatten_compose<'c>(
         component: &'c crate::Handle<crate::Expression>,
         is_vector: bool,
         expressions: &'c crate::Arena<crate::Expression>,
@@ -722,14 +724,35 @@ pub fn flatten_compose<'arenas>(
         std::slice::from_ref(component)
     }
 
-    // Expressions like `vec4(vec3(vec2(6, 7), 8), 9)` require us to flatten
-    // two levels.
+    /// Flatten `Splat` expressions if `is_vector` is true.
+    fn flatten_splat<'c>(
+        component: &'c crate::Handle<crate::Expression>,
+        is_vector: bool,
+        expressions: &'c crate::Arena<crate::Expression>,
+    ) -> impl Iterator<Item = crate::Handle<crate::Expression>> {
+        let mut expr = *component;
+        let mut count = 1;
+        if is_vector {
+            if let crate::Expression::Splat { size, value } = expressions[expr] {
+                expr = value;
+                count = size as usize;
+            }
+        }
+        std::iter::repeat(expr).take(count)
+    }
+
+    // Expressions like `vec4(vec3(vec2(6, 7), 8), 9)` require us to
+    // flatten up to two levels of `Compose` expressions.
+    //
+    // Expressions like `vec4(vec3(1.0), 1.0)` require us to flatten
+    // `Splat` expressions. Fortunately, the operand of a `Splat` must
+    // be a scalar, so we can stop there.
     components
         .iter()
-        .flat_map(move |component| flattener(component, is_vector, expressions))
-        .flat_map(move |component| flattener(component, is_vector, expressions))
+        .flat_map(move |component| flatten_compose(component, is_vector, expressions))
+        .flat_map(move |component| flatten_compose(component, is_vector, expressions))
+        .flat_map(move |component| flatten_splat(component, is_vector, expressions))
         .take(size)
-        .cloned()
 }
 
 #[test]

naga/tests/in/const-exprs.wgsl

@@ -9,6 +9,7 @@ fn main() {
     non_constant_initializers();
     splat_of_constant();
     compose_of_constant();
+    compose_of_splat();
 }
 
 // Swizzle the value of nested Compose expressions.
@@ -79,3 +80,7 @@ fn map_texture_kind(texture_kind: i32) -> u32 {
         default: { return 0u; }
     }
 }
+
+fn compose_of_splat() {
+    var x = vec4f(vec3f(1.0), 2.0).wzyx;
+}

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

@@ -57,6 +57,10 @@ void compose_of_constant() {
     ivec4 out_5 = ivec4(-4, -4, -4, -4);
 }
 
+void compose_of_splat() {
+    vec4 x_1 = vec4(2.0, 1.0, 1.0, 1.0);
+}
+
 uint map_texture_kind(int texture_kind) {
     switch(texture_kind) {
         case 0: {
@@ -81,6 +85,7 @@ void main() {
     non_constant_initializers();
     splat_of_constant();
     compose_of_constant();
+    compose_of_splat();
     return;
 }
 

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

@@ -61,6 +61,12 @@ void compose_of_constant()
 
 }
 
+void compose_of_splat()
+{
+    float4 x_1 = float4(2.0, 1.0, 1.0, 1.0);
+
+}
+
 uint map_texture_kind(int texture_kind)
 {
     switch(texture_kind) {
@@ -88,5 +94,6 @@ void main()
     non_constant_initializers();
     splat_of_constant();
     compose_of_constant();
+    compose_of_splat();
     return;
 }

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

@@ -61,6 +61,11 @@ void compose_of_constant(
     metal::int4 out_5 = metal::int4(-4, -4, -4, -4);
 }
 
+void compose_of_splat(
+) {
+    metal::float4 x_1 = metal::float4(2.0, 1.0, 1.0, 1.0);
+}
+
 uint map_texture_kind(
     int texture_kind
 ) {
@@ -88,5 +93,6 @@ kernel void main_(
     non_constant_initializers();
     splat_of_constant();
     compose_of_constant();
+    compose_of_splat();
     return;
 }

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

@@ -1,27 +1,27 @@
 ; SPIR-V
 ; Version: 1.1
 ; Generator: rspirv
-; Bound: 91
+; Bound: 100
 OpCapability Shader
 %1 = OpExtInstImport "GLSL.std.450"
 OpMemoryModel Logical GLSL450
-OpEntryPoint GLCompute %83 "main"
-OpExecutionMode %83 LocalSize 2 3 1
+OpEntryPoint GLCompute %91 "main"
+OpExecutionMode %91 LocalSize 2 3 1
 %2 = OpTypeVoid
 %3 = OpTypeInt 32 0
 %4 = OpTypeInt 32 1
 %5 = OpTypeVector %4 4
-%6 = OpTypeFloat 32
-%7 = OpTypeVector %6 4
+%7 = OpTypeFloat 32
+%6 = OpTypeVector %7 4
 %8 = OpConstant  %3  2
 %9 = OpConstant  %4  3
 %10 = OpConstant  %4  4
 %11 = OpConstant  %4  8
-%12 = OpConstant  %6  3.141
-%13 = OpConstant  %6  6.282
-%14 = OpConstant  %6  0.44444445
-%15 = OpConstant  %6  0.0
-%16 = OpConstantComposite  %7  %14 %15 %15 %15
+%12 = OpConstant  %7  3.141
+%13 = OpConstant  %7  6.282
+%14 = OpConstant  %7  0.44444445
+%15 = OpConstant  %7  0.0
+%16 = OpConstantComposite  %6  %14 %15 %15 %15
 %17 = OpConstant  %4  0
 %18 = OpConstant  %4  1
 %19 = OpConstant  %4  2
@@ -37,12 +37,16 @@ OpExecutionMode %83 LocalSize 2 3 1
 %48 = OpConstantNull  %5
 %59 = OpConstant  %4  -4
 %60 = OpConstantComposite  %5  %59 %59 %59 %59
-%70 = OpTypeFunction %3 %4
-%71 = OpConstant  %3  10
-%72 = OpConstant  %3  20
-%73 = OpConstant  %3  30
-%74 = OpConstant  %3  0
-%81 = OpConstantNull  %3
+%69 = OpConstant  %7  1.0
+%70 = OpConstant  %7  2.0
+%71 = OpConstantComposite  %6  %70 %69 %69 %69
+%73 = OpTypePointer Function %6
+%78 = OpTypeFunction %3 %4
+%79 = OpConstant  %3  10
+%80 = OpConstant  %3  20
+%81 = OpConstant  %3  30
+%82 = OpConstant  %3  0
+%89 = OpConstantNull  %3
 %21 = OpFunction  %2  None %22
 %20 = OpLabel
 %24 = OpVariable  %25  Function %23
@@ -99,33 +103,41 @@ OpBranch %66
 %66 = OpLabel
 OpReturn
 OpFunctionEnd
-%69 = OpFunction  %3  None %70
-%68 = OpFunctionParameter  %4
+%68 = OpFunction  %2  None %22
 %67 = OpLabel
-OpBranch %75
+%72 = OpVariable  %73  Function %71
+OpBranch %74
+%74 = OpLabel
+OpReturn
+OpFunctionEnd
+%77 = OpFunction  %3  None %78
+%76 = OpFunctionParameter  %4
 %75 = OpLabel
-OpSelectionMerge %76 None
-OpSwitch %68 %80 0 %77 1 %78 2 %79
-%77 = OpLabel
-OpReturnValue %71
-%78 = OpLabel
-OpReturnValue %72
-%79 = OpLabel
-OpReturnValue %73
-%80 = OpLabel
-OpReturnValue %74
-%76 = OpLabel
+OpBranch %83
+%83 = OpLabel
+OpSelectionMerge %84 None
+OpSwitch %76 %88 0 %85 1 %86 2 %87
+%85 = OpLabel
+OpReturnValue %79
+%86 = OpLabel
+OpReturnValue %80
+%87 = OpLabel
 OpReturnValue %81
-OpFunctionEnd
-%83 = OpFunction  %2  None %22
-%82 = OpLabel
-OpBranch %84
+%88 = OpLabel
+OpReturnValue %82
 %84 = OpLabel
-%85 = OpFunctionCall  %2  %21
-%86 = OpFunctionCall  %2  %28
-%87 = OpFunctionCall  %2  %33
-%88 = OpFunctionCall  %2  %38
-%89 = OpFunctionCall  %2  %58
-%90 = OpFunctionCall  %2  %64
+OpReturnValue %89
+OpFunctionEnd
+%91 = OpFunction  %2  None %22
+%90 = OpLabel
+OpBranch %92
+%92 = OpLabel
+%93 = OpFunctionCall  %2  %21
+%94 = OpFunctionCall  %2  %28
+%95 = OpFunctionCall  %2  %33
+%96 = OpFunctionCall  %2  %38
+%97 = OpFunctionCall  %2  %58
+%98 = OpFunctionCall  %2  %64
+%99 = OpFunctionCall  %2  %68
 OpReturn
 OpFunctionEnd

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

@@ -55,6 +55,11 @@ fn compose_of_constant() {
 
 }
 
+fn compose_of_splat() {
+    var x_1: vec4<f32> = vec4<f32>(2.0, 1.0, 1.0, 1.0);
+
+}
+
 fn map_texture_kind(texture_kind: i32) -> u32 {
     switch texture_kind {
         case 0: {
@@ -80,5 +85,6 @@ fn main() {
     non_constant_initializers();
     splat_of_constant();
     compose_of_constant();
+    compose_of_splat();
     return;
 }