fix(naga spv-out): OpConstantComposite does not support flattening

naga/src/back/spv/block.rs

@@ -813,17 +813,20 @@ impl BlockContext<'_> {
             crate::Expression::Compose { ty, ref components } => {
                 self.temp_list.clear();
                 if self.expression_constness.is_const(expr_handle) {
-                    self.temp_list.extend(
-                        crate::proc::flatten_compose(
-                            ty,
-                            components,
-                            &self.ir_function.expressions,
-                            &self.ir_module.types,
-                        )
-                        .map(|component| self.cached[component]),
-                    );
-                    self.writer
-                        .get_constant_composite(LookupType::Handle(ty), &self.temp_list)
+                    let flattened: Vec<_> = crate::proc::flatten_compose(
+                        ty,
+                        components,
+                        &self.ir_function.expressions,
+                        &self.ir_module.types,
+                    )
+                    .map(|component| (component, self.cached[component]))
+                    .collect();
+                    self.writer.get_flattened_constant_composite(
+                        ty,
+                        &flattened,
+                        &self.ir_module.types,
+                        &self.ir_function.expressions,
+                    )
                 } else {
                     self.temp_list
                         .extend(components.iter().map(|&component| self.cached[component]));

naga/src/back/spv/helpers.rs

@@ -3,7 +3,10 @@ use alloc::{vec, vec::Vec};
 use arrayvec::ArrayVec;
 use spirv::Word;
 
-use crate::{Handle, UniqueArena};
+use crate::{
+    back::spv::{LookupType, NumericType},
+    Arena, Handle, UniqueArena,
+};
 
 pub(super) fn bytes_to_words(bytes: &[u8]) -> Vec<Word> {
     bytes
@@ -227,3 +230,38 @@ pub enum BindingDecorations {
     },
     None,
 }
+
+impl super::Writer {
+    /// Produce an `OpConstantComposite` instruction and return the result ID.
+    ///
+    /// Unlike [`super::Writer::get_constant_composite`], this function accepts and expands
+    /// nested zero-value vectors in the arguments to a vector composite.
+    /// [`crate::proc::flatten_compose`] should still be used beforehand to flatten other types
+    /// of nested values. (It cannot flatten zero values, because it can only yield expressions
+    /// already in the arena.)
+    pub(in crate::back::spv) fn get_flattened_constant_composite(
+        &mut self,
+        ty: Handle<crate::Type>,
+        flattened: &[(Handle<crate::Expression>, Word)],
+        types: &UniqueArena<crate::Type>,
+        expressions: &Arena<crate::Expression>,
+    ) -> Word {
+        let mut component_ids = Vec::new();
+        if let crate::TypeInner::Vector { .. } = types[ty].inner {
+            for &(component, id) in flattened {
+                if let crate::Expression::ZeroValue(component_ty) = expressions[component] {
+                    if let crate::TypeInner::Vector { size, scalar } = types[component_ty].inner {
+                        let scalar_type_id = self.get_numeric_type_id(NumericType::Scalar(scalar));
+                        let null_id = self.get_constant_null(scalar_type_id);
+                        component_ids.extend(core::iter::repeat_n(null_id, size as usize));
+                        continue;
+                    }
+                }
+                component_ids.push(id);
+            }
+        } else {
+            component_ids.extend(flattened.iter().map(|&(_, id)| id));
+        }
+        self.get_constant_composite(LookupType::Handle(ty), &component_ids)
+    }
+}

naga/src/back/spv/writer.rs

@@ -2681,6 +2681,10 @@ impl Writer {
         instruction.to_words(&mut self.logical_layout.declarations);
     }
 
+    /// Produce an `OpConstantComposite` instruction and return the result ID.
+    ///
+    /// The length of `constituent_ids` must match the number of components in the composite
+    /// type (nesting vectors within vectors is not allowed).
     pub(super) fn get_constant_composite(
         &mut self,
         ty: LookupType,
@@ -2750,15 +2754,20 @@ impl Writer {
                 self.get_constant_null(type_id)
             }
             crate::Expression::Compose { ty, ref components } => {
-                let component_ids: Vec<_> = crate::proc::flatten_compose(
+                let flattened: Vec<_> = crate::proc::flatten_compose(
                     ty,
                     components,
                     &ir_module.global_expressions,
                     &ir_module.types,
                 )
-                .map(|component| self.constant_ids[component])
+                .map(|component| (component, self.constant_ids[component]))
                 .collect();
-                self.get_constant_composite(LookupType::Handle(ty), component_ids.as_slice())
+                self.get_flattened_constant_composite(
+                    ty,
+                    &flattened,
+                    &ir_module.types,
+                    &ir_module.global_expressions,
+                )
             }
             crate::Expression::Splat { size, value } => {
                 let value_id = self.constant_ids[value];

naga/tests/in/wgsl/constructors.wgsl

@@ -23,6 +23,16 @@ const cp1 = vec2(0u);
 const cp2 = mat2x2(vec2(0.), vec2(0.));
 const cp3 = array(0, 1, 2, 3);
 
+// complex composites: a vector ZeroValue or Splat as a constructor argument
+const ccz0 = vec4<f32>(vec2<f32>(), 1.0, 1.0);    // ZeroValue(vec2) component
+const ccz1 = vec4<f32>(vec2<f32>(1.0), 0.0, 1.0); // Splat(vec2) component
+const ccz2 = vec4<f32>(vec2<f32>(1.0, 0.0), 0.0, 1.0); // vec2 component
+
+// matrix composites: columns are vectors, so a ZeroValue/Splat column must be
+// kept as a single (vector) constituent, not expanded into scalars
+const ccm0 = mat2x2<f32>(vec2<f32>(), vec2<f32>(1.0, 2.0)); // ZeroValue(vec2) column
+const ccm1 = mat2x2<f32>(vec2<f32>(1.0), vec2<f32>());      // Splat + ZeroValue columns
+
 @compute @workgroup_size(1)
 fn main() {
     var foo: Foo;
@@ -56,6 +66,15 @@ fn main() {
     let cit1 = mat2x2(vec2(0.), vec2(0.));
     let cit2 = array(0, 1, 2, 3);
 
+    // complex composites: a vector ZeroValue or Splat as a constructor argument
+    let ccz3 = vec4<f32>(vec2<f32>(), 1.0, 1.0);
+    let ccz4 = vec4<f32>(vec2<f32>(1.0), 0.0, 1.0);
+    let ccz5 = vec4<f32>(vec2<f32>(1.0, 0.0), 0.0, 1.0);
+
+    // matrix composites: ZeroValue/Splat columns kept as vector constituents
+    let ccm2 = mat2x2<f32>(vec2<f32>(), vec2<f32>(1.0, 2.0));
+    let ccm3 = mat2x2<f32>(vec2<f32>(1.0), vec2<f32>());
+
     // identity constructors
     let ic0 = bool(bool());
     let ic1 = i32(i32());

naga/tests/out/glsl/wgsl-constructors.main.Compute.glsl

@@ -21,6 +21,11 @@ const mat2x2 cz5_ = mat2x2(0.0);
 const Foo cz6_[3] = Foo[3](Foo(vec4(0.0), 0), Foo(vec4(0.0), 0), Foo(vec4(0.0), 0));
 const Foo cz7_ = Foo(vec4(0.0), 0);
 const uvec2 cp1_ = uvec2(0u);
+const vec4 ccz0_ = vec4(vec2(0.0), 1.0, 1.0);
+const vec4 ccz1_ = vec4(vec2(1.0), 0.0, 1.0);
+const vec4 ccz2_ = vec4(vec2(1.0, 0.0), 0.0, 1.0);
+const mat2x2 ccm0_ = mat2x2(vec2(0.0), vec2(1.0, 2.0));
+const mat2x2 ccm1_ = mat2x2(vec2(1.0), vec2(0.0));
 
 
 void main() {
@@ -33,6 +38,11 @@ void main() {
     uvec2 cit0_ = uvec2(0u);
     mat2x2 cit1_ = mat2x2(vec2(0.0), vec2(0.0));
     int cit2_[4] = int[4](0, 1, 2, 3);
+    vec4 ccz3_ = vec4(vec2(0.0), 1.0, 1.0);
+    vec4 ccz4_ = vec4(vec2(1.0), 0.0, 1.0);
+    vec4 ccz5_ = vec4(vec2(1.0, 0.0), 0.0, 1.0);
+    mat2x2 ccm2_ = mat2x2(vec2(0.0), vec2(1.0, 2.0));
+    mat2x2 ccm3_ = mat2x2(vec2(1.0), vec2(0.0));
     uvec2 ic4_ = uvec2(0u, 0u);
     mat2x3 ic5_ = mat2x3(vec3(0.0, 0.0, 0.0), vec3(0.0, 0.0, 0.0));
     return;

naga/tests/out/hlsl/wgsl-constructors.hlsl

@@ -45,6 +45,10 @@ Foo ZeroValueFoo() {
     return (Foo)0;
 }
 
+float2 ZeroValuefloat2() {
+    return (float2)0;
+}
+
 static const float3 const1_ = (0.0).xxx;
 static const float2x2 const3_ = float2x2(float2(0.0, 1.0), float2(2.0, 3.0));
 static const float2x2 const4_[1] = Constructarray1_float2x2_(float2x2(float2(0.0, 1.0), float2(2.0, 3.0)));
@@ -57,6 +61,11 @@ static const float2x2 cz5_ = ZeroValuefloat2x2();
 static const Foo cz6_[3] = ZeroValuearray3_Foo_();
 static const Foo cz7_ = ZeroValueFoo();
 static const uint2 cp1_ = (0u).xx;
+static const float4 ccz0_ = float4(ZeroValuefloat2(), 1.0, 1.0);
+static const float4 ccz1_ = float4((1.0).xx, 0.0, 1.0);
+static const float4 ccz2_ = float4(float2(1.0, 0.0), 0.0, 1.0);
+static const float2x2 ccm0_ = float2x2(ZeroValuefloat2(), float2(1.0, 2.0));
+static const float2x2 ccm1_ = float2x2((1.0).xx, ZeroValuefloat2());
 
 Foo ConstructFoo(float4 arg0, int arg1) {
     Foo ret = (Foo)0;
@@ -88,6 +97,11 @@ void main()
     uint2 cit0_ = (0u).xx;
     float2x2 cit1_ = float2x2((0.0).xx, (0.0).xx);
     int cit2_[4] = Constructarray4_int_(int(0), int(1), int(2), int(3));
+    float4 ccz3_ = float4(ZeroValuefloat2(), 1.0, 1.0);
+    float4 ccz4_ = float4((1.0).xx, 0.0, 1.0);
+    float4 ccz5_ = float4(float2(1.0, 0.0), 0.0, 1.0);
+    float2x2 ccm2_ = float2x2(ZeroValuefloat2(), float2(1.0, 2.0));
+    float2x2 ccm3_ = float2x2((1.0).xx, ZeroValuefloat2());
     uint2 ic4_ = uint2(0u, 0u);
     float2x3 ic5_ = float2x3(float3(0.0, 0.0, 0.0), float3(0.0, 0.0, 0.0));
     return;

naga/tests/out/msl/wgsl-constructors.metal

@@ -30,6 +30,11 @@ constant metal::float2x2 cz5_ = metal::float2x2 {};
 constant type_10 cz6_ = type_10 {};
 constant Foo cz7_ = Foo {};
 constant metal::uint2 cp1_ = metal::uint2(0u);
+constant metal::float4 ccz0_ = metal::float4(metal::float2 {}, 1.0, 1.0);
+constant metal::float4 ccz1_ = metal::float4(metal::float2(1.0), 0.0, 1.0);
+constant metal::float4 ccz2_ = metal::float4(metal::float2(1.0, 0.0), 0.0, 1.0);
+constant metal::float2x2 ccm0_ = metal::float2x2(metal::float2 {}, metal::float2(1.0, 2.0));
+constant metal::float2x2 ccm1_ = metal::float2x2(metal::float2(1.0), metal::float2 {});
 
 kernel void main_(
 ) {
@@ -42,6 +47,11 @@ kernel void main_(
     metal::uint2 cit0_ = metal::uint2(0u);
     metal::float2x2 cit1_ = metal::float2x2(metal::float2(0.0), metal::float2(0.0));
     type_12 cit2_ = type_12 {{0, 1, 2, 3}};
+    metal::float4 ccz3_ = metal::float4(metal::float2 {}, 1.0, 1.0);
+    metal::float4 ccz4_ = metal::float4(metal::float2(1.0), 0.0, 1.0);
+    metal::float4 ccz5_ = metal::float4(metal::float2(1.0, 0.0), 0.0, 1.0);
+    metal::float2x2 ccm2_ = metal::float2x2(metal::float2 {}, metal::float2(1.0, 2.0));
+    metal::float2x2 ccm3_ = metal::float2x2(metal::float2(1.0), metal::float2 {});
     metal::uint2 ic4_ = metal::uint2(0u, 0u);
     metal::float2x3 ic5_ = metal::float2x3(metal::float3(0.0, 0.0, 0.0), metal::float3(0.0, 0.0, 0.0));
     return;

naga/tests/out/spv/wgsl-constructors.spvasm

@@ -1,12 +1,12 @@
 ; SPIR-V
 ; Version: 1.1
 ; Generator: rspirv
-; Bound: 70
+; Bound: 78
 OpCapability Shader
 %1 = OpExtInstImport "GLSL.std.450"
 OpMemoryModel Logical GLSL450
-OpEntryPoint GLCompute %41 "main"
-OpExecutionMode %41 LocalSize 1 1 1
+OpEntryPoint GLCompute %50 "main"
+OpExecutionMode %50 LocalSize 1 1 1
 OpMemberDecorate %6 0 Offset 0
 OpMemberDecorate %6 1 Offset 16
 OpDecorate %10 ArrayStride 16
@@ -50,37 +50,45 @@ OpDecorate %18 ArrayStride 4
 %37 = OpConstantNull  %6
 %38 = OpConstant  %12  0
 %39 = OpConstantComposite  %14  %38 %38
-%42 = OpTypeFunction %2
-%43 = OpConstantComposite  %4  %23 %23 %23 %23
-%44 = OpConstant  %5  1
-%45 = OpConstantComposite  %6  %43 %44
-%46 = OpConstantComposite  %9  %23 %21
-%47 = OpConstantComposite  %8  %46 %26
-%48 = OpConstantComposite  %4  %23 %21 %21 %21
-%49 = OpConstantComposite  %4  %21 %23 %21 %21
-%50 = OpConstantComposite  %4  %21 %21 %23 %21
-%51 = OpConstantComposite  %4  %21 %21 %21 %23
-%52 = OpConstantComposite  %17  %48 %49 %50 %51
-%53 = OpConstantComposite  %14  %38 %38
-%54 = OpConstantComposite  %9  %21 %21
-%55 = OpConstantComposite  %9  %21 %21
-%56 = OpConstantComposite  %8  %55 %55
-%57 = OpConstant  %5  0
-%58 = OpConstant  %5  2
-%59 = OpConstant  %5  3
-%60 = OpConstantComposite  %18  %57 %44 %58 %59
-%61 = OpConstantFalse  %13
-%62 = OpConstantComposite  %7  %21 %21 %21
-%63 = OpConstantComposite  %20  %62 %62
-%64 = OpConstantNull  %20
-%65 = OpConstantTrue  %13
-%67 = OpTypePointer Function %6
-%68 = OpConstantNull  %6
-%41 = OpFunction  %2  None %42
-%40 = OpLabel
-%66 = OpVariable  %67  Function %68
-OpBranch %69
-%69 = OpLabel
-OpStore %66 %45
+%40 = OpConstantNull  %9
+%41 = OpConstantComposite  %4  %33 %33 %23 %23
+%42 = OpConstantComposite  %9  %23 %23
+%43 = OpConstantComposite  %4  %23 %23 %21 %23
+%44 = OpConstantComposite  %9  %23 %21
+%45 = OpConstantComposite  %4  %23 %21 %21 %23
+%46 = OpConstantComposite  %9  %23 %24
+%47 = OpConstantComposite  %8  %40 %46
+%48 = OpConstantComposite  %8  %42 %40
+%51 = OpTypeFunction %2
+%52 = OpConstantComposite  %4  %23 %23 %23 %23
+%53 = OpConstant  %5  1
+%54 = OpConstantComposite  %6  %52 %53
+%55 = OpConstantComposite  %8  %44 %26
+%56 = OpConstantComposite  %4  %23 %21 %21 %21
+%57 = OpConstantComposite  %4  %21 %23 %21 %21
+%58 = OpConstantComposite  %4  %21 %21 %23 %21
+%59 = OpConstantComposite  %4  %21 %21 %21 %23
+%60 = OpConstantComposite  %17  %56 %57 %58 %59
+%61 = OpConstantComposite  %14  %38 %38
+%62 = OpConstantComposite  %9  %21 %21
+%63 = OpConstantComposite  %9  %21 %21
+%64 = OpConstantComposite  %8  %63 %63
+%65 = OpConstant  %5  0
+%66 = OpConstant  %5  2
+%67 = OpConstant  %5  3
+%68 = OpConstantComposite  %18  %65 %53 %66 %67
+%69 = OpConstantFalse  %13
+%70 = OpConstantComposite  %7  %21 %21 %21
+%71 = OpConstantComposite  %20  %70 %70
+%72 = OpConstantNull  %20
+%73 = OpConstantTrue  %13
+%75 = OpTypePointer Function %6
+%76 = OpConstantNull  %6
+%50 = OpFunction  %2  None %51
+%49 = OpLabel
+%74 = OpVariable  %75  Function %76
+OpBranch %77
+%77 = OpLabel
+OpStore %74 %54
 OpReturn
 OpFunctionEnd

naga/tests/out/spv/wgsl-constructors.spvasm.glsl

@@ -9,7 +9,7 @@ struct _6
 
 void main()
 {
-    _6 _66 = _6(vec4(0.0), 0);
-    _66 = _6(vec4(1.0), 1);
+    _6 _74 = _6(vec4(0.0), 0);
+    _74 = _6(vec4(1.0), 1);
 }
 

naga/tests/out/wgsl/wgsl-constructors.wgsl

@@ -15,6 +15,11 @@ const cz5_: mat2x2<f32> = mat2x2<f32>();
 const cz6_: array<Foo, 3> = array<Foo, 3>();
 const cz7_: Foo = Foo();
 const cp1_: vec2<u32> = vec2(0u);
+const ccz0_: vec4<f32> = vec4<f32>(vec2<f32>(), 1f, 1f);
+const ccz1_: vec4<f32> = vec4<f32>(vec2(1f), 0f, 1f);
+const ccz2_: vec4<f32> = vec4<f32>(vec2<f32>(1f, 0f), 0f, 1f);
+const ccm0_: mat2x2<f32> = mat2x2<f32>(vec2<f32>(), vec2<f32>(1f, 2f));
+const ccm1_: mat2x2<f32> = mat2x2<f32>(vec2(1f), vec2<f32>());
 
 @compute @workgroup_size(1, 1, 1) 
 fn main() {
@@ -28,6 +33,11 @@ fn main() {
     let cit0_ = vec2(0u);
     let cit1_ = mat2x2<f32>(vec2(0f), vec2(0f));
     let cit2_ = array<i32, 4>(0i, 1i, 2i, 3i);
+    let ccz3_ = vec4<f32>(vec2<f32>(), 1f, 1f);
+    let ccz4_ = vec4<f32>(vec2(1f), 0f, 1f);
+    let ccz5_ = vec4<f32>(vec2<f32>(1f, 0f), 0f, 1f);
+    let ccm2_ = mat2x2<f32>(vec2<f32>(), vec2<f32>(1f, 2f));
+    let ccm3_ = mat2x2<f32>(vec2(1f), vec2<f32>());
     let ic4_ = vec2<u32>(0u, 0u);
     let ic5_ = mat2x3<f32>(vec3<f32>(0f, 0f, 0f), vec3<f32>(0f, 0f, 0f));
     return;