[pull] master from tensorflow:master#1671
Merged
Merged
Conversation
…/Done Imported from GitHub PR openxla/xla#39382 **NFC**: this PR doesn't change what collective operations are executed and on what streams, it should be identical order of launched operations on identical stream. The only change is migration from per-operation thunks to structured concurrency based on start/done thunk wrapping. Migrate all GPU collective thunks to use the generic `AsyncStartThunk`/`AsyncDoneThunk` infrastructure, removing the per-collective async event management and simplifying the thunk hierarchy. ### Key changes: - **Remove `CollectiveAsyncEvents`**: Collective thunks no longer manage their own async events. The `AsyncStartThunk`/`AsyncDoneThunk` pair handles all structured concurrency. - **Remove per-collective Done thunks**: `CollectiveDoneCmd`, `NvshmemCollectiveDoneCmd`, and all collective-specific done command logic removed. `AsyncDoneThunk` handles completion for all collectives. - **Simplify command buffer commands**: `CollectiveCmd` and `NvshmemCollectiveCmd` now extend `Command` directly (not `TrackedCmd`). Removed `CollectiveAsyncEvents` from command constructors and execution. - **Consolidate `Thunk::Kind` enum**: Removed all collective `*Start` and `*Done` kind variants (e.g., `kAllReduceStart`, `kAllReduceDone`), keeping only base kinds (`kAllReduce`). Added `kNvshmemAllReduce`. Updated `ThunkKindProto` with reserved entries for removed values. - **Sync collective support**: When `IsGPUSyncCollective` is true, collective thunks are emitted directly without `AsyncStartThunk` wrapping, and `EmitCollectiveAsyncDone` returns an empty sequence. Copybara import of the project: -- a6568a645bc3d32e4e313a9076cf3fc9439b6d9a by Eugene Zhulenev <ezhulenev@openxla.org>: [xla:gpu] Migrate collective thunks to generic Async Start/Done Merging this change closes #39382 PiperOrigin-RevId: 891024570
…zation since there seems to be corner cases that are not covered. Reverts 09d30f3 PiperOrigin-RevId: 891080815
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to subscribe to this conversation on GitHub.
Already have an account?
Sign in.
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
See Commits and Changes for more details.
Created by
pull[bot] (v2.0.0-alpha.4)
Can you help keep this open source service alive? 💖 Please sponsor : )