[XLA:GPU] Split command buffer VA remapping into its own allocator class - #45567

…llocator class

Imported from GitHub PR #45567

## 📝 Summary of Changes

This PR separates the command-buffer virtual-address remapping implementation from `GpuExecutableBufferAllocator` into a dedicated `GpuExecutableVaRemapAllocator` subclass. It is a structural refactor of the existing `SKIP_TEMP` behavior; no functional change is intended.

- `GpuExecutableBufferAllocator` is now the common base implementation for allocation resolution and the `ALWAYS_UPDATE` policy. It also remains the runtime fallback for `SKIP_TEMP` when VA remapping is unavailable.
- `ExecutionScope` is now polymorphic and is returned as `std::unique_ptr<ExecutionScope>`. Protected hooks such as `PrepareReservation` and `AllocateTransientBuffer` provide the extension points needed by address-policy-specific subclasses.
- The new `GpuExecutableVaRemapAllocator` owns the `SKIP_TEMP` implementation:
  - per-executor remapping state and synchronization;
  - deterministic VA reservation sizing and layout;
  - VMM-backed allocation of command-buffer-referenced temporary buffers;
  - persistent allocation indices combining constants with VA-remapped temporary allocations; and
  - synchronization of pending VMM operations during allocator teardown.
- `GpuExecutableBufferAllocator::Create` selects the base allocator for `ALWAYS_UPDATE` and the VA-remapping subclass for `SKIP_TEMP`.
- GPU executable and PJRT call sites now use the factory and pointer-owned execution scopes.
- The BUILD target includes the new source/header pair and retains the direct dependencies required by both implementations.
- VMM test logging now enables both allocator translation units so existing assertions continue to observe the VA-reservation messages after the code move.

This PR is stacked on #45130, **[XLA:GPU] Update command buffers when allocation policy changes**. The follow-up policy is implemented separately in #45568.

## 🎯 Justification

The base allocator previously mixed general buffer-allocation mechanics with one specific VMM address-stability policy. Splitting those responsibilities makes the default path easier to reason about and gives additional update policies a narrow place to implement reservation, allocation, and execution behavior without further expanding the base class.

This cleanup is primarily preparation for profiled command-buffer allocation policies. Existing GPU command-buffer workloads should retain the same `ALWAYS_UPDATE` and `SKIP_TEMP` behavior.

## 🚀 Kind of Contribution

♻️ Cleanup

## 📊 Benchmark (for Performance Improvements)

Not applicable. This is a behavior-preserving refactor and makes no performance claim.

## 🧪 Unit Tests:

- Updated `GpuExecutableTest.CommandBufferAllocationPolicy` to construct allocators through the new factory and exercise the pointer-owned `ExecutionScope`.
- Existing assertions continue to cover the persistent-allocation sets for `ALWAYS_UPDATE` and `SKIP_TEMP`, including the non-VMM fallback path.

## 🧪 Execution Tests:

- Existing `VmmTest` command-buffer tests continue to cover the two-GEMM `SKIP_TEMP` path, VA reservation reuse, and fallback when the configured allocator does not provide VMM remapping.
- `ScopedBufferAllocatorVLog` now captures logs from `gpu_executable_va_remap_allocator.cc`, preserving the execution-test assertions whose log statements moved with the implementation.
Copybara import of the project:

--
789a0b4 by Shawn Wang <[email protected]>:

Update command buffers when allocation policy changes

--
1926cba by Shawn Wang <[email protected]>:

Restrict persistent allocation policy transitions

--
3a40f5d by Shawn Wang <[email protected]>:

Fix persistent allocation policy update fallback

--
6be84cf by Shawn Wang <[email protected]>:

Clarify persistent allocation validity names

--
3aa5205 by Shawn Wang <[email protected]>:

Remove debug-only allocator mode checks

--
e640997 by Shawn Wang <[email protected]>:

[XLA:GPU] Split command buffer VA remapping into its own allocator class

Separate the command buffer VA remapping implementation (SKIP_TEMP
update mode) from GpuExecutableBufferAllocator into a new
GpuExecutableVaRemapAllocator subclass in
gpu_executable_va_remap_allocator.{h,cc}:

- GpuExecutableBufferAllocator becomes the base class implementing the
  ALWAYS_UPDATE update mode (constants-only persistent allocations) and
  the runtime fallback behavior for SKIP_TEMP when VA remapping is
  unavailable. Its ExecutionScope is now polymorphic and exposes
  protected hooks (PrepareReservation, AllocateTransientBuffer) at the
  points where an allocation-address policy can intervene.
- GpuExecutableVaRemapAllocator implements SKIP_TEMP: the per-executor
  Remapping state, VA reservation layout, and temp-buffer remapping
  move to the new file.
- A GpuExecutableBufferAllocator::Create factory picks the class from
  the update mode; CreateExecutionScope now returns
  std::unique_ptr<ExecutionScope>.

No functional change intended.

--
5665403 by Shawn Wang <[email protected]>:

Fix VA remap allocator CI failures

--
7da1c61 by Shawn Wang <[email protected]>:

Restore platform-restricted GPU PJRT deps

Merging this change closes #45567

Reverts f212066

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45567 from shawnwang18:shawnw/cmd_buffer_va_remap_allocator 7da1c61
PiperOrigin-RevId: 947940077

…llocator class

Imported from GitHub PR openxla/xla#45567

## 📝 Summary of Changes

