[gpu][ds-fusion] Add handling for offset module in ds-fusion thunk #20794

shraiysh · 2024-12-20T22:02:58Z

This patch adds support for offset modules in ds fusion thunk. It also moves the ResourceRequests structure from gpu_executable.cc to gpu_executable.h because a valid implementation of the abstract class Thunk::ResourceRequests is required for calling Thunk::Prepare().

This is split from #20332 as per request.

shraiysh · 2024-12-28T16:52:27Z

The error in the CI job is unrelated. This PR is open for review.

shraiysh · 2025-01-02T17:31:02Z

@ezhulenev @bchetioui ping for review. Thanks.

bchetioui

Thanks @shraiysh! @beckerhe suggested moving ResourceRequests to its own header and source files.

At its core, it seems like this PR does two things:

Reformulate the definitions of ResourceRequests such that they are accessible through the header file;
Add some functionality.

To be honest, from the current diff, it's not super straightforward to figure out what is what. Would you mind splitting the pure code move (that'll be a trivial review for us), and the added functionality?

bchetioui · 2025-01-10T13:50:52Z

xla/service/gpu/gpu_executable.h

@@ -53,6 +53,42 @@ limitations under the License.
 namespace xla {
 namespace gpu {

+// Shared resources required for thunk initialization and execution.
+class ResourceRequests : public Thunk::ResourceRequests {


Let's maybe create dedicated source and header files for this class.

shraiysh · 2025-01-13T22:39:45Z

xla/service/gpu/runtime/dynamic_slice_thunk.h

-  using Offset = std::variant<int64_t, BufferAllocation::Slice>;
+  // transferred to host or (3) a temporary HloModule that computes the offset
+  // with induction variable as the input.
+  using Offset = std::variant<int64_t, BufferAllocation::Slice, HloModule*>;


Please update the docstring to account for this new variant constructor.

@bchetioui Added.

shraiysh · 2025-01-13T22:40:06Z

xla/service/gpu/runtime/dynamic_slice_thunk.h

+  // `indvar_update_` is a module of the prototype `(integer[]) -> integer[]`.
+  // It takes the current value of the induction variable, and returns the next
+  // value of the induction variable.
+  std::unique_ptr<HloModule> indvar_init_, indvar_update_;


Please clearly document all the new fields.

Done. Thanks for mentioning this.

…on thunk Imported from GitHub PR #20794 This patch adds support for offset modules in ds fusion thunk. It also moves the `ResourceRequests` structure from `gpu_executable.cc` to `gpu_executable.h` because a valid implementation of the abstract class `Thunk::ResourceRequests` is required for calling `Thunk::Prepare()`. This is split from #20332 as per request. Copybara import of the project: -- 53226e9 by Shraiysh Vaishay <[email protected]>: [gpu][ds-fusion] Add handling for offset module in ds-fusion thunk This patch adds support for offset modules in ds fusion thunk. It also moves the `ResourceRequests` structure from `gpu_executable.cc` to `gpu_executable.h` because a valid implementation of the abstract class `Thunk::ResourceRequests` is required for calling `Thunk::Prepare()`. -- e554ac6 by Shraiysh Vaishay <[email protected]>: Address comments and rebase -- 4ac8efa by Shraiysh Vaishay <[email protected]>: Address comments -- 8c23e2a by Shraiysh Vaishay <[email protected]>: Addressed comments. -- efbe6bf by Shraiysh Vaishay <[email protected]>: Addressed comments Merging this change closes #20794 FUTURE_COPYBARA_INTEGRATE_REVIEW=#20794 from shraiysh:ds_fusion_thunk_changes efbe6bf PiperOrigin-RevId: 718748248

…on thunk Imported from GitHub PR openxla/xla#20794 This patch adds support for offset modules in ds fusion thunk. It also moves the `ResourceRequests` structure from `gpu_executable.cc` to `gpu_executable.h` because a valid implementation of the abstract class `Thunk::ResourceRequests` is required for calling `Thunk::Prepare()`. This is split from #20332 as per request. Copybara import of the project: -- 53226e9428dff816819c192735b9569ef3d309ea by Shraiysh Vaishay <[email protected]>: [gpu][ds-fusion] Add handling for offset module in ds-fusion thunk This patch adds support for offset modules in ds fusion thunk. It also moves the `ResourceRequests` structure from `gpu_executable.cc` to `gpu_executable.h` because a valid implementation of the abstract class `Thunk::ResourceRequests` is required for calling `Thunk::Prepare()`. -- e554ac6b09b5023082210c5100a1f1a86c1c6605 by Shraiysh Vaishay <[email protected]>: Address comments and rebase -- 4ac8efa0c9885b0f21f526072f821ace40be4043 by Shraiysh Vaishay <[email protected]>: Address comments -- 8c23e2a4e2ff66fb7361c48c9d988ebb2261bc41 by Shraiysh Vaishay <[email protected]>: Addressed comments. -- efbe6bfa6f10f1ee998683161e398ae3adc0f183 by Shraiysh Vaishay <[email protected]>: Addressed comments Merging this change closes #20794 FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#20794 from shraiysh:ds_fusion_thunk_changes efbe6bfa6f10f1ee998683161e398ae3adc0f183 PiperOrigin-RevId: 718748248