[naga msl-out hlsl-out] Improve workaround for infinite loops causing undefined behaviour

[jamienicol]

Connections
Partially solves #6572 (only hlsl and msl)

Description

We must ensure that all loops emitted by the naga backends will terminate, in order to avoid undefined behaviour. This was previously implemented for the msl backend in #6545. However, the usage of volatile prevents the compiler from making other important optimizations. This patch improves the msl workaround and additionally implements it for hlsl. The spv implementation will be left for a follow up.

Rather than using volatile, this patch increments a counter on every loop iteration, breaking from the loop after 2^64 iterations. This ensures the compiler treats the loop as finite thereby avoiding undefined behaviour, whilst at the same time allowing for other optimizations and in reality not actually affecting execution.

Testing

Modified the hello_compute example to reproduce the undefined behaviour then ensured this patch avoids it.

Inspected metal and hlsl output in shaderplayground to ensure simple loop unrolling still works as expected

Checklist

• Run cargo fmt.
• Run taplo format.
• Run cargo clippy. If applicable, add:
  • --target wasm32-unknown-unknown
  • --target wasm32-unknown-emscripten
• Run cargo xtask test to run tests.
• Add change to CHANGELOG.md. See simple instructions inside file.

[cwfitzgerald]

Hey @rudderbucky can you try this PR to see how the shader performance is with the workaround in place?

[rudderbucky]

So it turns out I was able to find a couple workarounds to avoid the performance hit... I will give this a whirl around the weekend though, please remind me if I forget @cwfitzgerald

[Imberflur]

The checks probably need to be moved to the beginning of the loop as mentioned in #6528 (comment)

[jamienicol]

The checks probably need to be moved to the beginning of the loop as mentioned in #6528 (comment)

Ah of course. I'll update the patch

[jamienicol]

The checks probably need to be moved to the beginning of the loop as mentioned in #6528 (comment)

I think this may have been the cause of the additional macos failures. At least with the checks at the top of the loop we no longer have any new failures on CI. So I have removed the patch that updated the test expectations

[jamienicol]

@rudderbucky did you have a chance to test the performance with this PR?

[cwfitzgerald]

Comment on the form of the workaround

[cwfitzgerald]

This construct feels weird to me @jimblandy is there a reason this was the original suggestion? I would expect this to take the form of:

loop_bound.x += 1;
if (loop_bound.x == 4294967295u) {
    loop_bound.y += 1;
    if (loop_bound.y == 4294967295u) {
        break;
    }
}

I don't think this brings us into any uniformity issues compared to the above. Lifespan of variables should be the same too.

This brings us from 3 comparisons (4 as-written) and 2 additions in the hot path to just 1 comparison and 1 addition. While sure this isn't that big of a deal, we're going to be stacking this bad boy on every single loop, Additionally it may help loop bound analysis eliminate this if the first condition is simpler.

[jamienicol]

I thought I'd get some numbers to help verify this. Obvious caveat this is just one testcase on a couple of devices, but better than nothing.

I modified the hello_compute example to do 10 milllion loop iterations like so, and timed its duration:

@compute @workgroup_size(64)
fn doubleMe(@builtin(global_invocation_id) global_id: vec3<u32>) {
    var x: u32 = input[global_id.x];
    for (var i = 1u; i <= 10000000u; i++) {
      x = x + 1u;
    }
    output[global_id.x] = x;
}

	No loop bounding	Existing volatile workaround	Current patch workaround	Connor's suggestion
MSL M2 Macbook Pro	271ms	923ms	351ms	375ms
HLSL AMD Radeon Pro W6400	124ms	N/A	284ms	220ms

Metal seems to have a slight preference for the way the PR is currently written. DXC for connor's suggestion. Both are significantly better than the current situation, but still significantly worse than no loop bounding at all.

The current construct makes sense to me as "emulate a u64 counter with a vec2<u32>". But I think these results give me a slight preference for switching to @cwfitzgerald's suggestion. (Though I think we should be doing == 0u rather than == 4294967295u as the comparison occurs after the increment. Not that it will really matter in practice). Can anyone think of any specific shader constructs we should test this further with? Or are we happy enough to proceed based on this - it's clearly still a performance hit, but much better than the current situation.

One thing I noticed looking at Tint's code is they do some analysis of whether a loop is finite, and only emit the workaround if required. Perhaps long term we need to do something similar to really solve the performance issues.

[rudderbucky]

Hi... I tried setting my game up to hit this performance snag but haven't been able to get performance to significantly drop even after disabling a few performance improvements I made around the same time... maybe something else in wgpu 23 helped improve performance again... sorry about that.