AMDGPU Asynchronous Operations

Introduction

Asynchronous operations are operations whose completion is not tracked internally by the compiler. A thread that initiates one or more async operations can use asyncmarks to track their completion.

• Most DMA operations are asynchronous.

Asyncmarks

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 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 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, even though multiple sequences may be appended to with a single call to asyncmark (by naming multiple stages).

@llvm.amdgcn.asyncmark(i32 %K)

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.

@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.

Completion of Asyncmarks

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

• B names the stage of the sequence that M belongs to, and

• M is not in the current sequence at any operation Z that immediately follows Y in program-order.

Completion of Async Operations

An async operation executes outside the thread that initiated it, i.e., it is not related in program-order with any other operations from that thread. But the thread can use an asyncmark to ensure that the async operation is completed-at some later operation.

An async operation A initiated-by an instruction I is completed-at some wait.asyncmark() operation Y with stage mask B if there exists an asyncmark() operation X with mask C:

• I is program-ordered before X, and

• I belongs to a stage in C, and

• X is completed-at Y.

happens-before

When an instruction I initiates an async operation A, I happens-before A.

If A is completed-at a wait.asyncmark() operation Y, then A happens-before Y.

Examples

Most examples below use the zero mask, so every asyncmark tracks every async operation and there is only one sequence to reason about. The interleaved stages example shows what narrower masks add.

Uneven blocks of async operations

void foo(global int *g, local int *l) {
  // first block
  async_load_to_lds(l, g);
  async_load_to_lds(l, g);
  async_load_to_lds(l, g);
  asyncmark(0);

  // second block; longer
  async_load_to_lds(l, g);
  async_load_to_lds(l, g);
  async_load_to_lds(l, g);
  async_load_to_lds(l, g);
  async_load_to_lds(l, g);
  asyncmark(0);

  // third block; shorter
  async_load_to_lds(l, g);
  async_load_to_lds(l, g);
  asyncmark(0);

  // Wait for first block
  wait.asyncmark(2, 0);
}

Software pipeline

void foo(global int *g, local int *l) {
  // first block
  asyncmark(0);

  // second block
  asyncmark(0);

  // third block
  asyncmark(0);

  for (;;) {
    wait.asyncmark(2, 0);
    // use data

    // next block
    asyncmark(0);
  }

  // flush one block
  wait.asyncmark(2, 0);

  // flush one more block
  wait.asyncmark(1, 0);

  // flush last block
  wait.asyncmark(0, 0);
}

Ordinary function call

extern void bar(); // may or may not initiate async operations

void foo(global int *g, local int *l) {
    // first block
    asyncmark(0);

    // second block
    asyncmark(0);

    // function call
    bar();

    // third block
    asyncmark(0);

    // wait for the second block
    wait.asyncmark(1, 0);

    // wait for the third block, including bar()
    wait.asyncmark(0, 0);
}

Interleaved stages

Two kinds of async operation are in flight at once. Each is marked with a mask naming only its own stage, so each can be waited for without waiting for the other.

void foo(global int *g, local int *l, tensor_desc t) {
  // Start a long tensor load and mark it in the TENSOR stage alone.
  tensor_load_to_lds(l, t);
  asyncmark(TENSOR);

  // Start a short LDS load and mark it in its own stage alone.
  async_load_to_lds(l, g);
  asyncmark(GLOBAL_LOAD_ASYNC_TO_LDS);

  // Wait for the LDS load only. The tensor load is still in flight: this wait
  // does not name TENSOR, so it neither counts nor removes that asyncmark.
  wait.asyncmark(0, GLOBAL_LOAD_ASYNC_TO_LDS);

  // perform synchronization / use the data loaded by async_load_to_lds

  // Now wait for the tensor load.
  wait.asyncmark(0, TENSOR);
}

Had both marks used the zero mask, the first wait would have drained the tensor load as well, and the overlap would have been lost.

The same applies to two stages tracked by one hardware counter. Here both stages use the async counter, but the sequences are still separate:

