spv-out: don't dynamically index PushConstant vectors with OpAccessChain

Per VUID-RuntimeSpirv-None-04745, all accesses into a PushConstant-storage-class
variable that are arrays must use dynamically uniform indices; this restriction
also applies to vectors. Naga was emitting `OpAccessChain` with per-invocation
(non-uniform) indices directly into push-constant vectors, which is undefined
behavior under that rule -- observed on RADV as every invocation silently
reading element 0 of the vector.

Add a special case to `write_checked_load`, alongside the existing matCx2
ones, that loads the whole vector and uses `OpVectorExtractDynamic` instead,
which isn't subject to the restriction.

Fixes gfx-rs#8612.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>

Author: jimblandy

naga/src/back/spv/block.rs

@@ -429,6 +429,93 @@ impl BlockContext<'_> {
         block
     }
 
+    /// If `pointer` refers to a scalar reached by a dynamic (non-constant)
+    /// index into a vector in the [`Immediate`] (push-constant) address
+    /// space, write code to access the value, returning the ID of the
+    /// result. Else return `None`.
+    ///
+    /// `VUID-RuntimeSpirv-None-04745` requires:
+    ///
+    /// > All block members in a variable with a Storage Class of PushConstant
+    /// > declared as an array must only be accessed by dynamically uniform
+    /// > indices
+    ///
+    /// According to [upstream discussion], this requirement was intended to
+    /// allow drivers to always compile loads from push constants as scalar
+    /// loads; the omission of vectors seems to be an oversight.
+    ///
+    /// Naga IR, however, has no such restriction on indexing vectors in the
+    /// `Immediate` address space. So Naga must not emit `OpAccessChain` with
+    /// such an index directly into a `PushConstant` vector.
+    ///
+    /// Instead, this loads the whole vector -- a plain `OpLoad`, which isn't
+    /// subject to the restriction above -- and then extracts the desired
+    /// component from that loaded value with `OpVectorExtractDynamic`, exactly
+    /// as [`Self::write_vector_access()`] already does for by-value vectors.
+    ///
+    /// [`Immediate`]: crate::AddressSpace::Immediate
+    /// [upstream discussion]: https://gitlab.freedesktop.org/mesa/mesa/-/work_items/15891#note_3573369
+    fn maybe_write_immediate_vector_dynamic_access(
+        &mut self,
+        pointer: Handle<crate::Expression>,
+        block: &mut Block,
+    ) -> Result<Option<Word>, Error> {
+        // We're only interested in a scalar reached by indexing into a
+        // vector with a computed (not compile-time-constant) index.
+        let crate::Expression::Access {
+            base: vector_pointer,
+            index,
+        } = self.ir_function.expressions[pointer]
+        else {
+            return Ok(None);
+        };
+
+        // If the index is actually a compile-time constant, the plain
+        // access-chain path is fine: constants are as uniform as can be.
+        if let GuardedIndex::Known(_) =
+            GuardedIndex::from_expression(index, &self.ir_function.expressions, self.ir_module)
+        {
+            return Ok(None);
+        }
+
+        // Ensure `vector_pointer` is a pointer to a vector in the Immediate
+        // (push-constant) address space.
+        let vector_pointer_ty = self.fun_info[vector_pointer]
+            .ty
+            .inner_with(&self.ir_module.types);
+        if vector_pointer_ty.pointer_space() != Some(crate::AddressSpace::Immediate) {
+            return Ok(None);
+        }
+        let Some(vector_base_ty) = vector_pointer_ty.pointer_base_type() else {
+            return Ok(None);
+        };
+        let crate::TypeInner::Vector { size, scalar } =
+            *vector_base_ty.inner_with(&self.ir_module.types)
+        else {
+            return Ok(None);
+        };
+
+        let vector_type_id = self.get_numeric_type_id(NumericType::Vector { size, scalar });
+        let component_type_id = self.get_numeric_type_id(NumericType::Scalar(scalar));
+
+        let vector_load_id = self.write_checked_load(
+            vector_pointer,
+            block,
+            AccessTypeAdjustment::None,
+            vector_type_id,
+        )?;
+
+        let result_id = self.write_vector_access(
+            component_type_id,
+            vector_pointer,
+            Some(vector_load_id),
+            GuardedIndex::Expression(index),
+            block,
+        )?;
+
+        Ok(Some(result_id))
+    }
+
     /// If `pointer` refers to an access chain that contains a dynamic indexing
     /// of a two-row matrix in the [`Uniform`] address space, write code to
     /// access the value returning the ID of the result. Else return None.