This PR separates the command-buffer virtual-address remapping implementation from `GpuExecutableBufferAllocator` into a dedicated `GpuExecutableVaRemapAllocator` subclass. It is a structural refactor of the existing `SKIP_TEMP` behavior; no functional change is intended.

- `GpuExecutableBufferAllocator` is now the common base implementation for allocation resolution and the `ALWAYS_UPDATE` policy. It also remains the runtime fallback for `SKIP_TEMP` when VA remapping is unavailable.
- `ExecutionScope` is now polymorphic and is returned as `std::unique_ptr<ExecutionScope>`. Protected hooks such as `PrepareReservation` and `AllocateTransientBuffer` provide the extension points needed by address-policy-specific subclasses.
- The new `GpuExecutableVaRemapAllocator` owns the `SKIP_TEMP` implementation:
  - per-executor remapping state and synchronization;
  - deterministic VA reservation sizing and layout;
  - VMM-backed allocation of command-buffer-referenced temporary buffers;
  - persistent allocation indices combining constants with VA-remapped temporary allocations; and
  - synchronization of pending VMM operations during allocator teardown.
- `GpuExecutableBufferAllocator::Create` selects the base allocator for `ALWAYS_UPDATE` and the VA-remapping subclass for `SKIP_TEMP`.
- GPU executable and PJRT call sites now use the factory and pointer-owned execution scopes.
- The BUILD target includes the new source/header pair and retains the direct dependencies required by both implementations.
- VMM test logging now enables both allocator translation units so existing assertions continue to observe the VA-reservation messages after the code move.

This PR is stacked on #45130, **[XLA:GPU] Update command buffers when allocation policy changes**. The follow-up policy is implemented separately in #45568.

## 🎯 Justification

The base allocator previously mixed general buffer-allocation mechanics with one specific VMM address-stability policy. Splitting those responsibilities makes the default path easier to reason about and gives additional update policies a narrow place to implement reservation, allocation, and execution behavior without further expanding the base class.

This cleanup is primarily preparation for profiled command-buffer allocation policies. Existing GPU command-buffer workloads should retain the same `ALWAYS_UPDATE` and `SKIP_TEMP` behavior.

## 🚀 Kind of Contribution

♻️ Cleanup

## 📊 Benchmark (for Performance Improvements)

Not applicable. This is a behavior-preserving refactor and makes no performance claim.

## 🧪 Unit Tests:

- Updated `GpuExecutableTest.CommandBufferAllocationPolicy` to construct allocators through the new factory and exercise the pointer-owned `ExecutionScope`.
- Existing assertions continue to cover the persistent-allocation sets for `ALWAYS_UPDATE` and `SKIP_TEMP`, including the non-VMM fallback path.

## 🧪 Execution Tests:

- Existing `VmmTest` command-buffer tests continue to cover the two-GEMM `SKIP_TEMP` path, VA reservation reuse, and fallback when the configured allocator does not provide VMM remapping.
- `ScopedBufferAllocatorVLog` now captures logs from `gpu_executable_va_remap_allocator.cc`, preserving the execution-test assertions whose log statements moved with the implementation.
Copybara import of the project:

--
789a0b470b867a92fa7535d357f5557d601aa089 by Shawn Wang <[email protected]>:

Update command buffers when allocation policy changes

--
1926cbac84209d9b7dadd77133457cf8ed26f9eb by Shawn Wang <[email protected]>:

Restrict persistent allocation policy transitions

--
3a40f5dbcd14387534f8cc3ce73aeef57ff979ac by Shawn Wang <[email protected]>:

Fix persistent allocation policy update fallback

--
6be84cf6bb617265100b2ddb3261b3cd34358f20 by Shawn Wang <[email protected]>:

Clarify persistent allocation validity names

--
3aa5205db39fc9f109fea4cd868b02dbfab3019c by Shawn Wang <[email protected]>:

Remove debug-only allocator mode checks

--
e640997b3c02d5f5fad220431965d63374461ea0 by Shawn Wang <[email protected]>:

[XLA:GPU] Split command buffer VA remapping into its own allocator class

Separate the command buffer VA remapping implementation (SKIP_TEMP
update mode) from GpuExecutableBufferAllocator into a new
GpuExecutableVaRemapAllocator subclass in
gpu_executable_va_remap_allocator.{h,cc}:

- GpuExecutableBufferAllocator becomes the base class implementing the
  ALWAYS_UPDATE update mode (constants-only persistent allocations) and
  the runtime fallback behavior for SKIP_TEMP when VA remapping is
  unavailable. Its ExecutionScope is now polymorphic and exposes
  protected hooks (PrepareReservation, AllocateTransientBuffer) at the
  points where an allocation-address policy can intervene.
- GpuExecutableVaRemapAllocator implements SKIP_TEMP: the per-executor
  Remapping state, VA reservation layout, and temp-buffer remapping
  move to the new file.
- A GpuExecutableBufferAllocator::Create factory picks the class from
  the update mode; CreateExecutionScope now returns
  std::unique_ptr<ExecutionScope>.

No functional change intended.

--
5665403db4e1b7583251618533ca9378be006765 by Shawn Wang <[email protected]>:

Fix VA remap allocator CI failures

--
7da1c618865c4a910928df8dc03f60afd7466ab2 by Shawn Wang <[email protected]>:

Restore platform-restricted GPU PJRT deps

Merging this change closes #45567

Reverts 6c3c2e3

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#45567 from shawnwang18:shawnw/cmd_buffer_va_remap_allocator 7da1c618865c4a910928df8dc03f60afd7466ab2
PiperOrigin-RevId: 947940077

