[naga] resolve the size of override-sized arrays in backends

[teoxoy]

Connections
Fixes #6722
Closes #6746

Description
Implements Jim's suggestion of resolving the size of override-sized arrays in backends so that we don't have to rebuild the types arena when processing overrides.

Testing
The new and existing tests passed.

[kentslaney]

I'm confused about the status of this PR. In #6746 you said

I think it turned out well

but in #6788 you say

ultimately my attempt didn't work

Am I misreading or did something change? And where does this leave #6746?

[teoxoy]

but in #6788 you say

ultimately my attempt didn't work

In #6788 (comment) I was specifically talking about the issue Jim brought up in #6788; this PR doesn't touch the compactor.

And where does this leave #6746?

This PR obsoletes #6746.

[jimblandy]

Still reviewing.

[kentslaney]

@cwfitzgerald via Wgpu Users in reply to @junglie85

I assume I can parse a wgsl shaders through naga at build time to generate the module, then load that from disk (or embed bytes) at runtime? Does the module still need to be compiled to a specific backend format when creating the shader module from the naga module?

You can theoretically do it, as you can serialize a Naga module, but yeah we still need to translate it at runtime up hlsl/glsl/spirv/msl

Wouldn't this mean that it's better to resolve the size of override-sized arrays on the front end to minimize runtime overhead?

[cwfitzgerald]

Eh, way more work than the gain - that's not a very common workflow.

[jimblandy]

Is this an accurate description of the new process_overrides?

/// Replace all overrides in `module` with fully-evaluated constant expressions.
///
/// Given `pipeline_constants`, providing values for all overrides in
/// `module`:
///
/// - Replace all [`Override`] expressions with fully-evaluated
///   constant expressions.
///
/// - Replace all [`Override`][paso] array sizes with [`Expression`]
///   array sizes, referring to fully-evaluated constant expressions.
///
/// - Empty out the `module.overrides` arena.
///
/// Although the above is described in terms of changes to `module`'s
/// contents, this function only actually has shared access to
/// `module`. When changes are needed, this function clones `module`
/// and returns a [`Cow::Owned`] value. If no changes are needed, this
/// function returns a [`Cow::Borrowed`] value that just passes along
/// the original reference.
///
/// [`Override`]: Expression::Override
/// [paso]: crate::PendingArraySize::Override
/// [`Expression`]: crate::PendingArraySize::Expression

[teoxoy]

That does sound more accurate than the current doc. Do you want to push the commit? Or else I can fixup the first commit.

[jimblandy]

Great test case.

[teoxoy]

Kudos to @kentslaney.

[jimblandy]

UniqueArena::replace is specified to panic if the new value already exists in the arena. The WGSL front end also creates PendingArraySize::Expression array sizes. Why do we believe that these replace calls will never panic?

If this isn't a bug, it definitely deserves a comment.

[teoxoy]

The expression handle in PendingArraySize::Expression is the only handle to that expression.

I realize that this is not validated though, I will add some validation for this.

[teoxoy]

The expression handle in PendingArraySize::Expression is the only handle to that expression.

This is needed for type equivalency, arrays sized via an override expression should never be equal to other arrays.

[teoxoy]

The validation is not necessary, I missed the fact that we are working with a UniqueArena #6787 (comment).

[jimblandy]

There's something about this check that makes me uncomfortable. It seems like either the UniqueArena will ensure the expressions are unique automatically, in which case the check will never fail, or you've got array types with distinct base/stride but the same expression, which doesn't seem necessary to forbid.

Do we have a test case that exercises this check?

[teoxoy]

You are right, this validation is not needed. When thinking about your comment about UniqueArena::replace (#6787 (comment)) I thought it was needed.

I will remove it and add a comment instead.

[jimblandy]

I think we need to find a way to handle this without editing the type arena at all.

const X: array<u32, 2> = array(1, 2);
override Y = X[0u];
override Z = X[0u];

var<workgroup> aY: array<u32, Y>;
var<workgroup> aZ: array<u32, Z>;

This panics because the types aren't unique:

$ d=$HOME/play/naga-array-override && cd ~/rust/wgpu/naga-cli && RUST_BACKTRACE=1 cargo run $d/in.wgsl $d/out.metal
     Running `/home/jimb/rust/wgpu/target/debug/naga /home/jimb/play/naga-array-override/in.wgsl /home/jimb/play/naga-array-override/out.metal`
thread 'main' panicked at naga/src/arena/unique_arena.rs:153:9:
assertion failed: added && index == self.set.len() - 1
stack backtrace:
...
   3: naga::arena::unique_arena::UniqueArena<T>::replace
             at /home/jimb/rust/wgpu/naga/src/arena/unique_arena.rs:153:9
   4: naga::back::pipeline_constants::process_pending
             at /home/jimb/rust/wgpu/naga/src/back/pipeline_constants.rs:260:21
   5: naga::back::pipeline_constants::process_overrides
             at /home/jimb/rust/wgpu/naga/src/back/pipeline_constants.rs:210:5
   6: naga::write_output
             at ./src/bin/naga.rs:717:17
   7: naga::run
             at ./src/bin/naga.rs:582:9
   8: naga::main
             at ./src/bin/naga.rs:361:21

Constant evaluation handles AccessIndex { base, index } by simply returning the index'th subexpression of base, so both Y and Z end up with Override::init holding the same Handle<Expression>. Then process_pending tries to use that handle for both types.

[jimblandy]

I have different idea about how to solve this. I'll put together a PR tomorrow.

Leaving this open for the time being, since I'm not sure how well my idea will work.

[kentslaney]

is it worth taking another look at #6746?