[naga msl-out] Use iteration limits to avoid unbounded loops

[jimblandy]

Rather than injecting branches on volatile bools, Naga's backend for Metal Shading Language should avoid undefined behavior simply by imposing an iteration limit on every loop. Unlike the volatile bools, the optimizer will be able to reason about these limits and eliminate them when possible. Optimizations like unrolling are defeated by the volatile bool, but not by the iteration limit.

See a full description of the idea, with godbolt experiments, here: #6528 (comment)

[DemiMarie]

Will this cause problems due to slow 64-bit arithmetic? Is it possible for the GPU to actually run 2³² loop iterations before the device is lost?

[teoxoy]

I was proposing to add this to the WGSL spec first to avoid UB in both the Metal compiler and DXC. DXC doesn't make the same assumptions as Metal but it does technically change the behavior of the program as written in WGSL #6528 (comment).

My main argument is that given WGSL allows loops to be infinite and that we target MSL & HLSL (both being based on C++), therefore need to prove to downstream compilers that loops are finite. Currently the only way to do that is with volatiles but that's expensive, if we add an upper bound on the nr of loop iterations we should be in a better position performance wise. Loops can even be unrolled as you found in #6528 (comment).

Is it possible for the GPU to actually run 2³² loop iterations before the device is lost?

We could run some tests to see if this is the case. A more expensive way to make sure this happens is to add the mechanism ourselves but that's nontrivial. It might be enough to tell developers that loops have the upper bound and that they might hit it.

[DemiMarie]

A 2⁶⁴ - 1 limit allows assuming that the limit will never actually be reached, because the GPU cannot execute enough cycles to actually hit it in a human lifetime.

[kvark]

I believe I'm facing this issue as well in luminal-ai/luminal#158
It has a Metal backend that generates for loops, and we are generating fine for loops in WGSL, but they turn out to a bit of an abomination in MSL output:

        while(true) {
            if (!loop_init_1) {
                uint _e45 = loop_l;
                loop_l = _e45 + 1u;
            }
            loop_init_1 = false;
            uint _e32 = loop_l;
            if (_e32 < 512u) {
            } else {
                break;
            }
            {
                // body
            }
        }

This would explain why Blade's Metal backend produces slower kernels than the native Metal backend, and is unfortunate.

[cwfitzgerald]

@kvark this is unrelated actually - this is a byproduct of the fact that wgsl continuing blocks having potentially arbitrary code in them. As we can't guarantee that there's only a single expression, we can't use for loops in HLSL and MSL. I tried to solve this in the backends, but the IR generated with a separate Load and Store (causing the temporary variable in the continuing block) means it was a ton of work. We might need to make some small IR changes to make that happen. #6521 is the actual issue here. Definitely want to solve this in some way, just wasn't as easy as I had time for.