[Tracking] TMA Circular Buffering

[rdspring1]

Summary

This PR adds support for TMA circular buffering. It is based on #1484.

Description

• In the allocation pass, create shared memory allocations and operations around LoadStoreOp expression.
• In the circular buffer pass, clone operations to create the pre-prologue, prologue, main, epilogue, and post-epilogue for-loops.
• Pre-Prologue allocates share memory and initializes mbarriers.
• Prologue copies only the load operations.
• Main loop copies the load and computation operations and adds arrive_expected_tx for next stage and mbarrier_wait for current stage.
• Epilogue copies only the computation operations and adds mbarrier_wait for remaining stages in the pipeline.
• Post-Epilogue invalidated mbarriers.

Lowering Details

Description of changes in lowering passes.

• Prologue, Main, and Epilogue loops are created by TmaCircularBufferLoopCloner which is a child class of CircularBufferLoopCloner.
• PrePrologue and PostEpilogue loops are created by createCpAsyncBulkFixtures.
• The cuTensorMapEncodeTiled restricts the size of each box dimension to be <= 256. You need to launch multiple load operations to load larger tiles.
• We only allocate mbarriers for each stage, so the expected_transaction bytes is multiplied by the number of TMA loads per stage.
• The for-loop cloner must account for the nested for-loop structure used to launch multiple TMA loads before adding the mbarrier_wait for the stage.

Allocation Pass

// Created tokens, mbarriers, init, and inval operations in allocation pass.
for (circular_buffer_loop) {
    __shared__ tokens[num_stages];
    __shared__ mbarrier[num_stages];
    init(mbarrier);
    cp.async.bulk(data, mbarrier);
    inval(mbarrier);
}

Loop Structure

Description of for-loop structure for circular buffering.

Overview Circular Buffer Structure:

Pre-prologue loop:

• Allocate shared memory for mbarriers and mbarrier tokens
• Initialize mbarrier for all stages

Prologue loop:

• if selected_thread:
  • Issue cp async bulks for all but last stage

Main loop:

• if selected_thread:
  • Issue next cp async bulk for available stage
• All threads wait until tma operation arrives
• Copy body without
  • shared memory allocations
  • mbarrier_init exprs
  • mbarrier_inval exprs

Epilogue loop:

• All threads wait until tma operation arrives
• Copy body without
  • shared memory allocations
  • issuing cp async bulk operations
  • mbarrier_init exprs
  • mbarrier_inval exprs

Post-epilogue loop:

• if selected_thread:
• Invalidated mbarrier for all stages

Detailed Pseudo-Code:

Pre-Prologue loop:

__shared__ __mbarrier_t barriers[num_stages];
__shared__ __mbarrier_token_t tokens[num_stages];
if (threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0) {
  for (int64_t loop_index : irange(stages)) {
    mbarrier_init(mbarrier[loop_index], number_of_arrival_threads);
  }
}

Prologue loop:

if (threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0) {
  for (int64_t loop_index : irange(stages-1)) {
    tokens[loop_index] = mbarrier::arriveExpectTx(mbarrier[loop_index]);
    cpAsyncBulk(mbarriers[loop_index], ...);
  }
}

Main loop:

for (int64_t loop_index : irange(N-(stages-1))) {
  current_stage = loop_index % stage_depth
  load_stage = (loop_index + (stage_depth - 1)) % stage_depth)
  if (threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0) {
    token[load_stage] =
      mbarrier::arriveExpectTx(mbarrier[load_stage]);
    cpAsyncBulk(mbarrier[load_stage], ...);
  }
  mbarrier::wait(token[current_stage]);

  // Clone remaining operations
}

Epilogue loop:

for (int64_t loop_index : irange(N-(stages-1), N)) {
  current_stage = loop_index % stage_depth
  mbarrier::wait(token[current_stage]);

  // Clone remaining operations
}

Post-Epilogue loop:

if (threadIdx.x == 0 && threadIdx.y == 0 && threadIdx.z == 0) {
  for (int64_t loop_index : irange(stages)) {
    mbarrier_inval(mbarrier[loop_index]);
  }
}

Testing

• 2 to 4 pipeline stages.
• (128, 500, 1024) outer dimension.
• (128, 1024) inner dimension.

1. Single Dim including Unroll and Unswitch parallelizations.
2. Multiple Dim
3. Pointwise
4. Reduction
5. InnerPersistent
6. Matmul

Future PRs

• Replace TDX == 0 && TDY == 0 && TDZ == 0 with Hopper::elect_sync PTX instruction.
• Known Issue: Some LoadStore operations add additional TDY == 0 predicates that conflict with Hopper::elect_sync.

[rdspring1]

!build

[rdspring1]

!build