[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.

[cwfitzgerald]

So I put together a gpu side benchmark and unfortunately I don't think these numbers are representative 😅 The benchmark is #6987 and can be run with cargo bench "Loop Workaround". You currently need to divide the time by 100 to get the real number.

I'm currently getting 4.45us GPU time on my AMD laptop with clock speeds locked. I get the same number for both of these shaders:

@group(0) @binding(0) var<storage, read_write> data: array<u32>;

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

@group(0) @binding(0) var<storage, read_write> data: array<u32>;

@compute @workgroup_size(64)
fn addABunch(@builtin(global_invocation_id) global_id: vec3<u32>) {
    var x: u32 = data[global_id.x];
    data[global_id.x] = x * 2;
}

Ooops....

Going to look into this a smidge more to make sure the numbers are really what's going on and to see if there's some math we can do to preserve the loop...

[cwfitzgerald]

Alright, I've updated the PR to have this shader which the compiler can't see through the body of, but should be able to easily see through the bounds check of, and now it take a nice rock solid 74ms on my machine. I've pushed these changes to the PR I linked. If you rebase/cherry pick this on top of your changes, you should be able to see the difference.

@group(0) @binding(0) var<storage, read_write> data: array<u32>;

@compute @workgroup_size(64)
fn addABunch(@builtin(global_invocation_id) global_id: vec3<u32>) {
    var x: u32 = data[global_id.x];
    for (var i = 1u; i <= 100000u; i++) {
      x = u32(sin(f32(x * 120u)));
    }
    data[global_id.x] = x;
}

[cwfitzgerald]

Running it without locked clocks maxes out my gpu clocks and I get a stable 18ms runtime (locked clocks are 700mhz, boost clocks ~2800mhz)

[cwfitzgerald]

Ran this benchmark on my M1:

	No Workaround	Current Workaround	This PR	My Suggestion
M1 Mini	143ms	157ms	162ms	183ms
RTX 4070	7.24ms	-	9.5ms	10.1ms

So it seems like both this PR and my idea are significantly worse!

[cwfitzgerald]

I have pushed to my fork to make this easier to reproduce:

• trunk with benchmark: https://github.com/cwfitzgerald/wgpu/tree/ilub/old
• this pr suggestion: https://github.com/cwfitzgerald/wgpu/tree/ilub/benchmark
• my suggestion: https://github.com/cwfitzgerald/wgpu/tree/ilub/connor

[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.