feat(naga): constant evaluation for select

[ErichDonGubler]

Connections

• Blocks Fix typing for select #7572.

Testing

• Still deciding on test coverage here. WebGPU CTS is pretty thorough on this front; maybe we can defer?

Squash or Rebase?

Squash unless multiple commits.

Checklist

• If this contains user-facing changes, add a CHANGELOG.md entry.

[jimblandy]

I think it would be good to have just one or two smoke tests for this, like the ones already there at the bottom of constant_evaluator.rs. Nothing even close to exhaustive, just a witness that it was observed to have worked once or twice.

[jimblandy]

I could be mistaken on some of these, but I think there may be some simplifications possible.

[jimblandy]

I wonder, is this cast really necessary at all? That is, in #7572, we will end up doing all the type checking and automatic conversions in the WGSL front end, at which point this code can assume that its operands are well-typed. Then it only needs to throw errors if it encounters something it can't make progress with (like, a condition that's not bool or vecN<bool>).

[ErichDonGubler]

All I know is that if I don't have it, the tests cases I'm using fail, even when consuming #7572. 😅

[ErichDonGubler]

Concretely, this fails in constant evaluation if we don't cast:

const_assert select(42f, 9001i, true) == 9001f;

My understanding is that we want to match the first argument's type, independent of condition.

[jimblandy]

In the case where the condition is bool and not vecN<bool>, can't we just return reject or accept directly, without ever looking at its value?

[jimblandy]

Is it not possible to reuse any of the component_wise machinery for this?

[ErichDonGubler]

RE: component_wise_*: It wasn't possible previously because we're using a cast operation before performing a component-wise operation here; we already take on the scope of delving into all components before knowing we have the right types on LHS and RHS.

I might have something wrong here, though.

[ErichDonGubler]

At least with my current arguments, this simplification depends on whether we can remove cast at this conversation: #7602 (comment)

[jimblandy]

I played around with this a bit today and realized something that I hadn't really appreciated before, and which I think is part of why this has not been smooth (and why I've been giving not-so-great advice).

The select builtin is, I believe, unique among WGSL builtins in that it is the only function that completely ignores one of its arguments' values.

This undercuts some assumptions we've generally made in the WGSL front end and constant evaluator. Generally, if constant evaluation was able to take place at all, then that means the evaluator was able to find the values it needed to operate on, which in turn means that the call was well-typed. If the call is ill-typed, then constant evaluation fails, and we take the branch in try_eval_and_append that simply emits the Expression to the arena, unevaluated, where the validator will find it. Since the validator is responsible for fully type-checking the module, we'll get an error.

According to this pattern, we would simply rewrite select(a, ..., false) as a, without ever inspecting the ... at all. But doing so means that if there's a type error in that expression, the validator will never see it, and the error will go unreported.

Thus, for select only, if constant evaluation is possible, then somebody prior to constant evaluation must fully type-check the call. This would need to be either the WGSL front end or the constant evaluator itself.

If we decide to fully check it in the WGSL front end, here's what I came up with for naga::front::wgsl::lower::Lowerer::expr:

                        "select" => {
                            let mut args = ctx.prepare_args(arguments, 3, span);

                            let mut values = [
                                self.expression_for_abstract(args.next()?, ctx)?,
                                self.expression_for_abstract(args.next()?, ctx)?,
                            ];
                            let condition = self.expression(args.next()?, ctx)?;
                            for &value in &values {
                                ctx.grow_types(value)?;
                            }
                            let mut consensus_scalar = ctx
                                .automatic_conversion_consensus(&values)
                                .map_err(|_idx| Error::SelectArgumentsHaveNoCommonType {
                                    select_call: span,
                                })?;
                            if !ctx.is_const(condition) {
                                consensus_scalar = consensus_scalar.concretize();
                            }

                            ctx.convert_slice_to_common_leaf_scalar(&mut values, consensus_scalar)?;

                            let [reject, accept] = values;

                            args.finish()?;

                            ir::Expression::Select {
                                reject,
                                accept,
                                condition,
                            }
                        }

This also handles The Mysterious Rule in the WGSL overload resolution algorithm, with the call to is_const. (I had to move Scalar::concretize from naga::front::wgsl::lower::conversion to naga::proc::type_methods, which is probably where it always belonged.)

I also noticed that we had some really weird tests. wgsl_errors::select seems to think that this code should get a naga::valid::ExpressionError::SelectConditionNotABool error:

        fn select_pointers(which: bool) -> i32 {
            var x: i32 = 1;
            var y: i32 = 2;
            let p = select(&x, &y, which);
            return *p;
        }

There is no planet on which WGSL's select can operate on pointers. I have no idea who wrote this test. Probably me. I owe you a donut or some nice pastry when we're in Toronto.

[ErichDonGubler]

I'm going to combine this PR into #7572. There have been enough linked questions of behavior that I think it's better to test, review, and merge them as a unit.

[ErichDonGubler]

@jimblandy: I have incorporated your code, and modified it quite a bit so that diagnostics help the user better than the "weird" case you noted. 🙂