[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 | `ASYNC_LDS_STORE` | async stores from LDS |
+| 5 | `UNFORMATTED_BUFFER_GLOBAL_LOAD` | unformatted buffer and 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; marking an empty stage is simply a no-op.
+
+### Stage Masks
+
+Both intrinsics take a *stage mask*: an 11-bit value in which a set bit means
+"do not participate". An asyncmark leaves out the stages its mask names, and a
+wait disregards them. The mask `0` therefore names no stage and so covers all of
+them.
+
+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. Setting a bit above 10 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`.
+
+The sequences of distinct stages are independent: appending to one does not
+affect the length or contents of any other.
+
+### `@llvm.amdgcn.asyncmark(i32 %K)`
 
-### `@llvm.amdgcn.asyncmark()`
+Produces an asyncmark in every stage that the mask `K` does not name, and
+appends it to the current sequence of each. The sequences of the stages named by
+`K` are unaffected. `K` must be a constant stage mask.
 
-Produces an asyncmark and appends it to the current sequence.
+Each sequence receives its own asyncmark, and those asyncmarks are then removed
+independently by later waits.
 
-### `@llvm.amdgcn.wait.asyncmark(i16 %N)`
+### `@llvm.amdgcn.wait.asyncmark(i16 %N, i32 %K)`
 
-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`.
+For every stage that the mask `K` does not name, 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 named
+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:
----------------
RyanRio wrote:

I dont really understand what this is trying to say @ssahasra .

X -> M
...
Y
other-operations...

Is this just observing the side effect that completed-at implies the mark was removed from the sequence? Feels self-fulfilling.

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