[clang] [llvm] [AMDGPU] Implement AsyncMark Stages (PR #220442)

[Ryan Mitchell via cfe-commits]

================
@@ -16,31 +16,76 @@ internally by the compiler. A thread that initiates one or more async operations
 An *asyncmark* created by a thread can be used to track async operations
 initiated by that thread.
 
+### Stages
+
+A *stage* names a kind of async operation. Each async operation *belongs to* the
+one stage determined by the instruction that initiates it.
+
+The stages are:
+
+| Bit | Stage | Async operations |
+|---|---|---|
+| 0 | `TENSOR` | tensor loads and stores |
+| 1 | `GLOBAL_LOAD_ASYNC_TO_LDS` | global loads async to LDS |
+| 2 | `GLOBAL_LOAD_ASYNC_TO_LDS_MCAST` | multicast (cluster) global loads async to LDS |
+| 3 | `GLOBAL_STORE_ASYNC_FROM_LDS` | async global stores from LDS |
+| 5 | `BUFFER_GLOBAL_LOAD` | buffer loads to LDS and pre-gfx1250 global loads to LDS |
+
+Bits 4 and 6 through 10 are reserved for future async operations, and no
+operation belongs to them yet.
+
+Which async operations a given subtarget actually has is described in
+{ref}`AMDGPU DMA Operations <amdgpu-dma-operations>`. A stage exists on every
+subtarget that supports asyncmarks, whether or not that subtarget has any
+operation belonging to it.
+
+### Stage Masks
+
+Both intrinsics take a *stage mask*: an 11-bit value in which a set bit means
+"do not participate". In particular, an asyncmark *omits* the stages its 
+mask names, and a wait *ignores* them. The mask `0` therefore names no stage 
+and so omits/ignores none.
+
+A mask may set the bit of a reserved stage. Leaving out a stage whose operations
+do not exist yet is harmless, and lets a mask keep its meaning as the reserved
+bits are filled in. However, omitting/ignoring bits that are neither supported
+nor reserved is an error.
+
 ### Current Sequence
 
-The abstract machine maintains a sequence of asyncmarks during the execution of
-a function body, which excludes any asyncmarks produced by calls to other
-functions encountered in the currently executing function. The state of this
-sequence at each program point in the function is called the *current sequence*.
+The abstract machine maintains a separate sequence of asyncmarks *for each
+stage* during the execution of a function body, which excludes any asyncmarks
+produced by calls to other functions encountered in the currently executing
+function. The state of the sequence for a stage `S` at each program point in
+the function is called the *current sequence of* `S`.
 
-### `@llvm.amdgcn.asyncmark()`
+The sequences of distinct stages are independent: appending to one does not
+affect the length or contents of any other, even though multiple sequences may
+be appended to with a single call to asyncmark (by choosing to not *omit* 
+multiple stages).
 
-Produces an asyncmark and appends it to the current sequence.
+### `@llvm.amdgcn.asyncmark(i32 %K)`
 
-### `@llvm.amdgcn.wait.asyncmark(i16 %N)`
+Produces an asyncmark in every stage that the mask `K` does not *omit*, and
+appends it to the current sequence of each. The sequences of the stages 
+*omitted* by `K` are unaffected. `K` must be a constant stage mask.
 
-Ensures that the length of the current sequence is at most `N` by removing
-asyncmarks from the start of the sequence if it is more than `N`.
+### `@llvm.amdgcn.wait.asyncmark(i16 %N, i32 %K)`
+
+For every stage that the mask `K` does not *ignore*, ensures that the length
+of the current sequence of that stage is at most `N` by removing asyncmarks
+from the start of that sequence if it is more than `N`. The sequences of the 
+stages *not ignored* by `K` are unaffected. `K` must be a constant stage mask.
 
 ### Completion of Asyncmarks
 
-An `asyncmark()` operation `X` that produces an asyncmark `M` is
-*completed-at* a `wait.asyncmark()` operation `Y` in the same function body
-if:
+An asyncmark `M`, produced by an `asyncmark` operation `X`, is *completed-at* 
+a `wait.asyncmark()` operation `Y` with mask `B` in the same function body if:
 
 - `X` is *program-ordered* before `Y`, and
-- `M` is not in the current sequence at any operation `Z` that immediately
-  follows `Y` in *program-order*.
+- `B` does not ignore the stage of the sequence that `M` belongs to, and
----------------
RyanRio wrote:

Hm if we do that I think it's difficult if not impossible to reason about what wait asyncmark actually does to that virtual sequence. Ie does it go through it and just "update" the mask to not include the stage, until there are N updateable marks left?

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