[Migrated] We should override `thir_body` to inject loop merge points for structurization.

Issue automatically imported from old repo: https://github.com/EmbarkStudios/rust-gpu/issues/944

Originally creatd by eddyb on 2022-11-24T08:17:49Z 👍: 1

I was writing up this SPIR-T issue when I realized we can do much better than I ever thought we can:

• https://github.com/EmbarkStudios/spirt/issues/11

The "subgroup reconvergence"/"loop merges" problem statement

the fundamental ambiguity here (wrt reconvergence), that explicit merges solve, looks like this:

loop {
    // ...
    if cond {  // the "then" edge *leaves the loop*
        A();   // !!! it might be important that this STAYS INSIDE the loop!
        break; // this is just a linear A -> B edge in the CFG (outside the loop's cycle)
    }
}
B(); // !!! it might be important that this STAYS OUTSIDE the loop!

because the A() -> B() edge just looks like a redundant fusable edge (isomorphic to a single A(); B(); basic block), structurizers will very likely either produce:

• loop { ... if cond { break; } } A(); B(); (moving A() out of the loop)
  • A() no longer called for "early finishers" (losing a desired side-effect)
• loop { ... if cond { A(); B(); break; } } (moving B() into the loop)
  • B() now being called for "early finishers" (gaining an undesired side-effect)

and they're both bad because a non-uniform break (i.e. an "early finisher") will block on subgroup neighbors (semantically, at least, in hardware the now-inactive lanes will likely stick around while loop instructions continue being executed for remaining lanes, until they all eventually break)

An actual reasonable solution for once (what prompted me to open this issue)

A better approach might be to "just" make up a SPIR-T instruction that acts as "loop merge barrier", i.e.:

loop {
    // ...
    if cond {
        A();   // !!! definitely INSIDE the loop
        break; // just an edge to the `LoopMerge`, in the CFG
    }
}
asm!("spirt.ControlInst.LoopMerge");
B(); // !!! definitely OUTSIDE the loop

Each LoopMerge would be consumed by one loop structurization - if you have nested loops, you'd need to be careful not to forget to add them to each loop.

But wait, even if we make this ergonomic with macros... we'd still want to check that you don't use any instructions that "care" about e.g. subgroups, right? That could be e.g. a MIR check pass, right? Which we can inject... just like...

Yes, that's right, rustc_codegen_spirv has enough query overriding power that it could override thir_body, which MIR construction uses as the source of truth, to introduce additional nodes around loops, allowing Rust to be on par with GLSL wrt structured control-flow guarantees.

Additionally, this could be used to replace #[spirv(unroll_loops)] (once that's removed, see https://github.com/EmbarkStudios/rust-gpu/pull/940#issuecomment-1320004784 for more context), as it would allow us to check for a (new) #[spirv(unroll)] attribute on the loop expression, and then store that in the injected merge point.

That would be much nicer than fn-level catch-alls, and generally a good fit for SPIR-V expressivity.

Start with the linked rustc MIR construction reference and the thir_body query override in rustc_codegen_spirv. Investigate how loop nodes could carry LoopMerge and loop-level unroll information, including nested loops and subgroup-sensitive instructions. Done means loop merge points and the proposed unroll attribute are injected and validated without losing the intended control-flow semantics.