PR #41079: [XLA:GPU] Migrate MemzeroCmd into MemzeroThunk#116330
Draft
copybara-service[bot] wants to merge 1 commit intomasterfrom
Draft
PR #41079: [XLA:GPU] Migrate MemzeroCmd into MemzeroThunk#116330copybara-service[bot] wants to merge 1 commit intomasterfrom
copybara-service[bot] wants to merge 1 commit intomasterfrom
Conversation
Imported from GitHub PR openxla/xla#41079 Description Follows the same pattern as #40722 ([XLA:GPU] Migrate Memset32Cmd into Memset32BitValueThunk): eliminate the standalone MemzeroCmd wrapper class by making MemzeroThunk directly implement Command and provide its own Record() method. Before MemzeroThunk (a plain Thunk subclass) was converted by the command buffer emitter into a heap-allocated MemzeroCmd object via Convert<MemzeroThunk>(). This meant the same logical operation was split across two classes with duplicated state. After MemzeroThunk inherits from Command (using the protected three-argument constructor to preserve Kind::kMemzero), implements Record() directly, and is emplaced into the command sequence via static_cast — the same pattern as Memset32BitValueThunk, GemmThunk, KernelThunk, and DeviceToDeviceCopyThunk. MemzeroCmd is deleted entirely. Changes - memset_thunk.h/cc: MemzeroThunk now extends Command; Record() added - command_buffer_cmd.h/cc: MemzeroCmd class removed - command_buffer_cmd_emitter.cc: Convert(const MemzeroThunk&) helper removed; kMemzero case uses static_cast append - memset_thunk_test.cc: Added RecordCommandBuffer and RecordCommandBufferUpdate GPU tests for MemzeroThunk - command_buffer_thunk_test.cc, command_buffer_cmd_test.cc: Updated call sites Testing New MemzeroThunkTest::RecordCommandBuffer and MemzeroThunkTest::RecordCommandBufferUpdate tests cover the create and update paths. Copybara import of the project: -- 245fff2c7f3d4a2490391738e78ff6291589bec0 by Shawn Wang <shawnw@nvidia.com>: [XLA:GPU] Migrate MemzeroCmd into MemzeroThunk Follows the same pattern as the recent Memset32Cmd migration: make MemzeroThunk directly inherit Command and implement Record(), removing the separate MemzeroCmd wrapper class entirely. Merging this change closes #41079 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#41079 from shawnwang18:shawnw/migrate_memset_command_to_thunk 245fff2c7f3d4a2490391738e78ff6291589bec0 PiperOrigin-RevId: 902431513
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 join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
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.
PR #41079: [XLA:GPU] Migrate MemzeroCmd into MemzeroThunk
Imported from GitHub PR openxla/xla#41079
Description
Follows the same pattern as #40722 ([XLA:GPU] Migrate Memset32Cmd into Memset32BitValueThunk): eliminate the standalone MemzeroCmd wrapper class by making MemzeroThunk directly implement Command and provide its own Record() method.
Before
MemzeroThunk (a plain Thunk subclass) was converted by the command buffer emitter into a heap-allocated MemzeroCmd object via Convert(). This meant the same logical operation was split across two classes with duplicated state.
After
MemzeroThunk inherits from Command (using the protected three-argument constructor to preserve Kind::kMemzero), implements Record() directly, and is emplaced into the command sequence via static_cast — the same pattern as Memset32BitValueThunk, GemmThunk, KernelThunk, and DeviceToDeviceCopyThunk. MemzeroCmd is deleted entirely.
Changes
Testing
New MemzeroThunkTest::RecordCommandBuffer and MemzeroThunkTest::RecordCommandBufferUpdate tests cover the create and update paths.
Copybara import of the project:
--
245fff2c7f3d4a2490391738e78ff6291589bec0 by Shawn Wang shawnw@nvidia.com:
[XLA:GPU] Migrate MemzeroCmd into MemzeroThunk
Follows the same pattern as the recent Memset32Cmd migration: make
MemzeroThunk directly inherit Command and implement Record(), removing
the separate MemzeroCmd wrapper class entirely.
Merging this change closes #41079
FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#41079 from shawnwang18:shawnw/migrate_memset_command_to_thunk 245fff2c7f3d4a2490391738e78ff6291589bec0