[Mlir-commits] [mlir] [mlir][amdgpu] Introduce `assume_subgroup_uniform` op (PR #152740)

[Krzysztof Drewniak]

krzysz00 wrote:

1. Per discussion offline, this should be a `gpu.*` op that lets us also lower to SPIR-V, Nvidia, etc.
2. @arsenm The main motivation behind this op was to allow for `readfirstlane` to be _more_ speculatable than in usually is by letting you stick a "trust me" on it ... though we're also realizing that there's an abstraction we want around readfirstlane, the equivalent SPIR-V op, Nvidia's equivalent, etc. .

Is there a way to make `readfirstlane` more sepculatable?

To give higher-level context. The "I'm not reading all that" context is over in https://gist.github.com/Hardcode84/7d8c81d4081d8bb99a3bf74054ff7768, which has some IREE MLIR and some of the generated LLVM IR and assembly.

Even after applying https://github.com/llvm/llvm-project/pull/152581 to make the LDS subgroup indices recognizable as uniform, the `readfirstlane` for the move to m0 gets inserted inside the loop where the direct-to-LDS loads are, not outside the loop where the pointer is computed.

That is, we have
```
l0 = f(tidx >> 6, tidy) ; uniform, vgpr
l1 = l0 + C1
l2 = l0 + C2
...
for (k from 0 upto K) {
  load_lds(ptr addrspace(1/7) g(K, tidx, tidy, 0), ptr addrspace(3) l0, ...)
  load_lds(g(K, tidx, tidy, 1), l1, ...)
  ... [barriers, compute stuff, etc. this loop might be unrolled]
}
```
and because the l_i are computed in VGPRs, they only get `readfirstlane`'d automatically right ta the point where that transition is needed ... and then it doesn't get hoisted either because it's too late for that or because that's not a trivial rewrite or because no one's added logic for that.

So adding an MLIR-level abstraction around readfirstlane will not only allow us to make it platform-independent, but it'll allow us (by way of letting us stick an assertion of speculatability on it) to LICM it.

https://github.com/llvm/llvm-project/pull/152740