[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::back::spv::Function`.

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

Fixes gfx-rs#6358.
Fixes gfx-rs#4337.
Alternative to gfx-rs#6362.

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::{
- arena::Handle,
- proc::{index::GuardedIndex, TypeResolution},
- Statement,
-};
+use crate::{arena::Handle, proc::index::GuardedIndex, Statement};
 use spirv::Word;
 
 fn get_dimension(type_inner: &crate::TypeInner) -> Dimension {
@@ -345,34 +341,59 @@ 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 {
- base: ty_element, ..
- } => {
- let index_id = self.cached[index];
- let base_id = self.cached[base];
- let base_ty = match self.fun_info[base].ty {
- TypeResolution::Handle(handle) => handle,
- TypeResolution::Value(_) => {
- return Err(Error::Validation(
- "Array types should always be in the arena",
- ))
+ crate::TypeInner::Array { .. } | crate::TypeInner::Matrix { .. } => {
+ // See if `index` is known at compile time.
+ match GuardedIndex::from_expression(index, self.ir_function, self.ir_module)
+ {
+ GuardedIndex::Known(value) => {
+ // If `index` is known, we can just use `OpCompositeExtract`.
+ //
+ // We never need bounds checks for these cases: everything
+ // size is statically known and checked in validation.
+ let id = self.gen_id();
+ let base_id = self.cached[base];
+ block.body.push(Instruction::composite_extract(
+ result_type_id,
+ id,
+ base_id,
+ &[value],
+ ));
+ id
  }
- };
- let (id, variable) = self.writer.promote_access_expression_to_variable(
- result_type_id,
- base_id,
- base_ty,
- index_id,
- ty_element,
- block,
- )?;
- self.function.internal_variables.push(variable);
- id
+ 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,
+ )?
+ }
+ }
  }
- // wgpu#4337: Support `crate::TypeInner::Matrix`
  crate::TypeInner::BindingArray {
  base: binding_type, ..
  } => {
@@ -435,6 +456,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 +1769,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 +2029,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,10 +635,20 @@ 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(_)) {
- context.cache_expression_value(handle, &mut prelude)?;
+ 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;
+ }
+ _ => {}
  }
  }
 

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`.
- let dynamic_indexing_restricted = match *base_type {
- Ti::Matrix { .. } => true,
- Ti::Vector { .. }
+ match *base_type {
+ Ti::Matrix { .. }
+ | 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];
+}