Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[XLA:GPU] Support partially pipelined async send recv ops #17446

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

[XLA:GPU] Support partially pipelined async send recv ops #17446

wants to merge 1 commit into from
+452 −39

Commits on Oct 2, 2024

  1. [XLA:GPU] Support partially pipelined async send recv ops 

    This is needed for pipeline parallelism on GPU where the send/recv operations
are issued in one loop iteration and completed in the next. The same buffer
must be alive throughout the process and no copies can be inserted.

Avoid copies for these partially pipelined async send/recv ops. Insert the
required copies and controlflow constraints on the send/recv ops separately.
This is to ensure that the live times of the buffers do not overlap.

Send: For send, a copy is inserted on the operand, starting a new live range.
By enforcing this copy after the corresponding send/done, buffer live times are
disjoint.

Recv: For recv, a copy is inserted after recv-done, ending the live time of the
buffer. Bt enforcing the copy to be before the corresponding recv. buffer live
times are disjoint.
PiperOrigin-RevId: 676075106
    @frgossen @Google-ML-Automation
    frgossen authored and Google-ML-Automation committed Oct 2, 2024
    Configuration menu
    Copy the full SHA
    b886ebe View commit details
    Browse the repository at this point in the history