Add TMA support for circular buffering pass

[rdspring1]

Summary

This PR adds support for TMA circular buffering. It is stacked on #2824 and #2825.
Tracking branch: #2773

Description

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

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:

constexpr int64_t warp_size = 32;
bool first_warp = threadIdx.x < warp_size && threadIdx.y == 0 && threadIdx.z == 0;

Pre-Prologue loop:

__shared__ __mbarrier_t barriers[num_stages];
__shared__ __mbarrier_token_t tokens[num_stages];
if (first_warp && hopper::electSync()) {
  for (int64_t loop_index : irange(stages)) {
    mbarrier_init(mbarrier[loop_index], number_of_arrival_threads);
  }
}

Prologue loop:

for (int64_t loop_index : irange(stages-1)) {
  if (first_warp && hopper::electSync()) {
    tokens[loop_index] = mbarrier::arriveExpectTx(mbarrier[loop_index]);
    cpAsyncBulk(mbarriers[loop_index], ...);
  } else {
    token[load_stage] = mbarrier::arrive(mbarrier[load_stage]);
  }
}

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 (first_warp && hopper::electSync()) {
    token[load_stage] =
      mbarrier::arriveExpectTx(mbarrier[load_stage], expected_transaction_size);
    cpAsyncBulk(mbarrier[load_stage], ...);
  } else {
    token[load_stage] = mbarrier::arrive(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 (first_warp && hopper::electSync()) {
  for (int64_t loop_index : irange(stages)) {
    mbarrier_inval(mbarrier[loop_index]);
  }
}

Testing Setup

• 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

[csarofeen]

Awesome, detailed PR description. Thank you.

[jacobhinkle]

Just some minor comments from a first pass. I haven't looked at tests yet.

[jacobhinkle]

Nice improvement to the comment at line 34.

[zasdfgbnm]

Do we need a separate PredicateType::ElectSync predicate type? Should we just use whatever the original predicate type it has, and if the conditional happen to be tidx == 0 && tidy == 0 && tidz == 0, we convert it to the elec_sync() conditional? What do you think? @naoyam

[rdspring1]

FYI, I separated the ElectSync PredicateType changes into #2923.

[rdspring1]

!build