…llocator class

Imported from GitHub PR openxla/xla#45567

## 📝 Summary of Changes

This PR separates the command-buffer virtual-address remapping implementation from `GpuExecutableBufferAllocator` into a dedicated `GpuExecutableVaRemapAllocator` subclass. It is a structural refactor of the existing `SKIP_TEMP` behavior; no functional change is intended.

- `GpuExecutableBufferAllocator` is now the common base implementation for allocation resolution and the `ALWAYS_UPDATE` policy. It also remains the runtime fallback for `SKIP_TEMP` when VA remapping is unavailable.
- `ExecutionScope` is now polymorphic and is returned as `std::unique_ptr<ExecutionScope>`. Protected hooks such as `PrepareReservation` and `AllocateTransientBuffer` provide the extension points needed by address-policy-specific subclasses.
- The new `GpuExecutableVaRemapAllocator` owns the `SKIP_TEMP` implementation:
  - per-executor remapping state and synchronization;
  - deterministic VA reservation sizing and layout;
  - VMM-backed allocation of command-buffer-referenced temporary buffers;
  - persistent allocation indices combining constants with VA-remapped temporary allocations; and
  - synchronization of pending VMM operations during allocator teardown.
- `GpuExecutableBufferAllocator::Create` selects the base allocator for `ALWAYS_UPDATE` and the VA-remapping subclass for `SKIP_TEMP`.
- GPU executable and PJRT call sites now use the factory and pointer-owned execution scopes.
- The BUILD target includes the new source/header pair and retains the direct dependencies required by both implementations.
- VMM test logging now enables both allocator translation units so existing assertions continue to observe the VA-reservation messages after the code move.

This PR is stacked on #45130, **[XLA:GPU] Update command buffers when allocation policy changes**. The follow-up policy is implemented separately in #45568.

## 🎯 Justification

The base allocator previously mixed general buffer-allocation mechanics with one specific VMM address-stability policy. Splitting those responsibilities makes the default path easier to reason about and gives additional update policies a narrow place to implement reservation, allocation, and execution behavior without further expanding the base class.

This cleanup is primarily preparation for profiled command-buffer allocation policies. Existing GPU command-buffer workloads should retain the same `ALWAYS_UPDATE` and `SKIP_TEMP` behavior.

## 🚀 Kind of Contribution

♻️ Cleanup

## 📊 Benchmark (for Performance Improvements)

Not applicable. This is a behavior-preserving refactor and makes no performance claim.

## 🧪 Unit Tests:

- Updated `GpuExecutableTest.CommandBufferAllocationPolicy` to construct allocators through the new factory and exercise the pointer-owned `ExecutionScope`.
- Existing assertions continue to cover the persistent-allocation sets for `ALWAYS_UPDATE` and `SKIP_TEMP`, including the non-VMM fallback path.

## 🧪 Execution Tests:

- Existing `VmmTest` command-buffer tests continue to cover the two-GEMM `SKIP_TEMP` path, VA reservation reuse, and fallback when the configured allocator does not provide VMM remapping.
- `ScopedBufferAllocatorVLog` now captures logs from `gpu_executable_va_remap_allocator.cc`, preserving the execution-test assertions whose log statements moved with the implementation.
Copybara import of the project:

--
789a0b470b867a92fa7535d357f5557d601aa089 by Shawn Wang <[email protected]>:

Update command buffers when allocation policy changes

--
1926cbac84209d9b7dadd77133457cf8ed26f9eb by Shawn Wang <[email protected]>:

Restrict persistent allocation policy transitions

--
3a40f5dbcd14387534f8cc3ce73aeef57ff979ac by Shawn Wang <[email protected]>:

Fix persistent allocation policy update fallback

--
6be84cf6bb617265100b2ddb3261b3cd34358f20 by Shawn Wang <[email protected]>:

Clarify persistent allocation validity names

--
3aa5205db39fc9f109fea4cd868b02dbfab3019c by Shawn Wang <[email protected]>:

Remove debug-only allocator mode checks

--
e640997b3c02d5f5fad220431965d63374461ea0 by Shawn Wang <[email protected]>:

[XLA:GPU] Split command buffer VA remapping into its own allocator class

Separate the command buffer VA remapping implementation (SKIP_TEMP
update mode) from GpuExecutableBufferAllocator into a new
GpuExecutableVaRemapAllocator subclass in
gpu_executable_va_remap_allocator.{h,cc}:

- GpuExecutableBufferAllocator becomes the base class implementing the
  ALWAYS_UPDATE update mode (constants-only persistent allocations) and
  the runtime fallback behavior for SKIP_TEMP when VA remapping is
  unavailable. Its ExecutionScope is now polymorphic and exposes
  protected hooks (PrepareReservation, AllocateTransientBuffer) at the
  points where an allocation-address policy can intervene.
- GpuExecutableVaRemapAllocator implements SKIP_TEMP: the per-executor
  Remapping state, VA reservation layout, and temp-buffer remapping
  move to the new file.
- A GpuExecutableBufferAllocator::Create factory picks the class from
  the update mode; CreateExecutionScope now returns
  std::unique_ptr<ExecutionScope>.

No functional change intended.

--
5665403db4e1b7583251618533ca9378be006765 by Shawn Wang <[email protected]>:

Fix VA remap allocator CI failures

--
7da1c618865c4a910928df8dc03f60afd7466ab2 by Shawn Wang <[email protected]>:

Restore platform-restricted GPU PJRT deps

Merging this change closes #45567

PiperOrigin-RevId: 948290643

…ions

Imported from GitHub PR #45595

## 📝 Summary of Changes

This PR extends command-buffer virtual-address remapping to entry-computation parameter allocations. Parameter buffers remain caller-owned, so each execution aliases their physical memory into stable VA-reservation addresses instead of transferring ownership to the VMM allocator.

- Adds `ExecutionScope::ResolveParameterBuffer`; the base implementation returns the caller buffer, while the VA-remapping implementation installs and returns the stable reservation alias.
- Renames the `PrepareReservation` hook to `Prepare` and makes `ExecuteWithBufferAllocations` accept mutable `BufferAllocations`, allowing addresses to be restored after execution.
- Tracks per-execution parameter aliases with `RecordStepAlias`.
- Releases those aliases after every execution and rewrites remapped entries back to their caller-owned addresses before teardown and result handling. The execution-scope destructor provides cleanup for error paths.
- Adds `AddVaRemappedAllocationForTesting` and a host-only fake VMM allocator to exercise the parameter-remapping path without requiring a GPU.

This PR is stacked on #45567, **[XLA:GPU] Split command buffer VA remapping into its own allocator class**. It is part 1 of 3 splitting the former #45568 implementation into reviewable changes; #45596 adds output-buffer remapping and #45597 adds the `SKIP_PROFILED` mode.

## 🎯 Justification

`SKIP_TEMP` stabilizes only preallocated temporary buffers, but profiled command-buffer update policies also need to stabilize caller-owned parameters whose physical addresses remain constant across executions. Parameter ownership prevents allocating them directly through the command-buffer reservation.

Creating a per-execution alias gives recorded command buffers a stable VA while preserving the caller's original address and ownership. Restoring the original address after execution ensures teardown and result handling never observe a reservation address.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

Not applicable as a standalone change. This PR adds parameter-remapping support but does not select it in a production update mode. Performance evaluation belongs to the final `SKIP_PROFILED` policy in #45597.

## 🧪 Unit Tests:

- Added host-only `gpu_executable_va_remap_allocator_test.cc` coverage using a fake VMM allocator.
- Verifies that a caller-owned parameter receives the same stable reservation address across executions.
- Verifies persistent allocation-index handling, post-execution restoration to the caller address, and release of each per-execution mapping.
- Covers cleanup behavior without requiring a GPU.

## 🧪 Execution Tests:

No standalone execution test is added because no production update mode selects parameter allocations in this PR. End-to-end selection and execution coverage is provided by the stacked `SKIP_PROFILED` change in #45597.

Copybara import of the project:

--
a0de77a by Shawn Wang <[email protected]>:

[XLA:GPU] Support VA remapping of parameter buffer allocations

Allow command buffer VA remapping to give stable reservation addresses
to entry-computation-parameter allocations. Parameter buffers are owned
by the caller, so they are aliased into the VA reservation with Map()
each execution instead of being allocated through it:

- GpuExecutableBufferAllocator::ExecutionScope gains a
  ResolveParameterBuffer hook (base returns the caller buffer), and
  ExecuteWithBufferAllocations now takes a mutable BufferAllocations so
  subclasses can rewrite entries after execution. The Prepare hook is
  renamed from PrepareReservation.
- GpuExecutableVaRemapAllocator records each installed alias
  (RecordStepAlias) and, after every execution, UnMaps the aliases and
  rewrites the remapped BufferAllocations entries back to the
  caller-owned addresses so TearDown and result handling never see
  reservation addresses (ReleaseStepAliases; the scope destructor is a
  safety net for error paths).
- No production update mode selects parameter allocations yet
  (SKIP_TEMP remaps only preallocated temp buffers); the
  AddVaRemappedAllocationForTesting seam plus a host-only fake VMM
  allocator (gpu_executable_va_remap_allocator_test.cc) exercise the
  new path: stable reservation address across executions, persistent
  allocation indices, post-execution rewrite, and mapping release.

--
9e7c7fb by Shawn Wang <[email protected]>:

Fix VA remap allocator CI failures

--
7ec8063 by Shawn Wang <[email protected]>:

[XLA:GPU] Preserve owning buffers during VA remapping

--
b4db679 by Shawn Wang <[email protected]>:

[XLA:GPU] Return remapped buffer allocations by value

--
4d5aabb by Shawn Wang <[email protected]>:

[XLA:GPU] Rename remapped allocation helper

Merging this change closes #45595

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45595 from shawnwang18:shawnw/cmd_buffer_va_remap_params 4d5aabb
PiperOrigin-RevId: 948758452

…ions

Imported from GitHub PR #45595

## 📝 Summary of Changes

This PR extends command-buffer virtual-address remapping to entry-computation parameter allocations. Parameter buffers remain caller-owned, so each execution aliases their physical memory into stable VA-reservation addresses instead of transferring ownership to the VMM allocator.

- Adds `ExecutionScope::ResolveParameterBuffer`; the base implementation returns the caller buffer, while the VA-remapping implementation installs and returns the stable reservation alias.
- Renames the `PrepareReservation` hook to `Prepare` and makes `ExecuteWithBufferAllocations` accept mutable `BufferAllocations`, allowing addresses to be restored after execution.
- Tracks per-execution parameter aliases with `RecordStepAlias`.
- Releases those aliases after every execution and rewrites remapped entries back to their caller-owned addresses before teardown and result handling. The execution-scope destructor provides cleanup for error paths.
- Adds `AddVaRemappedAllocationForTesting` and a host-only fake VMM allocator to exercise the parameter-remapping path without requiring a GPU.

