Experimental Explicit Stream Annotation

[chaserileyroberts]

This is the first PR that is intended to support explicit stream annotations for GPU runtimes.

Why do we want/need this?

There are a few optimizations that are possible with existing hardware, but they are difficult to generate from XLA. The intention is that by allowing explicit annotation of what stream a subcomputation should run with, we can allow users to define their own stream assignment strategies

1. Communication-Communication overlap

Certain configurations of parallelization that would allow for overlapping all-gather and all-reduce operations on independent networking hardware (i.e., one operation is running on NVLinks and the other exclusively using IB). Right now there is no way to do this explicitly in JAX.

2. Compute-Compute overlap

Certain kernels do not utilize all of the SMs for the full duration of their computation, leaving some idling. Other independent compute kernels could utilize these SMs for better e2e performance. We do this already for some of our collective matmul implementations, but there is currently no way to do this explicitly in JAX.

What are the code changes?

• Added new experimental flag xla_gpu_experimental_stream_annotation
• Adding an exception to CallInliner.
• Set operation_queue_id and add wrapped async operations
• At runtime, the execution stream assignment will use the stream annotation and queue kernels accordingly.

Current known limitations

• Non-inlined calls aren't well supported by all passes. For example, a single Add will not correctly lower to a fusion.

[golechwierowicz]

Can you split the PR into:

1. Adding a new experimental flag + exception to call inliner.
2. Changes to stream_attribute_annotator + test.
3. Rest of the changes (execution stream assignment)

[chaserileyroberts]

Yes I will do that.