[Mlir-commits] [mlir] [mlir][nvgpu] Make `mbarrier.try_wait` fallable (PR #96508)

[llvmlistbot at llvm.org]

builtin-pb wrote:

# What guarantees uniformity of `@P1 bra.uni` in this `mbarrier.try_wait` loop?

I am formalizing compiled FlashAttention-4 PTX and encountered the following polling pattern, also present in the [inline PTX lowering discussed above](https://github.com/llvm/llvm-project/pull/96508#discussion_r1651235713):

```ptx
    .reg .pred P1;
LAB_WAIT:
    mbarrier.try_wait.parity.shared.b64 P1, [addr], phase, ticks;
    @P1 bra.uni DONE;
    bra.uni LAB_WAIT;
DONE:
```

Assume uniform `addr`, `phase`, and `ticks`, a correctly initialized and aligned shared-memory barrier, and valid phase/lifetime handling. **What guarantees identical `P1` values among the active warp threads executing the guarded `bra.uni`?**

The [PTX `bra.uni` contract](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#control-flow-instructions-bra) requires agreement at that instruction, and [§9.5](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#divergence-of-threads-in-control-constructs) puts the non-divergence guarantee on the compiler or code author. The [`try_wait` description](https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#parallel-synchronization-and-communication-instructions-mbarrier-test-wait-mbarrier-try-wait) allows a suspended thread to resume before phase completion after a time limit. I could not identify the rule connecting these contracts in the checked PTX/NVVM specifications, LLVM convergence documentation, or operation definitions.

Do threads executing one wait instruction together receive a common result? If so, what keeps separately resumed groups with different results from regrouping at the branch, including across polling iterations? Only threads executing the branch together need to agree; the question does not require all original waiters to finish together.

The [original operation definition](https://github.com/llvm/llvm-project/blob/b99d0b34400176cb9183113b96b245400caaf8d8/mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td#L326) states no uniformity precondition. The later [single-attempt intrinsic operation in #170285](https://github.com/llvm/llvm-project/pull/170285) is useful context, but I am trying to justify the already emitted loop. Which PTX/NVVM rule, additional caller precondition, or target-specific scheduling guarantee does this lowering rely on?

*AI-assisted documentation research and drafting; posted with the account owner's approval.*


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