[Mlir-commits] [clang] [llvm] [mlir] [AMDGPU] Implement AsyncMark Stages (PR #220442)

[Fabian Ritter]

================
@@ -16,31 +16,86 @@ 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
+
+`asyncmark` intrinsics take a *stage mask*: an 11-bit value in which a set bit
+names a stage.
+
+The zero mask, which sets no bit, is given the special meaning "every stage".
+This ensures if new stages are added that programs will continue to wait on all
+stages, if that was their intention.
+
+Bits not specified in the table above are reserved for future use. It is not an
+error to set them, but it could mean you have more conservative waits than
+necessary when the future stages are added.
+
+Setting a bit that is neither listed nor reserved is an error, and is rejected
+by the IR verifier.
+
+[Informational note --- Since the stages are independent, any mix of masks is
+valid. A program is easiest to reason about by considering the sequence of each
+stage separately.]
+
 ### 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 naming multiple stages).
 
-Produces an asyncmark and appends it to the current sequence.
+### `@llvm.amdgcn.asyncmark(i32 %K)`
 
-### `@llvm.amdgcn.wait.asyncmark(i16 %N)`
+Appends an asyncmark to the current sequence of every stage named in the mask
+`K`. The sequences of the stages that `K` does not name 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 named by the mask `K`, 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 longer than `N`. The sequences of the stages that `K`
+does not name are unaffected. `K` must be a constant stage mask.
+
+[Informational note --- Typically, when removing an asyncmark, this intrinsic
+*waits* for the async operations corresponding to that asyncmark to complete.]
 
 ### 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` names the stage of the sequence that `M` belongs to, and
----------------
ritter-x2a wrote:

What exactly is the intended model/wording here? 
(a) `llvm.amdgcn.asyncmark` produces a single asyncmark and puts it into `N` sequences,  or
(b) `llvm.amdgcn.asyncmark` produces `N` asyncmarks and puts one into each of the `N` relevant sequences?

The semantics may end up equivalent, but I think we should pick a way to talk about it and stick with it.
The text above says "Appends an asyncmark to the current sequence of every stage...", which sounds ambiguous to me, and this line says "`B` names the stage of the sequence that `M` belongs to...", which matches option (b). What Krzysztof wrote matches option (a).

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