MI450 WMMA and ds_load Scheduling via a DAG Mutation

Target: AMD GFX1250 / MI450, CoExecScheduler (pre-RA) · PR · Artifact

TLDR

We noticed that AMD's LLVM pre-RA scheduler, run on MI450, tends to bunch many ds_loads (LDS load instructions) at the top of a GEMM's hot loop. Upon further inspection of the CoExecScheduler, it was determined that because the scheduler greedily schedules instructions every time one is available, there are moments where the only instructions available (and therefore scheduled) are ds_loads for data consumed much, much later.

This does two bad things at once:

• it forces the WMMAs to sit behind long s_wait_dscnt stalls
• it inflates register pressure due to the unnecessarily increased live ranges

The DAG mutation in AMDGPUWMMASchedule.cpp adds hard constraints on where these loads may sit relative to the WMMAs that consume them. This forces the loads to be spread out, preventing the scheduler from greedily scheduling them far too early. The mutation is also (statically and naively) aware of register pressure, so these constraints form a window upon which each load can be scheduled without causing excessive spills.

Why do we want the DAG mutation?

To understand why the mutation is necessary, we should first look at the pre-RA scheduler. We should understand its strengths, and more importantly its weaknesses.

The CoExecScheduler

The DAG mutation is largely independent of any particular scheduler. There is an intention to switch the AMDGPU backend to the CoExecScheduler in the future, so that's what the DAG mutation applies to. CoExec is a pre-RA (pre register allocation) machine scheduler strategy for GFX1250; it replaces the generic bottom-up/top-down instruction scheduler run before register allocation. It is selected when amdgpu-sched-strategy == "coexec", is gated to GFX1250 only, and is assembled in AMDGPUCoExecSchedStrategy.cpp as a GCNScheduleDAGMILive driven by AMDGPUCoExecSchedStrategy plus two DAG mutations (three if we count ours).

Co-execution windows

The whole point of CoExec is that multi-cycle instructions (WMMAs, TRANS, multi-cycle VALU) can be scheduled alongside other instructions that run in parallel. Each multi-cycle op exposes a window of slots, and each slot has rules about which instruction types may legally and preferably co-execute there (the former is for correctness and thus cannot be violated, while the latter is for performance and can be violated).

How each ready instruction is chosen

Everything goes through AMDGPUCoExecSchedStrategy::tryCandidateCoexec. Given two candidates, it applies tiebreaking heuristics in order of how impactful the heuristic is to the efficiency of the generated program.

The following shows how pickNodeFromQueue runs. One pass is made over the whole queue available instructions (instructions that can be scheduled that cycle). Each instruction in the queue is compared against the current winning instruction. The loser of the comparison is discarded, which means it wont be scheduled for this cycle. Once the whole queue has been compared, the winner is chosen to be scheduled. Click any of the heuristics to read what it does.

The issue with this scheduler

This scheduler is greedy, like other schedulers. Whenever an instruction is available to be scheduled, an instruction will be scheduled. The problem is that there are moments when the only available instructions are ds_loads that will not be consumed for a very long time. None of the heuristics can help with this, because CoExec has no way to reject every available instruction and schedule nothing instead. Consequently, the loads are scheduled far too early, simply because the scheduler ran out of anything else that could (and most often would) win.

An example run on a real kernel

