Repository: gfx-rs/wgpu (language: rust)

Issue:
Vulkan Validation Error when using multiple dynamic indexing
<!-- Thank you for filing this! Please read the [debugging tips](https://github.com/gfx-rs/wgpu/wiki/Debugging-wgpu-Applications).
That may let you investigate on your own, or provide additional information that helps us to assist.-->

**Description**
When trying to run the code
```
@fragment
fn fs_main() -> @location(0) vec4f {
    let my_array = array(
        vec2f(0.0, 0.0),
        vec2f(0.0, 0.0),
    );

    let val_0 = my_array[index_0];
    let val_1 = my_array[index_0];

    return (val_0 * val_1).xxyy;
}
```
on the Vulkan backend, a Validation Error is raised:
```
Validation Error: [ VUID-VkShaderModuleCreateInfo-pCode-08737 ] | MessageID = 0xa5625282 | vkCreateShaderModule(): pCreateInfo->pCode (spirv-val produced an error):
    ID '28[%28]' has not been defined
      OpStore %28 %25
    The Vulkan spec states: If pCode is a pointer to SPIR-V code, pCode must adhere to the validation rules described by the Validation Rules within a Module section of the SPIR-V Environment appendix (https://vulkan.lunarg.com/doc/view/1.4.304.0/windows/1.4-extensions/vkspec.html#VUID-VkShaderModuleCreateInfo-pCode-08737)
```
The same fragment shader runs fine on DirectX 12 and OpenGL.

**Repro steps**
```
git clone https://github.com/lorinjacot/wgsl-indexing-bug.git
cd wgsl-indexing-bug
cargo run
```

**Expected vs observed behavior**
I expect the program to run and exit without any error. Currently, the program exits with the error `(exit code: 0xc0000005, STATUS_ACCESS_VIOLATION)` on the Vulkan backend.

**Platform**
OS: Windows 11 Home, version 24H2
`wgpu` version: 24.0.1
GPU: Integrated, AMD Ryzen 5 5500U with Radeon Graphics



Discussion / hints:
I hit this bug in a pretty simple use case. Using wgpu 24.0 and vulkan 1.4.303 on Linux.


```
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));
}
```

Was easy to work around by assigning `normals` to a `var` and using that instead.
Promoting to P1, since https://bugzilla.mozilla.org/show_bug.cgi?id=1930012 is P1.
I'm guessing this is a regression from #6390.
SPIR-V doesn't have instructions for accessing an array or matrix element with an index computed at runtime unless the array or matrix is stored in memory. In other words, it can apply an index to a pointer, but not to an array or matrix that has been computed by value - as with the `let` expression in the first test case, or the function argument in the second. (Oddly, SPIR-V does have `OpVectorExtractDynamic` for indexing vectors by value dynamically.)

However, WGSL and Naga IR certainly do permit this. So when the SPIR-V backend encounters an `Access` instruction applied to an array or matrix, it spills the expression's value to a generated temporary variable. Once the value is in memory, it can use the `OpAccessChain` instruction to do whatever indexing it needs. The `BlockContext::spill_to_internal_variable` function is responsible for spilling an expression's value to a temporary variable in preparation for indexing.

The bug is that `spill_to_internal_variable` doesn't check if there was already a local variable created for this expression; it always creates a new one each time. The local variables are actually accumulated in a table keyed by expression, and their declarations are only written out at the end, so all indexing operations but the last end up trying to spill the expression to a local variable that got replaced in the table before it could be emitted.

So the key ingredients for this bug are an array or matrix that is indexed with a dynamic index by value *twice*.

Fix patch:
diff --git a/naga/src/back/spv/block.rs b/naga/src/back/spv/block.rs
index 4f07dea3e2..36c433874f 100644
--- a/naga/src/back/spv/block.rs
+++ b/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.
diff --git a/naga/src/back/spv/mod.rs b/naga/src/back/spv/mod.rs
index bb2757766c..be63d3aac7 100644
--- a/naga/src/back/spv/mod.rs
+++ b/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


Test patch:
diff --git a/naga/tests/in/7048-multiple-dynamic-1.toml b/naga/tests/in/7048-multiple-dynamic-1.toml
new file mode 100644
index 0000000000..4883dbbd8d
--- /dev/null
+++ b/naga/tests/in/7048-multiple-dynamic-1.toml
@@ -0,0 +1,1 @@
+targets = "SPIRV"
diff --git a/naga/tests/in/7048-multiple-dynamic-1.wgsl b/naga/tests/in/7048-multiple-dynamic-1.wgsl
new file mode 100644
index 0000000000..2468d76b8c
--- /dev/null
+++ b/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;
+}
diff --git a/naga/tests/in/7048-multiple-dynamic-2.toml b/naga/tests/in/7048-multiple-dynamic-2.toml
new file mode 100644
index 0000000000..4883dbbd8d
--- /dev/null
+++ b/naga/tests/in/7048-multiple-dynamic-2.toml
@@ -0,0 +1,1 @@
+targets = "SPIRV"
diff --git a/naga/tests/in/7048-multiple-dynamic-2.wgsl b/naga/tests/in/7048-multiple-dynamic-2.wgsl
new file mode 100644
index 0000000000..0e92458b8f
--- /dev/null
+++ b/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;
+}
diff --git a/naga/tests/in/7048-multiple-dynamic-3.toml b/naga/tests/in/7048-multiple-dynamic-3.toml
new file mode 100644
index 0000000000..4883dbbd8d
--- /dev/null
+++ b/naga/tests/in/7048-multiple-dynamic-3.toml
@@ -0,0 +1,1 @@
+targets = "SPIRV"
diff --git a/naga/tests/in/7048-multiple-dynamic-3.wgsl b/naga/tests/in/7048-multiple-dynamic-3.wgsl
new file mode 100644
index 0000000000..90e4447784
--- /dev/null
+++ b/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));
+}
diff --git a/naga/tests/out/spv/7048-multiple-dynamic-1.spvasm b/naga/tests/out/spv/7048-multiple-dynamic-1.spvasm
new file mode 100644
index 0000000000..0e8ec75be0
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/naga/tests/out/spv/7048-multiple-dynamic-2.spvasm b/naga/tests/out/spv/7048-multiple-dynamic-2.spvasm
new file mode 100644
index 0000000000..504bd8f4eb
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/naga/tests/out/spv/7048-multiple-dynamic-3.spvasm b/naga/tests/out/spv/7048-multiple-dynamic-3.spvasm
new file mode 100644
index 0000000000..341b5093f2
--- /dev/null
+++ b/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
\ No newline at end of file
diff --git a/naga/tests/out/spv/index-by-value.spvasm b/naga/tests/out/spv/index-by-value.spvasm
index 94d609d375..a945feac10 100644
--- a/naga/tests/out/spv/index-by-value.spvasm
+++ b/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