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:TPU] Support output streaming and refactor TryOutputStreaming into a bottoms-up approach. #66989

Merged
merged 1 commit into from
May 10, 2024
Merged

[XLA:TPU] Support output streaming and refactor TryOutputStreaming into a bottoms-up approach. #66989

merged 1 commit into from
May 10, 2024

Commits on May 10, 2024

  1. [XLA:TPU] Support output streaming and refactor TryOutputStreaming in… 

    …to a bottoms-up approach.

Previously, output streaming took a top-down approach which indiscriminately checks if a MoveToHost custom call would trace down to an output marked with host memory space. This did not work when a dynamic-update-slice existed between the MTH call and the output. This CL fixes this problem by handling output streaming before other MTH calls, while also improving efficiency with the bottoms-up approach so we only trace a single path in the graph.

PiperOrigin-RevId: 632318740
    @jvstokes @tensorflower-gardener
    jvstokes authored and tensorflower-gardener committed May 10, 2024
    Configuration menu
    Copy the full SHA
    972a7e1 View commit details
    Browse the repository at this point in the history