[naga spv-out] Don't generate multiple temporaries for spill values. (#7239)

When spilling arrays and matrices so that we can access them with an
index computed at runtime, if we need to spill the same expression again,
don't wipe out our record of the temporary variable we used the first
time. Just re-use it.

Fixes #7048.

naga/src/back/spv/block.rs

@@ -2230,26 +2230,60 @@ impl BlockContext<'_> {
     }
 
     fn spill_to_internal_variable(&mut self, base: Handle<crate::Expression>, block: &mut Block) {
-        // Generate an internal variable of the appropriate type for `base`.
-        let variable_id = self.writer.id_gen.next();
-        let pointer_type_id = self
-            .writer
-            .get_resolution_pointer_id(&self.fun_info[base].ty, spirv::StorageClass::Function);
-        let variable = super::LocalVariable {
-            id: variable_id,
-            instruction: Instruction::variable(
-                pointer_type_id,
-                variable_id,
-                spirv::StorageClass::Function,
-                None,
-            ),
+        use indexmap::map::Entry;
+
+        // Make sure we have an internal variable to spill `base` to.
+        let spill_variable_id = match self.function.spilled_composites.entry(base) {
+            Entry::Occupied(preexisting) => preexisting.get().id,
+            Entry::Vacant(vacant) => {
+                // Generate a new internal variable of the appropriate
+                // type for `base`.
+                let pointer_type_id = self.writer.get_resolution_pointer_id(
+                    &self.fun_info[base].ty,
+                    spirv::StorageClass::Function,
+                );
+                let id = self.writer.id_gen.next();
+                vacant.insert(super::LocalVariable {
+                    id,
+                    instruction: Instruction::variable(
+                        pointer_type_id,
+                        id,
+                        spirv::StorageClass::Function,
+                        None,
+                    ),
+                });
+                id
+            }
         };
 
+        // Perform the store even if we already had a spill variable for `base`.
+        // Consider this code:
+        //
+        // var x = ...;
+        // var y = ...;
+        // var z = ...;
+        // for (i = 0; i<2; i++) {
+        //     let a = array(i, i, i);
+        //     if (i == 0) {
+        //         x += a[y];
+        //     } else [
+        //         x += a[z];
+        //     }
+        // }
+        //
+        // The value of `a` needs to be spilled so we can subscript it with `y` and `z`.
+        //
+        // When we generate SPIR-V for `a[y]`, we will create the spill
+        // variable, and store `a`'s value in it.
+        //
+        // When we generate SPIR-V for `a[z]`, we will notice that the spill
+        // variable for `a` has already been declared, but it is still essential
+        // that we store `a` into it, so that `a[z]` sees this iteration's value
+        // of `a`.
         let base_id = self.cached[base];
         block
             .body
-            .push(Instruction::store(variable.id, base_id, None));
-        self.function.spilled_composites.insert(base, variable);
+            .push(Instruction::store(spill_variable_id, base_id, None));
     }
 
     /// Generate an access to a spilled temporary, if necessary.

naga/src/back/spv/mod.rs

@@ -147,12 +147,17 @@ struct Function {
     /// List of local variables used as a counters to ensure that all loops are bounded.
     force_loop_bounding_vars: Vec<LocalVariable>,
 
-    /// A map taking an expression that yields a composite value (array, matrix)
-    /// to the temporary variables we have spilled it to, if any. Spilling
-    /// allows us to render an arbitrary chain of [`Access`] and [`AccessIndex`]
-    /// expressions as an `OpAccessChain` and an `OpLoad` (plus bounds checks).
-    /// This supports dynamic indexing of by-value arrays and matrices, which
-    /// SPIR-V does not.
+    /// A map from a Naga expression to the temporary SPIR-V variable we have
+    /// spilled its value to, if any.
+    ///
+    /// Naga IR lets us apply [`Access`] expressions to expressions whose value
+    /// is an array or matrix---not a pointer to such---but SPIR-V doesn't have
+    /// instructions that can do the same. So when we encounter such code, we
+    /// spill the expression's value to a generated temporary variable. That, we
+    /// can obtain a pointer to, and then use an `OpAccessChain` instruction to
+    /// do whatever series of [`Access`] and [`AccessIndex`] operations we need
+    /// (with bounds checks). Finally, we generate an `OpLoad` to get the final
+    /// value.
     ///
     /// [`Access`]: crate::Expression::Access
     /// [`AccessIndex`]: crate::Expression::AccessIndex

naga/tests/in/7048-multiple-dynamic-1.toml

@@ -0,0 +1 @@
+targets = "SPIRV"

naga/tests/in/7048-multiple-dynamic-1.wgsl

@@ -0,0 +1,8 @@
+fn f() {
+    let b = array<vec3f, 2>();
+    var poly = vec4f(0);
+    var k = 0;
+    var j = 0;
+
+    poly.x += b[j].y * b[k].z;
+}

naga/tests/in/7048-multiple-dynamic-2.toml

@@ -0,0 +1 @@
+targets = "SPIRV"

naga/tests/in/7048-multiple-dynamic-2.wgsl

@@ -0,0 +1,13 @@
+@fragment
+fn fs_main() -> @location(0) vec4f {
+    let my_array = array(
+        vec2f(0.0, 0.0),
+        vec2f(0.0, 0.0),
+    );
+    var index_0 = 0;
+
+    let val_0 = my_array[index_0];
+    let val_1 = my_array[index_0];
+
+    return (val_0 * val_1).xxyy;
+}

naga/tests/in/7048-multiple-dynamic-3.toml

@@ -0,0 +1 @@
+targets = "SPIRV"

naga/tests/in/7048-multiple-dynamic-3.wgsl

@@ -0,0 +1,16 @@
+struct QEFResult {
+        a: f32,
+        b: vec3<f32>,
+}
+
+fn foobar(normals: array<vec3f, 12>, count: u32) -> QEFResult {
+        for (var i = 0u; i < count; i++) {
+                var n0 = normals[i];
+        }
+
+        for (var j = 0u; j < count; j++) {
+                var n1 = normals[j];
+        }
+
+        return QEFResult(0.0, vec3(0.0));
+}

naga/tests/out/spv/7048-multiple-dynamic-1.spvasm

@@ -0,0 +1,52 @@
+; SPIR-V
+; Version: 1.1
+; Generator: rspirv
+; Bound: 39
+OpCapability Shader
+OpCapability Linkage
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpDecorate %5 ArrayStride 16
+%2 = OpTypeVoid
+%4 = OpTypeFloat 32
+%3 = OpTypeVector %4 3
+%7 = OpTypeInt 32 0
+%6 = OpConstant  %7  2
+%5 = OpTypeArray %3 %6
+%8 = OpTypeVector %4 4
+%9 = OpTypeInt 32 1
+%12 = OpTypeFunction %2
+%13 = OpConstantNull  %5
+%14 = OpConstant  %4  0.0
+%15 = OpConstantComposite  %8  %14 %14 %14 %14
+%16 = OpConstant  %9  0
+%18 = OpTypePointer Function %8
+%20 = OpTypePointer Function %9
+%23 = OpTypePointer Function %4
+%25 = OpTypePointer Function %5
+%27 = OpConstant  %7  1
+%34 = OpConstant  %7  0
+%11 = OpFunction  %2  None %12
+%10 = OpLabel
+%17 = OpVariable  %18  Function %15
+%19 = OpVariable  %20  Function %16
+%21 = OpVariable  %20  Function %16
+%26 = OpVariable  %25  Function
+OpBranch %22
+%22 = OpLabel
+%24 = OpLoad  %9  %21
+OpStore %26 %13
+%28 = OpAccessChain  %23  %26 %24 %27
+%29 = OpLoad  %4  %28
+%30 = OpLoad  %9  %19
+OpStore %26 %13
+%31 = OpAccessChain  %23  %26 %30 %6
+%32 = OpLoad  %4  %31
+%33 = OpFMul  %4  %29 %32
+%35 = OpAccessChain  %23  %17 %34
+%36 = OpLoad  %4  %35
+%37 = OpFAdd  %4  %36 %33
+%38 = OpAccessChain  %23  %17 %34
+OpStore %38 %37
+OpReturn
+OpFunctionEnd

naga/tests/out/spv/7048-multiple-dynamic-2.spvasm

@@ -0,0 +1,48 @@
+; SPIR-V
+; Version: 1.1
+; Generator: rspirv
+; Bound: 33
+OpCapability Shader
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpEntryPoint Fragment %13 "fs_main" %11
+OpExecutionMode %13 OriginUpperLeft
+OpDecorate %6 ArrayStride 8
+OpDecorate %11 Location 0
+%2 = OpTypeVoid
+%4 = OpTypeFloat 32
+%3 = OpTypeVector %4 4
+%5 = OpTypeVector %4 2
+%8 = OpTypeInt 32 0
+%7 = OpConstant  %8  2
+%6 = OpTypeArray %5 %7
+%9 = OpTypeInt 32 1
+%12 = OpTypePointer Output %3
+%11 = OpVariable  %12  Output
+%14 = OpTypeFunction %2
+%15 = OpConstant  %4  0.0
+%16 = OpConstantComposite  %5  %15 %15
+%17 = OpConstantComposite  %6  %16 %16
+%18 = OpConstant  %9  0
+%20 = OpTypePointer Function %9
+%23 = OpTypePointer Function %6
+%25 = OpTypePointer Function %5
+%13 = OpFunction  %2  None %14
+%10 = OpLabel
+%19 = OpVariable  %20  Function %18
+%24 = OpVariable  %23  Function
+OpBranch %21
+%21 = OpLabel
+%22 = OpLoad  %9  %19
+OpStore %24 %17
+%26 = OpAccessChain  %25  %24 %22
+%27 = OpLoad  %5  %26
+%28 = OpLoad  %9  %19
+OpStore %24 %17
+%29 = OpAccessChain  %25  %24 %28
+%30 = OpLoad  %5  %29
+%31 = OpFMul  %5  %27 %30
+%32 = OpVectorShuffle  %3  %31 %31 0 0 1 1
+OpStore %11 %32
+OpReturn
+OpFunctionEnd

naga/tests/out/spv/7048-multiple-dynamic-3.spvasm

@@ -0,0 +1,99 @@
+; SPIR-V
+; Version: 1.1
+; Generator: rspirv
+; Bound: 61
+OpCapability Shader
+OpCapability Linkage
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpMemberDecorate %5 0 Offset 0
+OpMemberDecorate %5 1 Offset 16
+OpDecorate %6 ArrayStride 16
+%2 = OpTypeVoid
+%3 = OpTypeFloat 32
+%4 = OpTypeVector %3 3
+%5 = OpTypeStruct %3 %4
+%8 = OpTypeInt 32 0
+%7 = OpConstant  %8  12
+%6 = OpTypeArray %4 %7
+%13 = OpTypeFunction %5 %6 %8
+%14 = OpConstant  %8  0
+%15 = OpConstant  %8  1
+%16 = OpConstant  %3  0.0
+%17 = OpConstantComposite  %4  %16 %16 %16
+%18 = OpConstantComposite  %5  %16 %17
+%20 = OpTypePointer Function %8
+%22 = OpTypePointer Function %4
+%23 = OpConstantNull  %4
+%26 = OpConstantNull  %4
+%33 = OpTypeBool
+%40 = OpTypePointer Function %6
+%12 = OpFunction  %5  None %13
+%10 = OpFunctionParameter  %6
+%11 = OpFunctionParameter  %8
+%9 = OpLabel
+%21 = OpVariable  %22  Function %23
+%25 = OpVariable  %22  Function %26
+%19 = OpVariable  %20  Function %14
+%24 = OpVariable  %20  Function %14
+%41 = OpVariable  %40  Function
+OpBranch %27
+%27 = OpLabel
+OpBranch %28
+%28 = OpLabel
+OpLoopMerge %29 %31 None
+OpBranch %30
+%30 = OpLabel
+%32 = OpLoad  %8  %19
+%34 = OpULessThan  %33  %32 %11
+OpSelectionMerge %35 None
+OpBranchConditional %34 %35 %36
+%36 = OpLabel
+OpBranch %29
+%35 = OpLabel
+OpBranch %37
+%37 = OpLabel
+%39 = OpLoad  %8  %19
+OpStore %41 %10
+%42 = OpAccessChain  %22  %41 %39
+%43 = OpLoad  %4  %42
+OpStore %21 %43
+OpBranch %38
+%38 = OpLabel
+OpBranch %31
+%31 = OpLabel
+%44 = OpLoad  %8  %19
+%45 = OpIAdd  %8  %44 %15
+OpStore %19 %45
+OpBranch %28
+%29 = OpLabel
+OpBranch %46
+%46 = OpLabel
+OpLoopMerge %47 %49 None
+OpBranch %48
+%48 = OpLabel
+%50 = OpLoad  %8  %24
+%51 = OpULessThan  %33  %50 %11
+OpSelectionMerge %52 None
+OpBranchConditional %51 %52 %53
+%53 = OpLabel
+OpBranch %47
+%52 = OpLabel
+OpBranch %54
+%54 = OpLabel
+%56 = OpLoad  %8  %24
+OpStore %41 %10
+%57 = OpAccessChain  %22  %41 %56
+%58 = OpLoad  %4  %57
+OpStore %25 %58
+OpBranch %55
+%55 = OpLabel
+OpBranch %49
+%49 = OpLabel
+%59 = OpLoad  %8  %24
+%60 = OpIAdd  %8  %59 %15
+OpStore %24 %60
+OpBranch %46
+%47 = OpLabel
+OpReturnValue %18
+OpFunctionEnd

naga/tests/out/spv/index-by-value.spvasm

@@ -24,7 +24,7 @@ OpDecorate %64 BuiltIn Position
 %11 = OpTypeMatrix %12 2
 %13 = OpTypeVector %10 4
 %18 = OpTypeFunction %3 %4 %3
-%21 = OpTypePointer Function %4
+%20 = OpTypePointer Function %4
 %22 = OpTypePointer Function %3
 %29 = OpTypeFunction %3 %3 %3
 %30 = OpConstant  %3  1
@@ -34,7 +34,7 @@ OpDecorate %64 BuiltIn Position
 %34 = OpConstant  %3  4
 %35 = OpConstantComposite  %7  %33 %34
 %36 = OpConstantComposite  %9  %32 %35
-%39 = OpTypePointer Function %9
+%38 = OpTypePointer Function %9
 %46 = OpTypeFunction %10 %3 %3
 %47 = OpConstant  %10  1.0
 %48 = OpConstant  %10  2.0
@@ -43,7 +43,7 @@ OpDecorate %64 BuiltIn Position
 %51 = OpConstantComposite  %12  %47 %48
 %52 = OpConstantComposite  %12  %49 %50
 %53 = OpConstantComposite  %11  %51 %52
-%56 = OpTypePointer Function %11
+%55 = OpTypePointer Function %11
 %57 = OpTypePointer Function %10
 %62 = OpTypePointer Input %6
 %61 = OpVariable  %62  Input
@@ -57,11 +57,11 @@ OpDecorate %64 BuiltIn Position
 %15 = OpFunctionParameter  %4
 %16 = OpFunctionParameter  %3
 %14 = OpLabel
-%20 = OpVariable  %21  Function
+%21 = OpVariable  %20  Function
 OpBranch %19
 %19 = OpLabel
-OpStore %20 %15
-%23 = OpAccessChain  %22  %20 %16
+OpStore %21 %15
+%23 = OpAccessChain  %22  %21 %16
 %24 = OpLoad  %3  %23
 OpReturnValue %24
 OpFunctionEnd
@@ -69,11 +69,11 @@ OpFunctionEnd
 %26 = OpFunctionParameter  %3
 %27 = OpFunctionParameter  %3
 %25 = OpLabel
-%38 = OpVariable  %39  Function
+%39 = OpVariable  %38  Function
 OpBranch %37
 %37 = OpLabel
-OpStore %38 %36
-%40 = OpAccessChain  %22  %38 %26 %27
+OpStore %39 %36
+%40 = OpAccessChain  %22  %39 %26 %27
 %41 = OpLoad  %3  %40
 OpReturnValue %41
 OpFunctionEnd
@@ -81,17 +81,17 @@ OpFunctionEnd
 %43 = OpFunctionParameter  %3
 %44 = OpFunctionParameter  %3
 %42 = OpLabel
-%55 = OpVariable  %56  Function
+%56 = OpVariable  %55  Function
 OpBranch %54
 %54 = OpLabel
-OpStore %55 %53
-%58 = OpAccessChain  %57  %55 %43 %44
+OpStore %56 %53
+%58 = OpAccessChain  %57  %56 %43 %44
 %59 = OpLoad  %10  %58
 OpReturnValue %59
 OpFunctionEnd
 %66 = OpFunction  %2  None %67
 %60 = OpLabel
-%71 = OpVariable  %21  Function
+%71 = OpVariable  %20  Function
 %63 = OpLoad  %6  %61
 OpBranch %70
 %70 = OpLabel