@@ -2773,7 +2860,11 @@ impl BlockContext<'_> {
         access_type_adjustment: AccessTypeAdjustment,
         result_type_id: Word,
     ) -> Result<Word, Error> {
-        if let Some(result_id) = self.maybe_write_uniform_matcx2_dynamic_access(pointer, block)? {
+        if let Some(result_id) = self.maybe_write_immediate_vector_dynamic_access(pointer, block)? {
+            Ok(result_id)
+        } else if let Some(result_id) =
+            self.maybe_write_uniform_matcx2_dynamic_access(pointer, block)?
+        {
             Ok(result_id)
         } else if let Some(result_id) =
             self.maybe_write_load_uniform_matcx2_struct_member(pointer, block)?

naga/tests/in/wgsl/push-constant-dynamic-vector-index.toml

@@ -0,0 +1,5 @@
+capabilities = "IMMEDIATES"
+targets = "SPIRV"
+
+[bounds_check_policies]
+index = "Restrict"

naga/tests/in/wgsl/push-constant-dynamic-vector-index.wgsl

@@ -0,0 +1,43 @@
+// Regression test for <https://github.com/gfx-rs/wgpu/issues/8612>.
+//
+// The Vulkan environment spec (`VUID-RuntimeSpirv-None-04745`) requires
+// that accesses into a `PushConstant`-storage-class variable that are
+// arrays use dynamically uniform indices; this restriction also applies
+// to vectors. So naga must not emit `OpAccessChain` with a non-constant
+// index directly into a vector living in the `immediate` (push-constant)
+// address space -- it must instead load the whole vector and use
+// `OpVectorExtractDynamic`.
+//
+// Naming (mirrors `mat_cx2.wgsl`):
+// V = vector field, M = matrix field, C = constant index, trailing
+// C/V = constant/variable (dynamic) index into the vector/column.
+
+struct Immediates {
+    v: vec4<f32>,
+    m: mat4x4<f32>,
+}
+
+var<immediate> im: Immediates;
+
+@group(0) @binding(0)
+var<storage, read_write> out: array<f32>;
+
+@compute @workgroup_size(1)
+fn main(@builtin(local_invocation_index) idx: u32) {
+    // Dynamically indexing a vector field directly. `im.v` is a
+    // `TypeInner::Pointer` to a vector (it has a concrete entry in the
+    // module's type arena, being a struct field).
+    let v_c = im.v[0];
+    let v_v = im.v[idx];
+
+    // Dynamically indexing a component of a matrix column. `im.m[0]` is a
+    // `TypeInner::ValuePointer` to a vector (it doesn't have its own type
+    // arena entry), which is a distinct code path from `v_v` above.
+    let m_cc = im.m[0][0];
+    let m_cv = im.m[0][idx];
+
+    out[0] = v_c;
+    out[1] = v_v;
+    out[2] = m_cc;
+    out[3] = m_cv;
+}

naga/tests/out/spv/wgsl-push-constant-dynamic-vector-index.spvasm

@@ -0,0 +1,77 @@
+; SPIR-V
+; Version: 1.1
+; Generator: rspirv
+; Bound: 50
+OpCapability Shader
+OpExtension "SPV_KHR_storage_buffer_storage_class"
+%1 = OpExtInstImport "GLSL.std.450"
+OpMemoryModel Logical GLSL450
+OpEntryPoint GLCompute %19 "main" %16
+OpExecutionMode %19 LocalSize 1 1 1
+OpMemberDecorate %6 0 Offset 0
+OpMemberDecorate %6 1 Offset 16
+OpMemberDecorate %6 1 ColMajor
+OpMemberDecorate %6 1 MatrixStride 16
+OpDecorate %7 ArrayStride 4
+OpDecorate %10 Block
+OpMemberDecorate %10 0 Offset 0
+OpDecorate %12 DescriptorSet 0
+OpDecorate %12 Binding 0
+OpDecorate %13 Block
+OpMemberDecorate %13 0 Offset 0
+OpDecorate %16 BuiltIn LocalInvocationIndex
+%2 = OpTypeVoid
+%3 = OpTypeFloat 32
+%4 = OpTypeVector %3 4
+%5 = OpTypeMatrix %4 4
+%6 = OpTypeStruct %4 %5
+%7 = OpTypeRuntimeArray %3
+%8 = OpTypeInt 32 0
+%10 = OpTypeStruct %6
+%11 = OpTypePointer PushConstant %10
+%9 = OpVariable  %11  PushConstant
+%13 = OpTypeStruct %7
+%14 = OpTypePointer StorageBuffer %13
+%12 = OpVariable  %14  StorageBuffer
+%17 = OpTypePointer Input %8
+%16 = OpVariable  %17  Input
+%20 = OpTypeFunction %2
+%21 = OpTypePointer PushConstant %6
+%22 = OpConstant  %8  0
+%24 = OpTypePointer StorageBuffer %7
+%27 = OpTypePointer PushConstant %4
+%28 = OpTypePointer PushConstant %3
+%33 = OpConstant  %8  3
+%36 = OpTypePointer PushConstant %5
+%37 = OpConstant  %8  1
+%44 = OpTypePointer StorageBuffer %3
+%47 = OpConstant  %8  2
+%19 = OpFunction  %2  None %20
+%15 = OpLabel
+%18 = OpLoad  %8  %16
+%23 = OpAccessChain  %21  %9 %22
+%25 = OpAccessChain  %24  %12 %22
+OpBranch %26
+%26 = OpLabel
+%29 = OpAccessChain  %28  %23 %22 %22
+%30 = OpLoad  %3  %29
+%31 = OpAccessChain  %27  %23 %22
+%32 = OpLoad  %4  %31
+%34 = OpExtInst  %8  %1 UMin %18 %33
+%35 = OpVectorExtractDynamic  %3  %32 %34
+%38 = OpAccessChain  %28  %23 %37 %22 %22
+%39 = OpLoad  %3  %38
+%40 = OpAccessChain  %27  %23 %37 %22
+%41 = OpLoad  %4  %40
+%42 = OpExtInst  %8  %1 UMin %18 %33
+%43 = OpVectorExtractDynamic  %3  %41 %42
+%45 = OpAccessChain  %44  %25 %22
+OpStore %45 %30
+%46 = OpAccessChain  %44  %25 %37
+OpStore %46 %35
+%48 = OpAccessChain  %44  %25 %47
+OpStore %48 %39
+%49 = OpAccessChain  %44  %25 %33
+OpStore %49 %43
+OpReturn
+OpFunctionEnd

naga/tests/out/spv/wgsl-push-constant-dynamic-vector-index.spvasm.glsl

@@ -0,0 +1,27 @@
+#version 460
+layout(local_size_x = 1, local_size_y = 1, local_size_z = 1) in;
+
+struct _6
+{
+    vec4 _m0;
+    mat4 _m1;
+};
+
+layout(set = 0, binding = 0, std430) buffer _13_12
+{
+    float _m0[];
+} _12;
+
+layout(push_constant, std430) uniform _10_9
+{
+    _6 _m0;
+} _9;
+
+void main()
+{
+    _12._m0[0u] = _9._m0._m0.x;
+    _12._m0[1u] = _9._m0._m0[min(gl_LocalInvocationIndex, 3u)];
+    _12._m0[2u] = _9._m0._m1[0u].x;
+    _12._m0[3u] = _9._m0._m1[0u][min(gl_LocalInvocationIndex, 3u)];
+}
+