void foo(global int *g, local int *l) {
  async_load_to_lds(l, g);
  asyncmark(GLOBAL_LOAD_ASYNC_TO_LDS);    // X

  async_store_from_lds(g, l);
  asyncmark(GLOBAL_STORE_ASYNC_FROM_LDS); // Y

  // Completes X. Y is in a different sequence and is not completed here, even
  // though both stages are tracked by the same counter.
  wait.asyncmark(0, GLOBAL_LOAD_ASYNC_TO_LDS);
}

A mask need not name just one stage. An asyncmark whose mask names several is appended to each of their sequences, and a wait whose mask names several trims each of them:

void foo(global int *g, local int *l, tensor_desc t) {
  tensor_load_to_lds(l, t);
  async_load_to_lds(l, g);

  // One asyncmark, appended to both sequences.
  asyncmark(TENSOR | GLOBAL_LOAD_ASYNC_TO_LDS);

  // Removes the mark in GLOBAL_LOAD_ASYNC_TO_LDS. The mark in TENSOR is a
  // separate asyncmark and stays where it is.
  wait.asyncmark(0, GLOBAL_LOAD_ASYNC_TO_LDS);

  // Removes the mark in TENSOR.
  wait.asyncmark(0, TENSOR);
}

Implementation notes

[This section is informational.]

Function Calls

In general, at a function call, if the caller uses sufficient waits to track its own async operations, the actions performed by the callee cannot affect correctness. But inlining such a call may result in redundant waits.

void foo() {
  ...
  asyncmark(0);         // X
  ...                   // no wait.asyncmark()
}

void bar() {
  asyncmark(0);         // B
  asyncmark(0);         // C
  foo();
  wait.asyncmark(1, 0); // D
}

Before inlining, it is unspecified whether X is completed-at D, while C is not completed-at D. The programmer can only rely on B being completed-at D.

void bar() {
  asyncmark(0);         // B
  asyncmark(0);         // C
  ...
  asyncmark(0);         // X
  ...                   // no wait.asyncmark()
  wait.asyncmark(1, 0); // D
}

After inlining, C is also completed-at D and X is not completed-at D.

Conversely, a wait.asyncmark call inside a callee cannot be used to track asyncmarks from the caller, since this wait.asyncmark can only observe the current sequence of the callee.

void foo() {
  ...                   // no asyncmark()
  wait.asyncmark(0, 0); // Y
  ...
}

void bar() {
  asyncmark(0);         // B
  asyncmark(0);         // C
  foo();
  wait.asyncmark(1, 0); // D
}

In the above example, it is unspecified whether B and C in bar() are completed-at Y, because they are not included in the sequence that can be examined at Y.

void bar() {
  asyncmark(0);         // B
  asyncmark(0);         // C
  ...                   // no asyncmark()
  wait.asyncmark(0, 0); // Y
  ...
  wait.asyncmark(1, 0); // D
}

After inlining, both B and C are completed-at Y.

Optimization

The implementation may eliminate/modify asyncmark/wait intrinsics in the following cases. These are just examples and not meant to be an exhaustive list.

1. An asyncmark operation which remains in a current sequence along every path that reaches the function exit.

   void foo() {
     ...
     asyncmark(0);         // X
     ...                   // no wait.asyncmark()
   }
   

   Here, the zero mask names every stage, and each independent mark remains, so the entire mark X can be eliminated. If only some sequences qualified, then the asyncmark operation could be modified to remove just those stages from its mask.

2. A wait.asyncmark which sees an empty sequence of asyncmarks along every path that reaches it.

   void foo() {
     ...                   // no asyncmark()
     wait.asyncmark(0, 0); // Y
     ...
   }
   

   Here, as with X above, every sequence qualifies, so Y can be eliminated. If only some sequences qualified, the wait.asyncmark could be modified to only wait for the relevant stages.

3. An asyncmark operation that has mask bits that don’t refer to any instruction.

   void foo() {
    ...
    global_load_async_to_lds
    asyncmark(0)
    ...
   }
   

   Here, the program can be refined to -

   void foo() {
     ...
     global_load_async_to_lds
     asyncmark(ASYNC_LOAD_LDS)
     ...
   }
   

   Since only that particular stage is actually relevant to ensuring the completion of async operations.