This PR is stacked on #45567, **[XLA:GPU] Split command buffer VA remapping into its own allocator class**. It is part 1 of 3 splitting the former #45568 implementation into reviewable changes; #45596 adds output-buffer remapping and #45597 adds the `SKIP_PROFILED` mode.

## 🎯 Justification

`SKIP_TEMP` stabilizes only preallocated temporary buffers, but profiled command-buffer update policies also need to stabilize caller-owned parameters whose physical addresses remain constant across executions. Parameter ownership prevents allocating them directly through the command-buffer reservation.

Creating a per-execution alias gives recorded command buffers a stable VA while preserving the caller's original address and ownership. Restoring the original address after execution ensures teardown and result handling never observe a reservation address.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

Not applicable as a standalone change. This PR adds parameter-remapping support but does not select it in a production update mode. Performance evaluation belongs to the final `SKIP_PROFILED` policy in #45597.

## 🧪 Unit Tests:

- Added host-only `gpu_executable_va_remap_allocator_test.cc` coverage using a fake VMM allocator.
- Verifies that a caller-owned parameter receives the same stable reservation address across executions.
- Verifies persistent allocation-index handling, post-execution restoration to the caller address, and release of each per-execution mapping.
- Covers cleanup behavior without requiring a GPU.

## 🧪 Execution Tests:

No standalone execution test is added because no production update mode selects parameter allocations in this PR. End-to-end selection and execution coverage is provided by the stacked `SKIP_PROFILED` change in #45597.

Copybara import of the project:

--
a0de77a by Shawn Wang <[email protected]>:

[XLA:GPU] Support VA remapping of parameter buffer allocations

Allow command buffer VA remapping to give stable reservation addresses
to entry-computation-parameter allocations. Parameter buffers are owned
by the caller, so they are aliased into the VA reservation with Map()
each execution instead of being allocated through it:

- GpuExecutableBufferAllocator::ExecutionScope gains a
  ResolveParameterBuffer hook (base returns the caller buffer), and
  ExecuteWithBufferAllocations now takes a mutable BufferAllocations so
  subclasses can rewrite entries after execution. The Prepare hook is
  renamed from PrepareReservation.
- GpuExecutableVaRemapAllocator records each installed alias
  (RecordStepAlias) and, after every execution, UnMaps the aliases and
  rewrites the remapped BufferAllocations entries back to the
  caller-owned addresses so TearDown and result handling never see
  reservation addresses (ReleaseStepAliases; the scope destructor is a
  safety net for error paths).
- No production update mode selects parameter allocations yet
  (SKIP_TEMP remaps only preallocated temp buffers); the
  AddVaRemappedAllocationForTesting seam plus a host-only fake VMM
  allocator (gpu_executable_va_remap_allocator_test.cc) exercise the
  new path: stable reservation address across executions, persistent
  allocation indices, post-execution rewrite, and mapping release.

--
9e7c7fb by Shawn Wang <[email protected]>:

Fix VA remap allocator CI failures

--
7ec8063 by Shawn Wang <[email protected]>:

[XLA:GPU] Preserve owning buffers during VA remapping

--
b4db679 by Shawn Wang <[email protected]>:

[XLA:GPU] Return remapped buffer allocations by value

--
4d5aabb by Shawn Wang <[email protected]>:

[XLA:GPU] Rename remapped allocation helper

Merging this change closes #45595

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45595 from shawnwang18:shawnw/cmd_buffer_va_remap_params 4d5aabb
PiperOrigin-RevId: 948758452

…ions

Imported from GitHub PR #45595

## 📝 Summary of Changes

This PR extends command-buffer virtual-address remapping to entry-computation parameter allocations. Parameter buffers remain caller-owned, so each execution aliases their physical memory into stable VA-reservation addresses instead of transferring ownership to the VMM allocator.

- Adds `ExecutionScope::ResolveParameterBuffer`; the base implementation returns the caller buffer, while the VA-remapping implementation installs and returns the stable reservation alias.
- Renames the `PrepareReservation` hook to `Prepare` and makes `ExecuteWithBufferAllocations` accept mutable `BufferAllocations`, allowing addresses to be restored after execution.
- Tracks per-execution parameter aliases with `RecordStepAlias`.
- Releases those aliases after every execution and rewrites remapped entries back to their caller-owned addresses before teardown and result handling. The execution-scope destructor provides cleanup for error paths.
- Adds `AddVaRemappedAllocationForTesting` and a host-only fake VMM allocator to exercise the parameter-remapping path without requiring a GPU.

This PR is stacked on #45567, **[XLA:GPU] Split command buffer VA remapping into its own allocator class**. It is part 1 of 3 splitting the former #45568 implementation into reviewable changes; #45596 adds output-buffer remapping and #45597 adds the `SKIP_PROFILED` mode.

## 🎯 Justification

`SKIP_TEMP` stabilizes only preallocated temporary buffers, but profiled command-buffer update policies also need to stabilize caller-owned parameters whose physical addresses remain constant across executions. Parameter ownership prevents allocating them directly through the command-buffer reservation.

Creating a per-execution alias gives recorded command buffers a stable VA while preserving the caller's original address and ownership. Restoring the original address after execution ensures teardown and result handling never observe a reservation address.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