The following is a snippet of a CoExecScheduler run; the buttons on the figure switch between the mxfp GEMM and the f16 GEMM. The x-axis is the order of scheduled WMMAs (so we are looking at the live ranges of loads relative to the position of WMMA instructions in the GEMM's hot loop). Each ds_load is also labeled with the first WMMA that consumes its data. For the sake of brevity, you can imagine that in between these loads being scheduled, so are some WMMAs that become available (which is why the "now" line moves across the x-axis after all). But this is not enough to keep the available queue saturated with instructions other than these loads. Whenever some other instruction is added to the available queue, they are quickly scheduled since they win out in the heuristics, leaving the queue with only these load instructions - and these load instructions won't be consumed until much later, so scheduling them creates really long live ranges!

TRITON_HIP_USE_COEXEC_SCHEDULER=1 \
python third_party/amd/python/examples/gluon/mxfp_gemm_gfx1250.py \
  -M 1024 -N 2048 -K 4096 \
  -BM 256 -BN 256 -BK 256 \
  --num_warps 4 --num_buffers 2 \
  --dtype_a float8_e4m3 --dtype_b float8_e4m3 \
  --scale_preshuffled --schedule sliceNK

TRITON_HIP_USE_COEXEC_SCHEDULER=1 \
python third_party/amd/python/examples/gluon/f16_gemm_gfx1250.py \
  -M 1024 -N 1024 -K 2048 \
  --block_m 256 --block_n 128 --block_k 256 \
  --num-warps 4 --num-buffers 3 \
  --prefetch-lds --prefetch-l2-distance 0 \
  --single-warp-schedule

while (Available.empty())
    bumpCycle();

The DAG mutation

The DAG mutation runs after the dependency DAG is built but before tryCandidateCoexec starts scheduling. The mutation makes four kinds of edits to the DAG. Three add edges, and the fourth corrects latencies on edges that already exists.

How the DAG mutation calculates each load fragment's window

The DAG mutation starts with every load fragment's window being as short and as late as possible while still having the load finish by the time its earliest consumer reads it. Based off this schedule, a live VGPR histogram is calculated, with its peak serving as the VGPR "budget" limit. Using this budget, the windows are eased earlier until it can't without violating the budget. This is to give the CoExecScheduler breathing room to make good scheduling decisions, while preventing it from scheduling the loads too early.

An example run on a real kernel, with the DAG mutation

With these added edges forming a window for every ds_load fragment, the same kernels where the CoExecScheduler was often forced to schedule ds_loads far too early instead has an empty available queue. Note how the ds_loads now have a WMMA it must follow first. That is the result of being assigned an incoming edge from an earlier WMMA. It forces said ds_loads from being put in the available queue early. This means that the CoExecScheduler is able to schedule nothing, as it finds itself with an empty queue rather than a queue with instructions that really shouldn't be scheduled at that moment.

while (Available.empty())
    bumpCycle();

Live range reduction and spreading out waits

The following chart shows how the DAG mutation affects the live ranges (and consequently, the VGPR usage) of the ds_load fragments, as well as how it affects the s_wait_dscnts. As shown, the live ranges are significantly reduced and the waits are spread out, which allows for better latency hiding! Click any of the load fragments or the waits to see where it is in the assembly code.

Results

The following are some metrics of the kernels we've been examining, determined by running the kernels on an AMD simulator of the MI450.

Conclusion

Based on the data, you can see that the (albeit minute amount of) kernels that have been run against this DAG mutation either stay relatively the same or improve massively. Some future work:

• To ensure the robustness of this pass before it’s merged, we’ll need to run this against a wider corpus of kernels.
• Some of the code in the DAG mutation, like FragInfo, may need to be more generalized.
• I ran a version of this DAG mutation that only kept the edges between an earlier WMMA → ds_load to ensure that the ds_loads are not scheduled too early and found that the results were very similar to what we see here. This indicates that many of the edges added are not entirely necessary because the CoExecScheduler handles them fine on its own. However, further testing on a wider corpus of kernels is needed to prove this with certainty. Note, setting the correct latency on the edge between a ds_load and its earliest WMMA consumer will need to be kept as well, as the CoExecScheduler provides a different latency than the default. So, if this DAG mutation is to be slimmed down, this would need to be kept as well.
• Ultimately, the CoExecScheduler itself should be changed to have non-greedy register pressure awareness and the ability to reject all instructions in its available queue and schedule nothing instead. This would make the DAG mutation obselete, except for the aforementioned correction of the latency on edges between ds_loads and their earliest consuming WMMAs.