[Mlir-commits] [mlir] [mlir][gpu] Store the all-reduce result from the first lane only (PR #221532)

[Alessandro Potenza]

alepot55 wrote:

I think that holds for the full-width branch, but the second stage goes through `isPartialSubgroup` with `numSubgroups` as the shuffle width, and there the xor partner can fall outside the active range: with 3 warps lane 1 xors with lane 3 at offset 2, `shfl.sync.bfly` reports it out of range (the clamp is `width - 1`), the accumulation is skipped, and lane 1 ends with its own partial plus lane 0's; lane 2 likewise never sees lane 1. Only lane 0 holds the full sum, which is what the comment on `createSubgroupReduce` says, and then every one of those lanes stores its value to `buffer[0]`. That is the case the interpreter flagged (7 of 501 programs, all with 3 warps); on an RTX PRO 6000 lane 0's store happened to be the one that survived. With a power-of-two number of warps every lane does hold the total, as you say. Would you rather I keep the predication as is, or gate it on the partial case only?

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