Not applicable as a standalone change. This PR adds parameter-remapping support but does not select it in a production update mode. Performance evaluation belongs to the final `SKIP_PROFILED` policy in #45597.

## 🧪 Unit Tests:

- Added host-only `gpu_executable_va_remap_allocator_test.cc` coverage using a fake VMM allocator.
- Verifies that a caller-owned parameter receives the same stable reservation address across executions.
- Verifies persistent allocation-index handling, post-execution restoration to the caller address, and release of each per-execution mapping.
- Covers cleanup behavior without requiring a GPU.

## 🧪 Execution Tests:

No standalone execution test is added because no production update mode selects parameter allocations in this PR. End-to-end selection and execution coverage is provided by the stacked `SKIP_PROFILED` change in #45597.

Copybara import of the project:

--
a0de77a by Shawn Wang <[email protected]>:

[XLA:GPU] Support VA remapping of parameter buffer allocations

Allow command buffer VA remapping to give stable reservation addresses
to entry-computation-parameter allocations. Parameter buffers are owned
by the caller, so they are aliased into the VA reservation with Map()
each execution instead of being allocated through it:

- GpuExecutableBufferAllocator::ExecutionScope gains a
  ResolveParameterBuffer hook (base returns the caller buffer), and
  ExecuteWithBufferAllocations now takes a mutable BufferAllocations so
  subclasses can rewrite entries after execution. The Prepare hook is
  renamed from PrepareReservation.
- GpuExecutableVaRemapAllocator records each installed alias
  (RecordStepAlias) and, after every execution, UnMaps the aliases and
  rewrites the remapped BufferAllocations entries back to the
  caller-owned addresses so TearDown and result handling never see
  reservation addresses (ReleaseStepAliases; the scope destructor is a
  safety net for error paths).
- No production update mode selects parameter allocations yet
  (SKIP_TEMP remaps only preallocated temp buffers); the
  AddVaRemappedAllocationForTesting seam plus a host-only fake VMM
  allocator (gpu_executable_va_remap_allocator_test.cc) exercise the
  new path: stable reservation address across executions, persistent
  allocation indices, post-execution rewrite, and mapping release.

--
9e7c7fb by Shawn Wang <[email protected]>:

Fix VA remap allocator CI failures

--
7ec8063 by Shawn Wang <[email protected]>:

[XLA:GPU] Preserve owning buffers during VA remapping

--
b4db679 by Shawn Wang <[email protected]>:

[XLA:GPU] Return remapped buffer allocations by value

--
4d5aabb by Shawn Wang <[email protected]>:

[XLA:GPU] Rename remapped allocation helper

Merging this change closes #45595

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45595 from shawnwang18:shawnw/cmd_buffer_va_remap_params 4d5aabb
PiperOrigin-RevId: 948758452

…ions

Imported from GitHub PR #45595

## 📝 Summary of Changes

This PR extends command-buffer virtual-address remapping to entry-computation parameter allocations. Parameter buffers remain caller-owned, so each execution aliases their physical memory into stable VA-reservation addresses instead of transferring ownership to the VMM allocator.

- Adds `ExecutionScope::ResolveParameterBuffer`; the base implementation returns the caller buffer, while the VA-remapping implementation installs and returns the stable reservation alias.
- Renames the `PrepareReservation` hook to `Prepare` and makes `ExecuteWithBufferAllocations` accept mutable `BufferAllocations`, allowing addresses to be restored after execution.
- Tracks per-execution parameter aliases with `RecordStepAlias`.
- Releases those aliases after every execution and rewrites remapped entries back to their caller-owned addresses before teardown and result handling. The execution-scope destructor provides cleanup for error paths.
- Adds `AddVaRemappedAllocationForTesting` and a host-only fake VMM allocator to exercise the parameter-remapping path without requiring a GPU.

This PR is stacked on #45567, **[XLA:GPU] Split command buffer VA remapping into its own allocator class**. It is part 1 of 3 splitting the former #45568 implementation into reviewable changes; #45596 adds output-buffer remapping and #45597 adds the `SKIP_PROFILED` mode.

## 🎯 Justification

`SKIP_TEMP` stabilizes only preallocated temporary buffers, but profiled command-buffer update policies also need to stabilize caller-owned parameters whose physical addresses remain constant across executions. Parameter ownership prevents allocating them directly through the command-buffer reservation.

Creating a per-execution alias gives recorded command buffers a stable VA while preserving the caller's original address and ownership. Restoring the original address after execution ensures teardown and result handling never observe a reservation address.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

Not applicable as a standalone change. This PR adds parameter-remapping support but does not select it in a production update mode. Performance evaluation belongs to the final `SKIP_PROFILED` policy in #45597.

## 🧪 Unit Tests:

- Added host-only `gpu_executable_va_remap_allocator_test.cc` coverage using a fake VMM allocator.
- Verifies that a caller-owned parameter receives the same stable reservation address across executions.
- Verifies persistent allocation-index handling, post-execution restoration to the caller address, and release of each per-execution mapping.
- Covers cleanup behavior without requiring a GPU.

## 🧪 Execution Tests:

No standalone execution test is added because no production update mode selects parameter allocations in this PR. End-to-end selection and execution coverage is provided by the stacked `SKIP_PROFILED` change in #45597.

Copybara import of the project:

--
a0de77a by Shawn Wang <[email protected]>:

[XLA:GPU] Support VA remapping of parameter buffer allocations

