[naga spv-out] Spill arrays and matrices for runtime indexing.

Improve handling of `Access` expressions whose base is an array or
matrix (not a pointer to such), and whose index is not known at
compile time. SPIR-V does not have instructions that can do this
directly, so spill such values to temporary variables, and perform the
accesses using `OpAccessChain` instructions applied to the
temporaries.

When performing chains of accesses like `a[i].x[j]`, do not reify
intermediate values; generate a single `OpAccessChain` for the entire
thing.

Remove special cases for arrays; the same code now handles arrays and
matrices.

Update validation to permit dynamic indexing of matrices.

For details, see the comments on the new tracking structures in
`naga:🔙:spv::Function`.

Add snapshot test `index-by-value.wgsl`.

Fixes #6358.
Fixes #4337.
Alternative to #6362.

CHANGELOG.md

@@ -113,6 +113,7 @@ By @bradwerth [#6216](https://github.com/gfx-rs/wgpu/pull/6216).
 - Accept global `var`s without explicit type. By @sagudev in [#6199](https://github.com/gfx-rs/wgpu/pull/6199).
 - Fix handling of phony statements, so they are actually emitted. By @sagudev in [#6328](https://github.com/gfx-rs/wgpu/pull/6328).
 - Added `gl_DrawID` to glsl and `DrawIndex` to spv. By @ChosenName in [#6325](https://github.com/gfx-rs/wgpu/pull/6325).
+- Matrices can now be indexed by value (#4337), and indexing arrays by value no longer causes excessive spilling (#6358). By @jimblandy in [#6390](https://github.com/gfx-rs/wgpu/pull/6390).
 
 #### General
 

naga/src/arena/handle_set.rs

@@ -83,6 +83,12 @@ impl<T> HandleSet<T> {
     }
 }
 
+impl<T> Default for HandleSet<T> {
+    fn default() -> Self {
+        Self::new()
+    }
+}
+
 pub trait ArenaType<T> {
     fn len(&self) -> usize;
 }

naga/src/back/spv/block.rs

@@ -6,11 +6,7 @@ use super::{
     index::BoundsCheckResult, selection::Selection, Block, BlockContext, Dimension, Error,
     Instruction, LocalType, LookupType, NumericType, ResultMember, Writer, WriterFlags,
 };
-use crate::{
+use crate::{arena::Handle, proc::index::GuardedIndex, Statement};
-    arena::Handle,
-    proc::{index::GuardedIndex, TypeResolution},
-    Statement,
-};
 use spirv::Word;
 
 fn get_dimension(type_inner: &crate::TypeInner) -> Dimension {
@@ -345,34 +341,64 @@ impl<'w> BlockContext<'w> {
                         // that actually dereferences the pointer.
                         0
                     }
+                    _ if self.function.spilled_accesses.contains(base) => {
+                        // As far as Naga IR is concerned, this expression does not yield
+                        // a pointer (we just checked, above), but this backend spilled it
+                        // to a temporary variable, so SPIR-V thinks we're accessing it
+                        // via a pointer.
+
+                        // Since the base expression was spilled, mark this access to it
+                        // as spilled, too.
+                        self.function.spilled_accesses.insert(expr_handle);
+                        self.maybe_access_spilled_composite(expr_handle, block, result_type_id)?
+                    }
                     crate::TypeInner::Vector { .. } => {
                         self.write_vector_access(expr_handle, base, index, block)?
                     }
-                    crate::TypeInner::Array {
+                    crate::TypeInner::Array { .. } | crate::TypeInner::Matrix { .. } => {
-                        base: ty_element, ..
+                        // See if `index` is known at compile time.
-                    } => {
+                        match GuardedIndex::from_expression(index, self.ir_function, self.ir_module)
-                        let index_id = self.cached[index];
+                        {
+                            GuardedIndex::Known(value) => {
+                                // If `index` is known and in bounds, we can just use
+                                // `OpCompositeExtract`.
+                                //
+                                // At the moment, validation rejects programs if this
+                                // index is out of bounds, so we don't need bounds checks.
+                                // However, that rejection is incorrect, since WGSL says
+                                // that `let` bindings are not constant expressions
+                                // (#6396). So eventually we will need to emulate bounds
+                                // checks here.
+                                let id = self.gen_id();
                                 let base_id = self.cached[base];
-                        let base_ty = match self.fun_info[base].ty {
+                                block.body.push(Instruction::composite_extract(
-                            TypeResolution::Handle(handle) => handle,
-                            TypeResolution::Value(_) => {
-                                return Err(Error::Validation(
-                                    "Array types should always be in the arena",
-                                ))
-                            }
-                        };
-                        let (id, variable) = self.writer.promote_access_expression_to_variable(
                                     result_type_id,
+                                    id,
                                     base_id,
-                            base_ty,
+                                    &[value],
-                            index_id,
+                                ));
-                            ty_element,
-                            block,
-                        )?;
-                        self.function.internal_variables.push(variable);
                                 id
                             }
-                    // wgpu#4337: Support `crate::TypeInner::Matrix`
+                            GuardedIndex::Expression(_) => {
+                                // We are subscripting an array or matrix that is not
+                                // behind a pointer, using an index computed at runtime.
+                                // SPIR-V has no instructions that do this, so the best we
+                                // can do is spill the value to a new temporary variable,
+                                // at which point we can get a pointer to that and just
+                                // use `OpAccessChain` in the usual way.
+                                self.spill_to_internal_variable(base, block);
+
+                                // Since the base was spilled, mark this access to it as
+                                // spilled, too.
+                                self.function.spilled_accesses.insert(expr_handle);
+                                self.maybe_access_spilled_composite(
+                                    expr_handle,
+                                    block,
+                                    result_type_id,
+                                )?
+                            }
+                        }
+                    }
                     crate::TypeInner::BindingArray {
                         base: binding_type, ..
                     } => {
@@ -435,6 +461,17 @@ impl<'w> BlockContext<'w> {
                         // that actually dereferences the pointer.
                         0
                     }
+                    _ if self.function.spilled_accesses.contains(base) => {
+                        // As far as Naga IR is concerned, this expression does not yield
+                        // a pointer (we just checked, above), but this backend spilled it
+                        // to a temporary variable, so SPIR-V thinks we're accessing it
+                        // via a pointer.
+
+                        // Since the base expression was spilled, mark this access to it
+                        // as spilled, too.
+                        self.function.spilled_accesses.insert(expr_handle);
+                        self.maybe_access_spilled_composite(expr_handle, block, result_type_id)?
+                    }
                     crate::TypeInner::Vector { .. }
                     | crate::TypeInner::Matrix { .. }
                     | crate::TypeInner::Array { .. }
@@ -1737,6 +1774,14 @@ impl<'w> BlockContext<'w> {
 
         self.temp_list.clear();
         let root_id = loop {
+            // If `expr_handle` was spilled, then the temporary variable has exactly
+            // the value we want to start from.
+            if let Some(spilled) = self.function.spilled_composites.get(&expr_handle) {
+                // The root id of the `OpAccessChain` instruction is the temporary
+                // variable we spilled the composite to.
+                break spilled.id;
+            }
+
             expr_handle = match self.ir_function.expressions[expr_handle] {
                 crate::Expression::Access { base, index } => {
                     is_non_uniform_binding_array |=
@@ -1989,6 +2034,71 @@ impl<'w> BlockContext<'w> {
         }
     }
 
+    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,
+            ),
+        };
+
+        let base_id = self.cached[base];
+        block
+            .body
+            .push(Instruction::store(variable.id, base_id, None));
+        self.function.spilled_composites.insert(base, variable);
+    }
+
+    /// Generate an access to a spilled temporary, if necessary.
+    ///
+    /// Given `access`, an [`Access`] or [`AccessIndex`] expression that refers
+    /// to a component of a composite value that has been spilled to a temporary
+    /// variable, determine whether other expressions are going to use
+    /// `access`'s value:
+    ///
+    /// - If so, perform the access and cache that as the value of `access`.
+    ///
+    /// - Otherwise, generate no code and cache no value for `access`.
+    ///
+    /// Return `Ok(0)` if no value was fetched, or `Ok(id)` if we loaded it into
+    /// the instruction given by `id`.
+    ///
+    /// [`Access`]: crate::Expression::Access
+    /// [`AccessIndex`]: crate::Expression::AccessIndex
+    fn maybe_access_spilled_composite(
+        &mut self,
+        access: Handle<crate::Expression>,
+        block: &mut Block,
+        result_type_id: Word,
+    ) -> Result<Word, Error> {
+        let access_uses = self.function.access_uses.get(&access).map_or(0, |r| *r);
+        if access_uses == self.fun_info[access].ref_count {
+            // This expression is only used by other `Access` and
+            // `AccessIndex` expressions, so we don't need to cache a
+            // value for it yet.
+            Ok(0)
+        } else {
+            // There are other expressions that are going to expect this
+            // expression's value to be cached, not just other `Access` or
+            // `AccessIndex` expressions. We must actually perform the
+            // access on the spill variable now.
+            self.write_checked_load(
+                access,
+                block,
+                AccessTypeAdjustment::IntroducePointer(spirv::StorageClass::Function),
+                result_type_id,
+            )
+        }
+    }
+
     /// Build the instructions for matrix - matrix column operations
     #[allow(clippy::too_many_arguments)]
     fn write_matrix_matrix_column_op(

naga/src/back/spv/mod.rs

@@ -144,7 +144,41 @@ struct Function {
     signature: Option<Instruction>,
     parameters: Vec<FunctionArgument>,
     variables: crate::FastHashMap<Handle<crate::LocalVariable>, LocalVariable>,
-    internal_variables: 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.
+    ///
+    /// [`Access`]: crate::Expression::Access
+    /// [`AccessIndex`]: crate::Expression::AccessIndex
+    spilled_composites: crate::FastIndexMap<Handle<crate::Expression>, LocalVariable>,
+
+    /// A set of expressions that are either in [`spilled_composites`] or refer
+    /// to some component/element of such.
+    ///
+    /// [`spilled_composites`]: Function::spilled_composites
+    spilled_accesses: crate::arena::HandleSet<crate::Expression>,
+
+    /// A map taking each expression to the number of [`Access`] and
+    /// [`AccessIndex`] expressions that uses it as a base value. If an
+    /// expression has no entry, its count is zero: it is never used as a
+    /// [`Access`] or [`AccessIndex`] base.
+    ///
+    /// We use this, together with [`ExpressionInfo::ref_count`], to recognize
+    /// the tips of chains of [`Access`] and [`AccessIndex`] expressions that
+    /// access spilled values --- expressions in [`spilled_composites`]. We
+    /// defer generating code for the chain until we reach its tip, so we can
+    /// handle it with a single instruction.
+    ///
+    /// [`Access`]: crate::Expression::Access
+    /// [`AccessIndex`]: crate::Expression::AccessIndex
+    /// [`ExpressionInfo::ref_count`]: crate::valid::ExpressionInfo
+    /// [`spilled_composites`]: Function::spilled_composites
+    access_uses: crate::FastHashMap<Handle<crate::Expression>, usize>,
+
     blocks: Vec<TerminatedBlock>,
     entry_point_context: Option<EntryPointContext>,
 }

naga/src/back/spv/writer.rs

@@ -32,7 +32,7 @@ impl Function {
                 for local_var in self.variables.values() {
                     local_var.instruction.to_words(sink);
                 }
-                for internal_var in self.internal_variables.iter() {
+                for internal_var in self.spilled_composites.values() {
                     internal_var.instruction.to_words(sink);
                 }
             }
@@ -138,54 +138,6 @@ impl Writer {
         self.capabilities_used.insert(spirv::Capability::Shader);
     }
 
-    #[allow(clippy::too_many_arguments)]
-    pub(super) fn promote_access_expression_to_variable(
-        &mut self,
-        result_type_id: Word,
-        container_id: Word,
-        container_ty: Handle<crate::Type>,
-        index_id: Word,
-        element_ty: Handle<crate::Type>,
-        block: &mut Block,
-    ) -> Result<(Word, LocalVariable), Error> {
-        let pointer_type_id = self.get_pointer_id(container_ty, spirv::StorageClass::Function);
-
-        let variable = {
-            let id = self.id_gen.next();
-            LocalVariable {
-                id,
-                instruction: Instruction::variable(
-                    pointer_type_id,
-                    id,
-                    spirv::StorageClass::Function,
-                    None,
-                ),
-            }
-        };
-        block
-            .body
-            .push(Instruction::store(variable.id, container_id, None));
-
-        let element_pointer_id = self.id_gen.next();
-        let element_pointer_type_id =
-            self.get_pointer_id(element_ty, spirv::StorageClass::Function);
-        block.body.push(Instruction::access_chain(
-            element_pointer_type_id,
-            element_pointer_id,
-            variable.id,
-            &[index_id],
-        ));
-        let id = self.id_gen.next();
-        block.body.push(Instruction::load(
-            result_type_id,
-            id,
-            element_pointer_id,
-            None,
-        ));
-
-        Ok((id, variable))
-    }
-
     /// Indicate that the code requires any one of the listed capabilities.
     ///
     /// If nothing in `capabilities` appears in the available capabilities
@@ -683,11 +635,21 @@ impl Writer {
                 .insert(handle, LocalVariable { id, instruction });
         }
 
-        // cache local variable expressions
         for (handle, expr) in ir_function.expressions.iter() {
-            if matches!(*expr, crate::Expression::LocalVariable(_)) {
+            match *expr {
+                crate::Expression::LocalVariable(_) => {
+                    // Cache the `OpVariable` instruction we generated above as
+                    // the value of this expression.
                     context.cache_expression_value(handle, &mut prelude)?;
                 }
+                crate::Expression::Access { base, .. }
+                | crate::Expression::AccessIndex { base, .. } => {
+                    // Count references to `base` by `Access` and `AccessIndex`
+                    // instructions. See `access_uses` for details.
+                    *context.function.access_uses.entry(base).or_insert(0) += 1;
+                }
+                _ => {}
+            }
         }
 
         let next_id = context.gen_id();

naga/src/valid/expression.rs

@@ -19,8 +19,6 @@ pub enum ExpressionError {
     NegativeIndex(Handle<crate::Expression>),
     #[error("Accessing index {1} is out of {0:?} bounds")]
     IndexOutOfBounds(Handle<crate::Expression>, u32),
-    #[error("The expression {0:?} may only be indexed by a constant")]
-    IndexMustBeConstant(Handle<crate::Expression>),
     #[error("Function argument {0:?} doesn't exist")]
     FunctionArgumentDoesntExist(u32),
     #[error("Loading of {0:?} can't be done")]
@@ -238,10 +236,9 @@ impl super::Validator {
         let stages = match *expression {
             E::Access { base, index } => {
                 let base_type = &resolver[base];
-                // See the documentation for `Expression::Access`.
+                match *base_type {
-                let dynamic_indexing_restricted = match *base_type {
+                    Ti::Matrix { .. }
-                    Ti::Matrix { .. } => true,
+                    | Ti::Vector { .. }
-                    Ti::Vector { .. }
                     | Ti::Array { .. }
                     | Ti::Pointer { .. }
                     | Ti::ValuePointer { size: Some(_), .. }
@@ -262,9 +259,6 @@ impl super::Validator {
                         return Err(ExpressionError::InvalidIndexType(index));
                     }
                 }
-                if dynamic_indexing_restricted && function.expressions[index].is_dynamic_index() {
-                    return Err(ExpressionError::IndexMustBeConstant(base));
-                }
 
                 // If we know both the length and the index, we can do the
                 // bounds check now.

naga/tests/in/index-by-value.wgsl

@@ -0,0 +1,13 @@
+fn index_arg_array(a: array<i32, 5>, i: i32) -> i32 {
+    return a[i];
+}
+
+fn index_let_array(i: i32, j: i32) -> i32 {
+  let a = array<array<i32, 2>, 2>(array(1, 2), array(3, 4));
+  return a[i][j];
+}
+
+fn index_let_matrix(i: i32, j: i32) -> f32 {
+  let a = mat2x2<f32>(1, 2, 3, 4);
+  return a[i][j];
+}

naga/tests/out/ir/index-by-value.compact.ron

@@ -0,0 +1,278 @@
+(
+    types: [
+        (
+            name: None,
+            inner: Scalar((
+                kind: Sint,
+                width: 4,
+            )),
+        ),
+        (
+            name: None,
+            inner: Array(
+                base: 0,
+                size: Constant(5),
+                stride: 4,
+            ),
+        ),
+        (
+            name: None,
+            inner: Array(
+                base: 0,
+                size: Constant(2),
+                stride: 4,
+            ),
+        ),
+        (
+            name: None,
+            inner: Array(
+                base: 2,
+                size: Constant(2),
+                stride: 8,
+            ),
+        ),
+        (
+            name: None,
+            inner: Scalar((
+                kind: Float,
+                width: 4,
+            )),
+        ),
+        (
+            name: None,
+            inner: Matrix(
+                columns: Bi,
+                rows: Bi,
+                scalar: (
+                    kind: Float,
+                    width: 4,
+                ),
+            ),
+        ),
+        (
+            name: None,
+            inner: Vector(
+                size: Bi,
+                scalar: (
+                    kind: Float,
+                    width: 4,
+                ),
+            ),
+        ),
+    ],
+    special_types: (
+        ray_desc: None,
+        ray_intersection: None,
+        predeclared_types: {},
+    ),
+    constants: [],
+    overrides: [],
+    global_variables: [],
+    global_expressions: [],
+    functions: [
+        (
+            name: Some("index_arg_array"),
+            arguments: [
+                (
+                    name: Some("a"),
+                    ty: 1,
+                    binding: None,
+                ),
+                (
+                    name: Some("i"),
+                    ty: 0,
+                    binding: None,
+                ),
+            ],
+            result: Some((
+                ty: 0,
+                binding: None,
+            )),
+            local_variables: [],
+            expressions: [
+                FunctionArgument(0),
+                FunctionArgument(1),
+                Access(
+                    base: 0,
+                    index: 1,
+                ),
+            ],
+            named_expressions: {
+                0: "a",
+                1: "i",
+            },
+            body: [
+                Emit((
+                    start: 2,
+                    end: 3,
+                )),
+                Return(
+                    value: Some(2),
+                ),
+            ],
+        ),
+        (
+            name: Some("index_let_array"),
+            arguments: [
+                (
+                    name: Some("i"),
+                    ty: 0,
+                    binding: None,
+                ),
+                (
+                    name: Some("j"),
+                    ty: 0,
+                    binding: None,
+                ),
+            ],
+            result: Some((
+                ty: 0,
+                binding: None,
+            )),
+            local_variables: [],
+            expressions: [
+                FunctionArgument(0),
+                FunctionArgument(1),
+                Literal(I32(1)),
+                Literal(I32(2)),
+                Compose(
+                    ty: 2,
+                    components: [
+                        2,
+                        3,
+                    ],
+                ),
+                Literal(I32(3)),
+                Literal(I32(4)),
+                Compose(
+                    ty: 2,
+                    components: [
+                        5,
+                        6,
+                    ],
+                ),
+                Compose(
+                    ty: 3,
+                    components: [
+                        4,
+                        7,
+                    ],
+                ),
+                Access(
+                    base: 8,
+                    index: 0,
+                ),
+                Access(
+                    base: 9,
+                    index: 1,
+                ),
+            ],
+            named_expressions: {
+                0: "i",
+                1: "j",
+                8: "a",
+            },
+            body: [
+                Emit((
+                    start: 0,
+                    end: 0,
+                )),
+                Emit((
+                    start: 0,
+                    end: 0,
+                )),
+                Emit((
+                    start: 4,
+                    end: 5,
+                )),
+                Emit((
+                    start: 7,
+                    end: 9,
+                )),
+                Emit((
+                    start: 9,
+                    end: 11,
+                )),
+                Return(
+                    value: Some(10),
+                ),
+            ],
+        ),
+        (
+            name: Some("index_let_matrix"),
+            arguments: [
+                (
+                    name: Some("i"),
+                    ty: 0,
+                    binding: None,
+                ),
+                (
+                    name: Some("j"),
+                    ty: 0,
+                    binding: None,
+                ),
+            ],
+            result: Some((
+                ty: 4,
+                binding: None,
+            )),
+            local_variables: [],
+            expressions: [
+                FunctionArgument(0),
+                FunctionArgument(1),
+                Literal(F32(1.0)),
+                Literal(F32(2.0)),
+                Literal(F32(3.0)),
+                Literal(F32(4.0)),
+                Compose(
+                    ty: 6,
+                    components: [
+                        2,
+                        3,
+                    ],
+                ),
+                Compose(
+                    ty: 6,
+                    components: [
+                        4,
+                        5,
+                    ],
+                ),
+                Compose(
+                    ty: 5,
+                    components: [
+                        6,
+                        7,
+                    ],
+                ),
+                Access(
+                    base: 8,
+                    index: 0,
+                ),
+                Access(
+                    base: 9,
+                    index: 1,
+                ),
+            ],
+            named_expressions: {
+                0: "i",
+                1: "j",
+                8: "a",
+            },
+            body: [
+                Emit((
+                    start: 6,
+                    end: 9,
+                )),
+                Emit((
+                    start: 9,
+                    end: 11,
+                )),
+                Return(
+                    value: Some(10),
+                ),
+            ],
+        ),
+    ],
+    entry_points: [],
+)

naga/tests/out/ir/index-by-value.ron

@@ -0,0 +1,278 @@
+(
+    types: [
+        (
+            name: None,
+            inner: Scalar((
+                kind: Sint,
+                width: 4,
+            )),
+        ),
+        (
+            name: None,
+            inner: Array(
+                base: 0,
+                size: Constant(5),
+                stride: 4,
+            ),
+        ),
+        (
+            name: None,
+            inner: Array(
+                base: 0,
+                size: Constant(2),
+                stride: 4,
+            ),
+        ),
+        (
+            name: None,
+            inner: Array(
+                base: 2,
+                size: Constant(2),
+                stride: 8,
+            ),
+        ),
+        (
+            name: None,
+            inner: Scalar((
+                kind: Float,
+                width: 4,
+            )),
+        ),
+        (
+            name: None,
+            inner: Matrix(
+                columns: Bi,
+                rows: Bi,
+                scalar: (
+                    kind: Float,
+                    width: 4,
+                ),
+            ),
+        ),
+        (
+            name: None,
+            inner: Vector(
+                size: Bi,
+                scalar: (
+                    kind: Float,
+                    width: 4,
+                ),
+            ),
+        ),
+    ],
+    special_types: (
+        ray_desc: None,
+        ray_intersection: None,
+        predeclared_types: {},
+    ),
+    constants: [],
+    overrides: [],
+    global_variables: [],
+    global_expressions: [],
+    functions: [
+        (
+            name: Some("index_arg_array"),
+            arguments: [
+                (
+                    name: Some("a"),
+                    ty: 1,
+                    binding: None,
+                ),
+                (
+                    name: Some("i"),
+                    ty: 0,
+                    binding: None,
+                ),
+            ],
+            result: Some((
+                ty: 0,
+                binding: None,
+            )),
+            local_variables: [],
+            expressions: [
+                FunctionArgument(0),
+                FunctionArgument(1),
+                Access(
+                    base: 0,
+                    index: 1,
+                ),
+            ],
+            named_expressions: {
+                0: "a",
+                1: "i",
+            },
+            body: [
+                Emit((
+                    start: 2,
+                    end: 3,
+                )),
+                Return(
+                    value: Some(2),
+                ),
+            ],
+        ),
+        (
+            name: Some("index_let_array"),
+            arguments: [
+                (
+                    name: Some("i"),
+                    ty: 0,
+                    binding: None,
+                ),
+                (
+                    name: Some("j"),
+                    ty: 0,
+                    binding: None,
+                ),
+            ],
+            result: Some((
+                ty: 0,
+                binding: None,
+            )),
+            local_variables: [],
+            expressions: [
+                FunctionArgument(0),
+                FunctionArgument(1),
+                Literal(I32(1)),
+                Literal(I32(2)),
+                Compose(
+                    ty: 2,
+                    components: [
+                        2,
+                        3,
+                    ],
+                ),
+                Literal(I32(3)),
+                Literal(I32(4)),
+                Compose(
+                    ty: 2,
+                    components: [
+                        5,
+                        6,
+                    ],
+                ),
+                Compose(
+                    ty: 3,
+                    components: [
+                        4,
+                        7,
+                    ],
+                ),
+                Access(
+                    base: 8,
+                    index: 0,
+                ),
+                Access(
+                    base: 9,
+                    index: 1,
+                ),
+            ],
+            named_expressions: {
+                0: "i",
+                1: "j",
+                8: "a",
+            },
+            body: [
+                Emit((
+                    start: 0,
+                    end: 0,
+                )),
+                Emit((
+                    start: 0,
+                    end: 0,
+                )),
+                Emit((
+                    start: 4,
+                    end: 5,
+                )),
+                Emit((
+                    start: 7,
+                    end: 9,
+                )),
+                Emit((
+                    start: 9,
+                    end: 11,
+                )),
+                Return(
+                    value: Some(10),
+                ),
+            ],
+        ),
+        (
+            name: Some("index_let_matrix"),
+            arguments: [
+                (
+                    name: Some("i"),
+                    ty: 0,
+                    binding: None,
+                ),
+                (
+                    name: Some("j"),
+                    ty: 0,
+                    binding: None,
+                ),
+            ],
+            result: Some((
+                ty: 4,
+                binding: None,
+            )),
+            local_variables: [],
+            expressions: [
+                FunctionArgument(0),
+                FunctionArgument(1),
+                Literal(F32(1.0)),
+                Literal(F32(2.0)),
+                Literal(F32(3.0)),
+                Literal(F32(4.0)),
+                Compose(
+                    ty: 6,
+                    components: [
+                        2,
+                        3,
+                    ],
+                ),
+                Compose(
+                    ty: 6,
+                    components: [
+                        4,
+                        5,
+                    ],
+                ),
+                Compose(
+                    ty: 5,
+                    components: [
+                        6,
+                        7,
+                    ],
+                ),
+                Access(
+                    base: 8,
+                    index: 0,
+                ),
+                Access(
+                    base: 9,
+                    index: 1,
+                ),
+            ],
+            named_expressions: {
+                0: "i",
+                1: "j",
+                8: "a",
+            },
+            body: [
+                Emit((
+                    start: 6,
+                    end: 9,
+                )),
+                Emit((
+                    start: 9,
+                    end: 11,
+                )),
+                Return(
+                    value: Some(10),
+                ),
+            ],
+        ),
+    ],
+    entry_points: [],
+)

naga/tests/out/spv/access.spvasm

@@ -240,7 +240,7 @@ OpDecorate %343 BuiltIn Position
 %215 = OpConstantComposite  %31  %66 %66 %66 %66
 %216 = OpConstantComposite  %34  %214 %215
 %222 = OpTypeFunction %5 %32 %5
-%224 = OpTypePointer Function %32
+%225 = OpTypePointer Function %32
 %231 = OpTypeFunction %3 %37
 %238 = OpTypeFunction %2 %37
 %244 = OpTypeFunction %3 %40
@@ -434,11 +434,11 @@ OpFunctionEnd
 %219 = OpFunctionParameter  %32
 %220 = OpFunctionParameter  %5
 %218 = OpLabel
-%225 = OpVariable  %224  Function
+%224 = OpVariable  %225  Function
 OpBranch %223
 %223 = OpLabel
-OpStore %225 %219
+OpStore %224 %219
-%226 = OpAccessChain  %88  %225 %220
+%226 = OpAccessChain  %88  %224 %220
 %227 = OpLoad  %5  %226
 OpReturnValue %227
 OpFunctionEnd
@@ -481,7 +481,7 @@ OpFunctionEnd
 %260 = OpFunction  %2  None %60
 %254 = OpLabel
 %272 = OpVariable  %27  Function %265
-%273 = OpVariable  %224  Function %274
+%273 = OpVariable  %225  Function %274
 %257 = OpLoad  %3  %255
 %261 = OpAccessChain  %61  %49 %41
 %263 = OpAccessChain  %262  %52 %41
@@ -549,7 +549,7 @@ OpReturn
 OpFunctionEnd
 %344 = OpFunction  %2  None %60
 %340 = OpLabel
-%348 = OpVariable  %224  Function
+%348 = OpVariable  %225  Function
 %342 = OpLoad  %3  %341
 OpBranch %347
 %347 = OpLabel

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

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

naga/tests/snapshots.rs

@@ -936,6 +936,7 @@ fn convert_wgsl() {
             Targets::SPIRV | Targets::METAL | Targets::GLSL | Targets::HLSL | Targets::WGSL,
         ),
         ("6220-break-from-loop", Targets::SPIRV),
+        ("index-by-value", Targets::SPIRV | Targets::IR),
     ];
 
     for &(name, targets) in inputs.iter() {

naga/tests/wgsl_errors.rs

@@ -1358,21 +1358,6 @@ fn missing_bindings2() {
 
 #[test]
 fn invalid_access() {
-    check_validation! {
-        "
-        fn matrix_by_value(m: mat4x4<f32>, i: i32) -> vec4<f32> {
-            return m[i];
-        }
-        ":
-        Err(naga::valid::ValidationError::Function {
-            source: naga::valid::FunctionError::Expression {
-                source: naga::valid::ExpressionError::IndexMustBeConstant(_),
-                ..
-            },
-            ..
-        })
-    }
-
     check_validation! {
         r#"
             fn main() -> f32 {
@@ -1415,6 +1400,15 @@ fn valid_access() {
         ":
         Ok(_)
     }
+
+    check_validation! {
+        "
+        fn matrix_by_value(m: mat4x4<f32>, i: i32) -> vec4<f32> {
+            return m[i];
+        }
+        ":
+        Ok(_)
+    }
 }
 
 #[test]