Allow command buffer VA remapping to give stable reservation addresses
to entry-computation-parameter allocations. Parameter buffers are owned
by the caller, so they are aliased into the VA reservation with Map()
each execution instead of being allocated through it:

- GpuExecutableBufferAllocator::ExecutionScope gains a
  ResolveParameterBuffer hook (base returns the caller buffer), and
  ExecuteWithBufferAllocations now takes a mutable BufferAllocations so
  subclasses can rewrite entries after execution. The Prepare hook is
  renamed from PrepareReservation.
- GpuExecutableVaRemapAllocator records each installed alias
  (RecordStepAlias) and, after every execution, UnMaps the aliases and
  rewrites the remapped BufferAllocations entries back to the
  caller-owned addresses so TearDown and result handling never see
  reservation addresses (ReleaseStepAliases; the scope destructor is a
  safety net for error paths).
- No production update mode selects parameter allocations yet
  (SKIP_TEMP remaps only preallocated temp buffers); the
  AddVaRemappedAllocationForTesting seam plus a host-only fake VMM
  allocator (gpu_executable_va_remap_allocator_test.cc) exercise the
  new path: stable reservation address across executions, persistent
  allocation indices, post-execution rewrite, and mapping release.

--
9e7c7fb by Shawn Wang <[email protected]>:

Fix VA remap allocator CI failures

--
7ec8063 by Shawn Wang <[email protected]>:

[XLA:GPU] Preserve owning buffers during VA remapping

--
b4db679 by Shawn Wang <[email protected]>:

[XLA:GPU] Return remapped buffer allocations by value

--
4d5aabb by Shawn Wang <[email protected]>:

[XLA:GPU] Rename remapped allocation helper

Merging this change closes #45595

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45595 from shawnwang18:shawnw/cmd_buffer_va_remap_params 4d5aabb
PiperOrigin-RevId: 948758452

…ions

Imported from GitHub PR #45595

## 📝 Summary of Changes

This PR extends command-buffer virtual-address remapping to entry-computation parameter allocations. Parameter buffers remain caller-owned, so each execution aliases their physical memory into stable VA-reservation addresses instead of transferring ownership to the VMM allocator.

- Adds `ExecutionScope::ResolveParameterBuffer`; the base implementation returns the caller buffer, while the VA-remapping implementation installs and returns the stable reservation alias.
- Renames the `PrepareReservation` hook to `Prepare` and makes `ExecuteWithBufferAllocations` accept mutable `BufferAllocations`, allowing addresses to be restored after execution.
- Tracks per-execution parameter aliases with `RecordStepAlias`.
- Releases those aliases after every execution and rewrites remapped entries back to their caller-owned addresses before teardown and result handling. The execution-scope destructor provides cleanup for error paths.
- Adds `AddVaRemappedAllocationForTesting` and a host-only fake VMM allocator to exercise the parameter-remapping path without requiring a GPU.

This PR is stacked on #45567, **[XLA:GPU] Split command buffer VA remapping into its own allocator class**. It is part 1 of 3 splitting the former #45568 implementation into reviewable changes; #45596 adds output-buffer remapping and #45597 adds the `SKIP_PROFILED` mode.

## 🎯 Justification

`SKIP_TEMP` stabilizes only preallocated temporary buffers, but profiled command-buffer update policies also need to stabilize caller-owned parameters whose physical addresses remain constant across executions. Parameter ownership prevents allocating them directly through the command-buffer reservation.

Creating a per-execution alias gives recorded command buffers a stable VA while preserving the caller's original address and ownership. Restoring the original address after execution ensures teardown and result handling never observe a reservation address.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

Not applicable as a standalone change. This PR adds parameter-remapping support but does not select it in a production update mode. Performance evaluation belongs to the final `SKIP_PROFILED` policy in #45597.

## 🧪 Unit Tests:

- Added host-only `gpu_executable_va_remap_allocator_test.cc` coverage using a fake VMM allocator.
- Verifies that a caller-owned parameter receives the same stable reservation address across executions.
- Verifies persistent allocation-index handling, post-execution restoration to the caller address, and release of each per-execution mapping.
- Covers cleanup behavior without requiring a GPU.

## 🧪 Execution Tests:

No standalone execution test is added because no production update mode selects parameter allocations in this PR. End-to-end selection and execution coverage is provided by the stacked `SKIP_PROFILED` change in #45597.

Copybara import of the project:

--
a0de77a by Shawn Wang <[email protected]>:

[XLA:GPU] Support VA remapping of parameter buffer allocations

Allow command buffer VA remapping to give stable reservation addresses
to entry-computation-parameter allocations. Parameter buffers are owned
by the caller, so they are aliased into the VA reservation with Map()
each execution instead of being allocated through it:

- GpuExecutableBufferAllocator::ExecutionScope gains a
  ResolveParameterBuffer hook (base returns the caller buffer), and
  ExecuteWithBufferAllocations now takes a mutable BufferAllocations so
  subclasses can rewrite entries after execution. The Prepare hook is
  renamed from PrepareReservation.
- GpuExecutableVaRemapAllocator records each installed alias
  (RecordStepAlias) and, after every execution, UnMaps the aliases and
  rewrites the remapped BufferAllocations entries back to the
  caller-owned addresses so TearDown and result handling never see
  reservation addresses (ReleaseStepAliases; the scope destructor is a
  safety net for error paths).
- No production update mode selects parameter allocations yet
  (SKIP_TEMP remaps only preallocated temp buffers); the
  AddVaRemappedAllocationForTesting seam plus a host-only fake VMM
  allocator (gpu_executable_va_remap_allocator_test.cc) exercise the
  new path: stable reservation address across executions, persistent
  allocation indices, post-execution rewrite, and mapping release.

--
9e7c7fb by Shawn Wang <[email protected]>:

Fix VA remap allocator CI failures

--
7ec8063 by Shawn Wang <[email protected]>:

[XLA:GPU] Preserve owning buffers during VA remapping

--
b4db679 by Shawn Wang <[email protected]>:

[XLA:GPU] Return remapped buffer allocations by value

--
4d5aabb by Shawn Wang <[email protected]>:

[XLA:GPU] Rename remapped allocation helper

Merging this change closes #45595

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45595 from shawnwang18:shawnw/cmd_buffer_va_remap_params 4d5aabb
PiperOrigin-RevId: 948758452

…ions

Imported from GitHub PR #45595

## 📝 Summary of Changes

This PR extends command-buffer virtual-address remapping to entry-computation parameter allocations. Parameter buffers remain caller-owned, so each execution aliases their physical memory into stable VA-reservation addresses instead of transferring ownership to the VMM allocator.

- Adds `ExecutionScope::ResolveParameterBuffer`; the base implementation returns the caller buffer, while the VA-remapping implementation installs and returns the stable reservation alias.
- Renames the `PrepareReservation` hook to `Prepare` and makes `ExecuteWithBufferAllocations` accept mutable `BufferAllocations`, allowing addresses to be restored after execution.
- Tracks per-execution parameter aliases with `RecordStepAlias`.
- Releases those aliases after every execution and rewrites remapped entries back to their caller-owned addresses before teardown and result handling. The execution-scope destructor provides cleanup for error paths.
- Adds `AddVaRemappedAllocationForTesting` and a host-only fake VMM allocator to exercise the parameter-remapping path without requiring a GPU.

This PR is stacked on #45567, **[XLA:GPU] Split command buffer VA remapping into its own allocator class**. It is part 1 of 3 splitting the former #45568 implementation into reviewable changes; #45596 adds output-buffer remapping and #45597 adds the `SKIP_PROFILED` mode.

## 🎯 Justification

`SKIP_TEMP` stabilizes only preallocated temporary buffers, but profiled command-buffer update policies also need to stabilize caller-owned parameters whose physical addresses remain constant across executions. Parameter ownership prevents allocating them directly through the command-buffer reservation.

Creating a per-execution alias gives recorded command buffers a stable VA while preserving the caller's original address and ownership. Restoring the original address after execution ensures teardown and result handling never observe a reservation address.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

Not applicable as a standalone change. This PR adds parameter-remapping support but does not select it in a production update mode. Performance evaluation belongs to the final `SKIP_PROFILED` policy in #45597.

## 🧪 Unit Tests:

- Added host-only `gpu_executable_va_remap_allocator_test.cc` coverage using a fake VMM allocator.
- Verifies that a caller-owned parameter receives the same stable reservation address across executions.
- Verifies persistent allocation-index handling, post-execution restoration to the caller address, and release of each per-execution mapping.
- Covers cleanup behavior without requiring a GPU.

## 🧪 Execution Tests:

No standalone execution test is added because no production update mode selects parameter allocations in this PR. End-to-end selection and execution coverage is provided by the stacked `SKIP_PROFILED` change in #45597.

Copybara import of the project:

--
a0de77a by Shawn Wang <[email protected]>:

[XLA:GPU] Support VA remapping of parameter buffer allocations

Allow command buffer VA remapping to give stable reservation addresses
to entry-computation-parameter allocations. Parameter buffers are owned
by the caller, so they are aliased into the VA reservation with Map()
each execution instead of being allocated through it:

- GpuExecutableBufferAllocator::ExecutionScope gains a
  ResolveParameterBuffer hook (base returns the caller buffer), and
  ExecuteWithBufferAllocations now takes a mutable BufferAllocations so
  subclasses can rewrite entries after execution. The Prepare hook is
  renamed from PrepareReservation.
- GpuExecutableVaRemapAllocator records each installed alias
  (RecordStepAlias) and, after every execution, UnMaps the aliases and
  rewrites the remapped BufferAllocations entries back to the
  caller-owned addresses so TearDown and result handling never see
  reservation addresses (ReleaseStepAliases; the scope destructor is a
  safety net for error paths).
- No production update mode selects parameter allocations yet
  (SKIP_TEMP remaps only preallocated temp buffers); the
  AddVaRemappedAllocationForTesting seam plus a host-only fake VMM
  allocator (gpu_executable_va_remap_allocator_test.cc) exercise the
  new path: stable reservation address across executions, persistent
  allocation indices, post-execution rewrite, and mapping release.

--
9e7c7fb by Shawn Wang <[email protected]>:

Fix VA remap allocator CI failures

--
7ec8063 by Shawn Wang <[email protected]>:

[XLA:GPU] Preserve owning buffers during VA remapping

--
b4db679 by Shawn Wang <[email protected]>:

[XLA:GPU] Return remapped buffer allocations by value

--
4d5aabb by Shawn Wang <[email protected]>:

[XLA:GPU] Rename remapped allocation helper

Merging this change closes #45595

COPYBARA_INTEGRATE_REVIEW=#45595 from shawnwang18:shawnw/cmd_buffer_va_remap_params 4d5aabb
PiperOrigin-RevId: 948990342