[XLA:GPU] Support VA remapping of output buffer allocations - #45596

…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

Imported from GitHub PR #45596

## 📝 Summary of Changes

This PR extends command-buffer virtual-address remapping to allocations that
escape to the caller (`maybe_live_out`) while preserving the const ownership
model introduced by #45595.

- Keeps selected live-out allocations at normal allocator-owned addresses in
  the owning `BufferAllocations` table.
- Extends `GetRemappedBufferAllocations` to map finalized selected input/output
  addresses into their stable reservation slices immediately before execution.
- Rewrites only the non-owning execution copy, so result handling and teardown
  continue to see external, deallocatable addresses.
- Handles donated and copy-protected aliases by mapping after output setup has
  finalized the owning address table.
- Extends the host-only fake-VMM tests to cover stable execution aliases,
  logical external sizes, output deallocation, and deferred alias cleanup.

This PR is stacked on #45595, **[XLA:GPU] Support VA remapping of parameter
buffer allocations**. It is part 2 of 3 splitting the former #45568
implementation; #45597 adds the `SKIP_PROFILED` mode.

## 🎯 Justification

Live-out buffers can outlive an execution scope, so the caller must retain a
normal allocator-owned address that remains valid and can be passed to
`Deallocate()`. At the same time, recorded command buffers need a stable
reservation address.

Building a separate execution allocation table provides both properties
without mutating the owning table or removing its const qualifier. Deferring
the mapping until execution also ensures donation and copy protection have
already selected the final external address.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

Not applicable as a standalone change. This PR adds output-buffer remapping
support but does not independently enable a new command-buffer update mode.
Performance evaluation belongs to the final `SKIP_PROFILED` policy in #45597.

## 🧪 Unit Tests:

- Extends the host-only fake-VMM allocator test with a non-granularity-aligned
  remapped live-out allocation.
- Verifies the owning table contains the logical-size external address before
  and after execution.
- Verifies execution receives the same stable reservation address across runs
  and maps the correct raw physical allocation.
- Verifies the external address is deallocatable and all aliases are released.

## 🧪 Execution Tests:

No standalone execution test is added for this enabling layer. End-to-end
live-out selection, execution, and output deallocation are covered by the
stacked `SKIP_PROFILED` change in #45597.

Copybara import of the project:

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

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

Keep selected maybe_live_out allocations at normal allocator-owned addresses in the owning BufferAllocations table. Immediately before execution, GetRemappedBufferAllocations maps the finalized selected input/output addresses into their stable reservation slices and rewrites only a non-owning execution copy.

This preserves const ownership semantics, handles donation and copy protection before mapping, and leaves output addresses deallocatable by callers. Extend the host-only VMM test to verify stable execution addresses, logical external sizes, unchanged owning buffers, and alias cleanup across executions.

Merging this change closes #45596

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45596 from shawnwang18:shawnw/cmd_buffer_va_remap_outputs b7b79af
PiperOrigin-RevId: 950641350

Imported from GitHub PR openxla/xla#45596

## 📝 Summary of Changes

This PR extends command-buffer virtual-address remapping to allocations that
escape to the caller (`maybe_live_out`) while preserving the const ownership
model introduced by #45595.

- Keeps selected live-out allocations at normal allocator-owned addresses in
  the owning `BufferAllocations` table.
- Extends `GetRemappedBufferAllocations` to map finalized selected input/output
  addresses into their stable reservation slices immediately before execution.
- Rewrites only the non-owning execution copy, so result handling and teardown
  continue to see external, deallocatable addresses.
- Handles donated and copy-protected aliases by mapping after output setup has
  finalized the owning address table.
- Extends the host-only fake-VMM tests to cover stable execution aliases,
  logical external sizes, output deallocation, and deferred alias cleanup.

This PR is stacked on #45595, **[XLA:GPU] Support VA remapping of parameter
buffer allocations**. It is part 2 of 3 splitting the former #45568
implementation; #45597 adds the `SKIP_PROFILED` mode.

## 🎯 Justification

Live-out buffers can outlive an execution scope, so the caller must retain a
normal allocator-owned address that remains valid and can be passed to
`Deallocate()`. At the same time, recorded command buffers need a stable
reservation address.

Building a separate execution allocation table provides both properties
without mutating the owning table or removing its const qualifier. Deferring
the mapping until execution also ensures donation and copy protection have
already selected the final external address.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

Not applicable as a standalone change. This PR adds output-buffer remapping
support but does not independently enable a new command-buffer update mode.
Performance evaluation belongs to the final `SKIP_PROFILED` policy in #45597.

## 🧪 Unit Tests:

- Extends the host-only fake-VMM allocator test with a non-granularity-aligned
  remapped live-out allocation.
- Verifies the owning table contains the logical-size external address before
  and after execution.
- Verifies execution receives the same stable reservation address across runs
  and maps the correct raw physical allocation.
- Verifies the external address is deallocatable and all aliases are released.

## 🧪 Execution Tests:

No standalone execution test is added for this enabling layer. End-to-end
live-out selection, execution, and output deallocation are covered by the
stacked `SKIP_PROFILED` change in #45597.

Copybara import of the project:

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

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

Keep selected maybe_live_out allocations at normal allocator-owned addresses in the owning BufferAllocations table. Immediately before execution, GetRemappedBufferAllocations maps the finalized selected input/output addresses into their stable reservation slices and rewrites only a non-owning execution copy.

This preserves const ownership semantics, handles donation and copy protection before mapping, and leaves output addresses deallocatable by callers. Extend the host-only VMM test to verify stable execution addresses, logical external sizes, unchanged owning buffers, and alias cleanup across executions.

Merging this change closes #45596

Reverts changelist 949495726

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#45596 from shawnwang18:shawnw/cmd_buffer_va_remap_outputs b7b79af2174271e65f2c7d8e8bac527a72e97af9
PiperOrigin-RevId: 950641350

Imported from GitHub PR #45596

## 📝 Summary of Changes

This PR extends command-buffer virtual-address remapping to allocations that
escape to the caller (`maybe_live_out`) while preserving the const ownership
model introduced by #45595.

- Keeps selected live-out allocations at normal allocator-owned addresses in
  the owning `BufferAllocations` table.
- Extends `GetRemappedBufferAllocations` to map finalized selected input/output
  addresses into their stable reservation slices immediately before execution.
- Rewrites only the non-owning execution copy, so result handling and teardown
  continue to see external, deallocatable addresses.
- Handles donated and copy-protected aliases by mapping after output setup has
  finalized the owning address table.
- Extends the host-only fake-VMM tests to cover stable execution aliases,
  logical external sizes, output deallocation, and deferred alias cleanup.

This PR is stacked on #45595, **[XLA:GPU] Support VA remapping of parameter
buffer allocations**. It is part 2 of 3 splitting the former #45568
implementation; #45597 adds the `SKIP_PROFILED` mode.

## 🎯 Justification

Live-out buffers can outlive an execution scope, so the caller must retain a
normal allocator-owned address that remains valid and can be passed to
`Deallocate()`. At the same time, recorded command buffers need a stable
reservation address.

Building a separate execution allocation table provides both properties
without mutating the owning table or removing its const qualifier. Deferring
the mapping until execution also ensures donation and copy protection have
already selected the final external address.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

Not applicable as a standalone change. This PR adds output-buffer remapping
support but does not independently enable a new command-buffer update mode.
Performance evaluation belongs to the final `SKIP_PROFILED` policy in #45597.

## 🧪 Unit Tests:

- Extends the host-only fake-VMM allocator test with a non-granularity-aligned
  remapped live-out allocation.
- Verifies the owning table contains the logical-size external address before
  and after execution.
- Verifies execution receives the same stable reservation address across runs
  and maps the correct raw physical allocation.
- Verifies the external address is deallocatable and all aliases are released.

## 🧪 Execution Tests:

No standalone execution test is added for this enabling layer. End-to-end
live-out selection, execution, and output deallocation are covered by the
stacked `SKIP_PROFILED` change in #45597.

Copybara import of the project:

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

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

Keep selected maybe_live_out allocations at normal allocator-owned addresses in the owning BufferAllocations table. Immediately before execution, GetRemappedBufferAllocations maps the finalized selected input/output addresses into their stable reservation slices and rewrites only a non-owning execution copy.

This preserves const ownership semantics, handles donation and copy protection before mapping, and leaves output addresses deallocatable by callers. Extend the host-only VMM test to verify stable execution addresses, logical external sizes, unchanged owning buffers, and alias cleanup across executions.

Merging this change closes #45596

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45596 from shawnwang18:shawnw/cmd_buffer_va_remap_outputs b7b79af
PiperOrigin-RevId: 950641350

Imported from GitHub PR openxla/xla#45596

## 📝 Summary of Changes

This PR extends command-buffer virtual-address remapping to allocations that
escape to the caller (`maybe_live_out`) while preserving the const ownership
model introduced by #45595.

- Keeps selected live-out allocations at normal allocator-owned addresses in
  the owning `BufferAllocations` table.
- Extends `GetRemappedBufferAllocations` to map finalized selected input/output
  addresses into their stable reservation slices immediately before execution.
- Rewrites only the non-owning execution copy, so result handling and teardown
  continue to see external, deallocatable addresses.
- Handles donated and copy-protected aliases by mapping after output setup has
  finalized the owning address table.
- Extends the host-only fake-VMM tests to cover stable execution aliases,
  logical external sizes, output deallocation, and deferred alias cleanup.

This PR is stacked on #45595, **[XLA:GPU] Support VA remapping of parameter
buffer allocations**. It is part 2 of 3 splitting the former #45568
implementation; #45597 adds the `SKIP_PROFILED` mode.

## 🎯 Justification

Live-out buffers can outlive an execution scope, so the caller must retain a
normal allocator-owned address that remains valid and can be passed to
`Deallocate()`. At the same time, recorded command buffers need a stable
reservation address.

Building a separate execution allocation table provides both properties
without mutating the owning table or removing its const qualifier. Deferring
the mapping until execution also ensures donation and copy protection have
already selected the final external address.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

Not applicable as a standalone change. This PR adds output-buffer remapping
support but does not independently enable a new command-buffer update mode.
Performance evaluation belongs to the final `SKIP_PROFILED` policy in #45597.

## 🧪 Unit Tests:

- Extends the host-only fake-VMM allocator test with a non-granularity-aligned
  remapped live-out allocation.
- Verifies the owning table contains the logical-size external address before
  and after execution.
- Verifies execution receives the same stable reservation address across runs
  and maps the correct raw physical allocation.
- Verifies the external address is deallocatable and all aliases are released.

## 🧪 Execution Tests:

No standalone execution test is added for this enabling layer. End-to-end
live-out selection, execution, and output deallocation are covered by the
stacked `SKIP_PROFILED` change in #45597.

Copybara import of the project:

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

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

Keep selected maybe_live_out allocations at normal allocator-owned addresses in the owning BufferAllocations table. Immediately before execution, GetRemappedBufferAllocations maps the finalized selected input/output addresses into their stable reservation slices and rewrites only a non-owning execution copy.

This preserves const ownership semantics, handles donation and copy protection before mapping, and leaves output addresses deallocatable by callers. Extend the host-only VMM test to verify stable execution addresses, logical external sizes, unchanged owning buffers, and alias cleanup across executions.

Merging this change closes #45596

Reverts changelist 949495726

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#45596 from shawnwang18:shawnw/cmd_buffer_va_remap_outputs b7b79af2174271e65f2c7d8e8bac527a72e97af9
PiperOrigin-RevId: 950641350

Also add parametric unit tests that covers all currently supported data types in the CuDNNFusionCompiler.

Reverts changelist 949495726

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#45596 from shawnwang18:shawnw/cmd_buffer_va_remap_outputs b7b79af2174271e65f2c7d8e8bac527a72e97af9
PiperOrigin-RevId: 948295087

CuDNN cannot compile dilated window input tensors and we have to add a padding operation before calling CuDNN.

This change updates the conv_padding_legalization.cc to add a pad operation for raw convolution operations. Currently conv_padding_legalization is used for convolution CustomCalls.
This changes also disallows the algebraic simplifier to fold the pad into the convolution, which would revert conv_padding_legalization.
Note: I cannot place conv_padding_legalization in the fusion pipeline, where ConvFusionRewriter lives, as we need to run layout assignment on the padding op.

An example failing HLO is (focus on `lhs_dilate=`):
```
ROOT %convolution-base-dilated.10 = f32[128,512,14,14]{0,1,3,2} convolution(%parameter.0, %parameter.1), window={size=3x3 pad=1_2x1_2 lhs_dilate=2x2 rhs_reversal=1x1}, dim_labels=bf01_01oi->bf01
```

Reverts changelist 949495726

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#45596 from shawnwang18:shawnw/cmd_buffer_va_remap_outputs b7b79af2174271e65f2c7d8e8bac527a72e97af9
PiperOrigin-RevId: 923250586

Shape inference is used in the HLOVerifier which currently treats all convolutions as forward convolutions.

Reverts changelist 949495726

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#45596 from shawnwang18:shawnw/cmd_buffer_va_remap_outputs b7b79af2174271e65f2c7d8e8bac527a72e97af9
PiperOrigin-RevId: 949061528

Imported from GitHub PR #45597

## 📝 Summary of Changes

This PR adds the `SKIP_PROFILED` command-buffer update mode. The mode
automatically includes command-buffer-referenced preallocated temp allocations
in the final VA-remap set while profiling other eligible non-constant
allocations during the first three successful executions. Stable profile
candidates are combined with the temp allocations so subsequent command-buffer
executions can treat the final set as persistent and avoid unnecessary address
updates.

- Adds `DebugOptions::SKIP_PROFILED` and routes it through the GPU executable
  buffer-allocation policy.
- Adds a per-executor profile state machine to
  `GpuExecutableVaRemapAllocator`: `kProfiling` transitions exactly once to
  `kActive` when an automatic temp or safe profile candidate exists, or to
  `kDisabled` when the final remap set is empty.
- Automatically selects preallocated temp buffers and profiles eligible
  parameter and live-out candidates while filtering unstable addresses,
  non-VMM-backed parameters, and incompatible shared parameter buffers.
- Counts only successful executions toward the profiling window and passes
  `std::nullopt` for persistent allocation indices while profiling, followed by
  one absent-to-present transition to the final persistent set.
- Uses the parameter-remapping support from #45595 and output-buffer remapping
  support from #45596. Parameters and outputs remain externally owned in the
  const owning allocation table; finalized external addresses are mapped into
  an execution-only allocation-table copy immediately before execution.
- Preserves the base allocator fallback when VMM remapping is unavailable or
  there are neither automatic temp buffers nor stable profile candidates.

This PR is stacked on #45596, **[XLA:GPU] Support VA remapping of output buffer
allocations**. It is part 3 of 3 and, together with #45595 and #45596,
supersedes the closed pre-split PR #45568.

## 🎯 Justification

`SKIP_TEMP` can stabilize preallocated temporary buffers, but parameters and
live-out buffers may also keep stable physical addresses across repeated
executions. Updating command-buffer node addresses for those stable allocations
adds avoidable runtime overhead.

`SKIP_PROFILED` discovers address stability for each executable/executor pair
instead of assuming every dynamic allocation is stable. Workloads with
repeatable allocation behavior can therefore stabilize more
command-buffer-visible addresses, while unstable or unsupported workloads retain
the existing fallback behavior. Mapping finalized owning addresses immediately
before execution also keeps donation and copy-protection changes compatible with
the const owning-allocation API introduced by the lower PRs.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

No benchmark result is claimed in this draft. The PR adds the policy and
correctness coverage; a public-HLO performance benchmark can be added separately
when reporting workload-specific speedups.

## 🧪 Unit Tests:

- Extended `GpuExecutableTest.CommandBufferAllocationPolicy` expectations for
  automatic temp-buffer selection in `SKIP_PROFILED`.
- Added host-only fake-VMM coverage for profiling, transition, active remapping,
  and disabled fallback states.
- Verifies a temp-only executable transitions to an active persistent temp set,
  and an unstable parameter cannot disqualify an automatically selected temp.
- Verifies stable parameter and output selection, including non-granularity-sized
  externally owned outputs.
- Verifies failed executions do not advance the three-successful-execution
  profiling window.
- Verifies fallback for unstable and non-VMM-backed parameters and incompatible
  shared parameter buffers.
- Verifies copy protection both during profiling and after transition: the
  owning table retains the finalized external address while execution uses the
  stable reservation address.

## 🧪 Execution Tests:

- Added one-device `VmmTest` coverage using a two-GEMM command-buffer workload
  across profiling, transition, and remapped executions, checking numerical
  correctness on every run.
- Verifies one-time profile selection, automatic temp inclusion, VA-reservation
  creation and reuse, and successful deallocation of externally returned output
  buffers.
- Added a non-VMM allocator execution test verifying the workload remains
  correct and creates no VA reservation when remapping is unavailable.

Copybara import of the project:

--
8b5c4d7 by Shawn Wang <[email protected]>:

[XLA:GPU] Add SKIP_PROFILED command buffer update mode

Add SKIP_PROFILED to profile command-buffer-referenced non-constant allocation addresses for three successful executions, then VA-remap the stable candidates so recorded command buffers can skip updating them.

The per-executor profile transitions once from profiling to active or disabled, preserving the required absent-to-present persistent-index transition. Active executions allocate selected temps at reservation addresses and map finalized selected input/output addresses into an execution-only BufferAllocations copy, leaving the owning table unchanged.

Filter unstable candidates and parameter addresses that are not VMM-backed or are shared by multiple parameter allocations. Add flag parsing, allocation-policy coverage, GPU end-to-end coverage, and fallback behavior without a VMM allocator.

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

Add host-only unit tests for SKIP_PROFILED

Cover the SKIP_PROFILED profile state machine with the fake-VMM harness through a test-only profile-candidate seam.

Verify stable parameters and outputs transition to execution-only reservation addresses while owning buffers remain external, failed executions do not advance profiling, and unstable, non-VMM-backed, or shared parameter buffers disable remapping.

Verify copy-protected addresses are observed after output finalization and can be remapped both during a stable profile and when copy protection first appears after the transition.

--
153df49 by Shawn Wang <[email protected]>:

Remove profile candidate count accessor

--
16473a0 by Shawn Wang <[email protected]>:

[XLA:GPU] Automatically remap temp buffers in SKIP_PROFILED

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

[XLA:GPU] Add direct VA remap allocator includes

Merging this change closes #45597

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45597 from shawnwang18:shawnw/cmd_buffer_skip_profiled_split 7600262
PiperOrigin-RevId: 950773148

And use the same logic for deciding whether to use Triton, and whether to normalize for Triton.

Before this change, it wasn't in sync for per-Ampere GPUs.

Reverts changelist 949495726

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#45596 from shawnwang18:shawnw/cmd_buffer_va_remap_outputs b7b79af2174271e65f2c7d8e8bac527a72e97af9
PiperOrigin-RevId: 950731345

Also add parametric unit tests that covers all currently supported data types in the CuDNNFusionCompiler.

Reverts changelist 949495726

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#45596 from shawnwang18:shawnw/cmd_buffer_va_remap_outputs b7b79af2174271e65f2c7d8e8bac527a72e97af9
PiperOrigin-RevId: 948295087

And use the same logic for deciding whether to use Triton, and whether to normalize for Triton.

Before this change, it wasn't in sync for per-Ampere GPUs.

Reverts changelist 949495726

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#45596 from shawnwang18:shawnw/cmd_buffer_va_remap_outputs b7b79af2174271e65f2c7d8e8bac527a72e97af9
PiperOrigin-RevId: 950731345

Imported from GitHub PR #45597

## 📝 Summary of Changes

This PR adds the `SKIP_PROFILED` command-buffer update mode. The mode
automatically includes command-buffer-referenced preallocated temp allocations
in the final VA-remap set while profiling other eligible non-constant
allocations during the first three successful executions. Stable profile
candidates are combined with the temp allocations so subsequent command-buffer
executions can treat the final set as persistent and avoid unnecessary address
updates.

- Adds `DebugOptions::SKIP_PROFILED` and routes it through the GPU executable
  buffer-allocation policy.
- Adds a per-executor profile state machine to
  `GpuExecutableVaRemapAllocator`: `kProfiling` transitions exactly once to
  `kActive` when an automatic temp or safe profile candidate exists, or to
  `kDisabled` when the final remap set is empty.
- Automatically selects preallocated temp buffers and profiles eligible
  parameter and live-out candidates while filtering unstable addresses,
  non-VMM-backed parameters, and incompatible shared parameter buffers.
- Counts only successful executions toward the profiling window and passes
  `std::nullopt` for persistent allocation indices while profiling, followed by
  one absent-to-present transition to the final persistent set.
- Uses the parameter-remapping support from #45595 and output-buffer remapping
  support from #45596. Parameters and outputs remain externally owned in the
  const owning allocation table; finalized external addresses are mapped into
  an execution-only allocation-table copy immediately before execution.
- Preserves the base allocator fallback when VMM remapping is unavailable or
  there are neither automatic temp buffers nor stable profile candidates.

This PR is stacked on #45596, **[XLA:GPU] Support VA remapping of output buffer
allocations**. It is part 3 of 3 and, together with #45595 and #45596,
supersedes the closed pre-split PR #45568.

## 🎯 Justification

`SKIP_TEMP` can stabilize preallocated temporary buffers, but parameters and
live-out buffers may also keep stable physical addresses across repeated
executions. Updating command-buffer node addresses for those stable allocations
adds avoidable runtime overhead.

`SKIP_PROFILED` discovers address stability for each executable/executor pair
instead of assuming every dynamic allocation is stable. Workloads with
repeatable allocation behavior can therefore stabilize more
command-buffer-visible addresses, while unstable or unsupported workloads retain
the existing fallback behavior. Mapping finalized owning addresses immediately
before execution also keeps donation and copy-protection changes compatible with
the const owning-allocation API introduced by the lower PRs.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

No benchmark result is claimed in this draft. The PR adds the policy and
correctness coverage; a public-HLO performance benchmark can be added separately
when reporting workload-specific speedups.

## 🧪 Unit Tests:

- Extended `GpuExecutableTest.CommandBufferAllocationPolicy` expectations for
  automatic temp-buffer selection in `SKIP_PROFILED`.
- Added host-only fake-VMM coverage for profiling, transition, active remapping,
  and disabled fallback states.
- Verifies a temp-only executable transitions to an active persistent temp set,
  and an unstable parameter cannot disqualify an automatically selected temp.
- Verifies stable parameter and output selection, including non-granularity-sized
  externally owned outputs.
- Verifies failed executions do not advance the three-successful-execution
  profiling window.
- Verifies fallback for unstable and non-VMM-backed parameters and incompatible
  shared parameter buffers.
- Verifies copy protection both during profiling and after transition: the
  owning table retains the finalized external address while execution uses the
  stable reservation address.

## 🧪 Execution Tests:

- Added one-device `VmmTest` coverage using a two-GEMM command-buffer workload
  across profiling, transition, and remapped executions, checking numerical
  correctness on every run.
- Verifies one-time profile selection, automatic temp inclusion, VA-reservation
  creation and reuse, and successful deallocation of externally returned output
  buffers.
- Added a non-VMM allocator execution test verifying the workload remains
  correct and creates no VA reservation when remapping is unavailable.

Copybara import of the project:

--
8b5c4d7 by Shawn Wang <[email protected]>:

[XLA:GPU] Add SKIP_PROFILED command buffer update mode

Add SKIP_PROFILED to profile command-buffer-referenced non-constant allocation addresses for three successful executions, then VA-remap the stable candidates so recorded command buffers can skip updating them.

The per-executor profile transitions once from profiling to active or disabled, preserving the required absent-to-present persistent-index transition. Active executions allocate selected temps at reservation addresses and map finalized selected input/output addresses into an execution-only BufferAllocations copy, leaving the owning table unchanged.

Filter unstable candidates and parameter addresses that are not VMM-backed or are shared by multiple parameter allocations. Add flag parsing, allocation-policy coverage, GPU end-to-end coverage, and fallback behavior without a VMM allocator.

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

Add host-only unit tests for SKIP_PROFILED

Cover the SKIP_PROFILED profile state machine with the fake-VMM harness through a test-only profile-candidate seam.

Verify stable parameters and outputs transition to execution-only reservation addresses while owning buffers remain external, failed executions do not advance profiling, and unstable, non-VMM-backed, or shared parameter buffers disable remapping.

Verify copy-protected addresses are observed after output finalization and can be remapped both during a stable profile and when copy protection first appears after the transition.

--
153df49 by Shawn Wang <[email protected]>:

Remove profile candidate count accessor

--
16473a0 by Shawn Wang <[email protected]>:

[XLA:GPU] Automatically remap temp buffers in SKIP_PROFILED

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

[XLA:GPU] Add direct VA remap allocator includes

Merging this change closes #45597

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45597 from shawnwang18:shawnw/cmd_buffer_skip_profiled_split 7600262
PiperOrigin-RevId: 950773148

Imported from GitHub PR #45597

## 📝 Summary of Changes

This PR adds the `SKIP_PROFILED` command-buffer update mode. The mode
automatically includes command-buffer-referenced preallocated temp allocations
in the final VA-remap set while profiling other eligible non-constant
allocations during the first three successful executions. Stable profile
candidates are combined with the temp allocations so subsequent command-buffer
executions can treat the final set as persistent and avoid unnecessary address
updates.

- Adds `DebugOptions::SKIP_PROFILED` and routes it through the GPU executable
  buffer-allocation policy.
- Adds a per-executor profile state machine to
  `GpuExecutableVaRemapAllocator`: `kProfiling` transitions exactly once to
  `kActive` when an automatic temp or safe profile candidate exists, or to
  `kDisabled` when the final remap set is empty.
- Automatically selects preallocated temp buffers and profiles eligible
  parameter and live-out candidates while filtering unstable addresses,
  non-VMM-backed parameters, and incompatible shared parameter buffers.
- Counts only successful executions toward the profiling window and passes
  `std::nullopt` for persistent allocation indices while profiling, followed by
  one absent-to-present transition to the final persistent set.
- Uses the parameter-remapping support from #45595 and output-buffer remapping
  support from #45596. Parameters and outputs remain externally owned in the
  const owning allocation table; finalized external addresses are mapped into
  an execution-only allocation-table copy immediately before execution.
- Preserves the base allocator fallback when VMM remapping is unavailable or
  there are neither automatic temp buffers nor stable profile candidates.

This PR is stacked on #45596, **[XLA:GPU] Support VA remapping of output buffer
allocations**. It is part 3 of 3 and, together with #45595 and #45596,
supersedes the closed pre-split PR #45568.

## 🎯 Justification

`SKIP_TEMP` can stabilize preallocated temporary buffers, but parameters and
live-out buffers may also keep stable physical addresses across repeated
executions. Updating command-buffer node addresses for those stable allocations
adds avoidable runtime overhead.

`SKIP_PROFILED` discovers address stability for each executable/executor pair
instead of assuming every dynamic allocation is stable. Workloads with
repeatable allocation behavior can therefore stabilize more
command-buffer-visible addresses, while unstable or unsupported workloads retain
the existing fallback behavior. Mapping finalized owning addresses immediately
before execution also keeps donation and copy-protection changes compatible with
the const owning-allocation API introduced by the lower PRs.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

No benchmark result is claimed in this draft. The PR adds the policy and
correctness coverage; a public-HLO performance benchmark can be added separately
when reporting workload-specific speedups.

## 🧪 Unit Tests:

- Extended `GpuExecutableTest.CommandBufferAllocationPolicy` expectations for
  automatic temp-buffer selection in `SKIP_PROFILED`.
- Added host-only fake-VMM coverage for profiling, transition, active remapping,
  and disabled fallback states.
- Verifies a temp-only executable transitions to an active persistent temp set,
  and an unstable parameter cannot disqualify an automatically selected temp.
- Verifies stable parameter and output selection, including non-granularity-sized
  externally owned outputs.
- Verifies failed executions do not advance the three-successful-execution
  profiling window.
- Verifies fallback for unstable and non-VMM-backed parameters and incompatible
  shared parameter buffers.
- Verifies copy protection both during profiling and after transition: the
  owning table retains the finalized external address while execution uses the
  stable reservation address.

## 🧪 Execution Tests:

- Added one-device `VmmTest` coverage using a two-GEMM command-buffer workload
  across profiling, transition, and remapped executions, checking numerical
  correctness on every run.
- Verifies one-time profile selection, automatic temp inclusion, VA-reservation
  creation and reuse, and successful deallocation of externally returned output
  buffers.
- Added a non-VMM allocator execution test verifying the workload remains
  correct and creates no VA reservation when remapping is unavailable.

Copybara import of the project:

--
8b5c4d7 by Shawn Wang <[email protected]>:

[XLA:GPU] Add SKIP_PROFILED command buffer update mode

Add SKIP_PROFILED to profile command-buffer-referenced non-constant allocation addresses for three successful executions, then VA-remap the stable candidates so recorded command buffers can skip updating them.

The per-executor profile transitions once from profiling to active or disabled, preserving the required absent-to-present persistent-index transition. Active executions allocate selected temps at reservation addresses and map finalized selected input/output addresses into an execution-only BufferAllocations copy, leaving the owning table unchanged.

Filter unstable candidates and parameter addresses that are not VMM-backed or are shared by multiple parameter allocations. Add flag parsing, allocation-policy coverage, GPU end-to-end coverage, and fallback behavior without a VMM allocator.

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

Add host-only unit tests for SKIP_PROFILED

Cover the SKIP_PROFILED profile state machine with the fake-VMM harness through a test-only profile-candidate seam.

Verify stable parameters and outputs transition to execution-only reservation addresses while owning buffers remain external, failed executions do not advance profiling, and unstable, non-VMM-backed, or shared parameter buffers disable remapping.

Verify copy-protected addresses are observed after output finalization and can be remapped both during a stable profile and when copy protection first appears after the transition.

--
153df49 by Shawn Wang <[email protected]>:

Remove profile candidate count accessor

--
16473a0 by Shawn Wang <[email protected]>:

[XLA:GPU] Automatically remap temp buffers in SKIP_PROFILED

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

[XLA:GPU] Add direct VA remap allocator includes

Merging this change closes #45597

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45597 from shawnwang18:shawnw/cmd_buffer_skip_profiled_split 7600262
PiperOrigin-RevId: 950773148

Imported from GitHub PR #45597

## 📝 Summary of Changes

This PR adds the `SKIP_PROFILED` command-buffer update mode. The mode
automatically includes command-buffer-referenced preallocated temp allocations
in the final VA-remap set while profiling other eligible non-constant
allocations during the first three successful executions. Stable profile
candidates are combined with the temp allocations so subsequent command-buffer
executions can treat the final set as persistent and avoid unnecessary address
updates.

- Adds `DebugOptions::SKIP_PROFILED` and routes it through the GPU executable
  buffer-allocation policy.
- Adds a per-executor profile state machine to
  `GpuExecutableVaRemapAllocator`: `kProfiling` transitions exactly once to
  `kActive` when an automatic temp or safe profile candidate exists, or to
  `kDisabled` when the final remap set is empty.
- Automatically selects preallocated temp buffers and profiles eligible
  parameter and live-out candidates while filtering unstable addresses,
  non-VMM-backed parameters, and incompatible shared parameter buffers.
- Counts only successful executions toward the profiling window and passes
  `std::nullopt` for persistent allocation indices while profiling, followed by
  one absent-to-present transition to the final persistent set.
- Uses the parameter-remapping support from #45595 and output-buffer remapping
  support from #45596. Parameters and outputs remain externally owned in the
  const owning allocation table; finalized external addresses are mapped into
  an execution-only allocation-table copy immediately before execution.
- Preserves the base allocator fallback when VMM remapping is unavailable or
  there are neither automatic temp buffers nor stable profile candidates.

This PR is stacked on #45596, **[XLA:GPU] Support VA remapping of output buffer
allocations**. It is part 3 of 3 and, together with #45595 and #45596,
supersedes the closed pre-split PR #45568.

## 🎯 Justification

`SKIP_TEMP` can stabilize preallocated temporary buffers, but parameters and
live-out buffers may also keep stable physical addresses across repeated
executions. Updating command-buffer node addresses for those stable allocations
adds avoidable runtime overhead.

`SKIP_PROFILED` discovers address stability for each executable/executor pair
instead of assuming every dynamic allocation is stable. Workloads with
repeatable allocation behavior can therefore stabilize more
command-buffer-visible addresses, while unstable or unsupported workloads retain
the existing fallback behavior. Mapping finalized owning addresses immediately
before execution also keeps donation and copy-protection changes compatible with
the const owning-allocation API introduced by the lower PRs.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

No benchmark result is claimed in this draft. The PR adds the policy and
correctness coverage; a public-HLO performance benchmark can be added separately
when reporting workload-specific speedups.

## 🧪 Unit Tests:

- Extended `GpuExecutableTest.CommandBufferAllocationPolicy` expectations for
  automatic temp-buffer selection in `SKIP_PROFILED`.
- Added host-only fake-VMM coverage for profiling, transition, active remapping,
  and disabled fallback states.
- Verifies a temp-only executable transitions to an active persistent temp set,
  and an unstable parameter cannot disqualify an automatically selected temp.
- Verifies stable parameter and output selection, including non-granularity-sized
  externally owned outputs.
- Verifies failed executions do not advance the three-successful-execution
  profiling window.
- Verifies fallback for unstable and non-VMM-backed parameters and incompatible
  shared parameter buffers.
- Verifies copy protection both during profiling and after transition: the
  owning table retains the finalized external address while execution uses the
  stable reservation address.

## 🧪 Execution Tests:

- Added one-device `VmmTest` coverage using a two-GEMM command-buffer workload
  across profiling, transition, and remapped executions, checking numerical
  correctness on every run.
- Verifies one-time profile selection, automatic temp inclusion, VA-reservation
  creation and reuse, and successful deallocation of externally returned output
  buffers.
- Added a non-VMM allocator execution test verifying the workload remains
  correct and creates no VA reservation when remapping is unavailable.

Copybara import of the project:

--
8b5c4d7 by Shawn Wang <[email protected]>:

[XLA:GPU] Add SKIP_PROFILED command buffer update mode

Add SKIP_PROFILED to profile command-buffer-referenced non-constant allocation addresses for three successful executions, then VA-remap the stable candidates so recorded command buffers can skip updating them.

The per-executor profile transitions once from profiling to active or disabled, preserving the required absent-to-present persistent-index transition. Active executions allocate selected temps at reservation addresses and map finalized selected input/output addresses into an execution-only BufferAllocations copy, leaving the owning table unchanged.

Filter unstable candidates and parameter addresses that are not VMM-backed or are shared by multiple parameter allocations. Add flag parsing, allocation-policy coverage, GPU end-to-end coverage, and fallback behavior without a VMM allocator.

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

Add host-only unit tests for SKIP_PROFILED

Cover the SKIP_PROFILED profile state machine with the fake-VMM harness through a test-only profile-candidate seam.

Verify stable parameters and outputs transition to execution-only reservation addresses while owning buffers remain external, failed executions do not advance profiling, and unstable, non-VMM-backed, or shared parameter buffers disable remapping.

Verify copy-protected addresses are observed after output finalization and can be remapped both during a stable profile and when copy protection first appears after the transition.

--
153df49 by Shawn Wang <[email protected]>:

Remove profile candidate count accessor

--
16473a0 by Shawn Wang <[email protected]>:

[XLA:GPU] Automatically remap temp buffers in SKIP_PROFILED

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

[XLA:GPU] Add direct VA remap allocator includes

Merging this change closes #45597

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45597 from shawnwang18:shawnw/cmd_buffer_skip_profiled_split 7600262
PiperOrigin-RevId: 950773148

…ntinuous_profiling is enabled to bypass the 2GB protobuf limit.

Reverts changelist 949495726

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#45596 from shawnwang18:shawnw/cmd_buffer_va_remap_outputs b7b79af2174271e65f2c7d8e8bac527a72e97af9
PiperOrigin-RevId: 948107319

Imported from GitHub PR #45597

## 📝 Summary of Changes

This PR adds the `SKIP_PROFILED` command-buffer update mode. The mode
automatically includes command-buffer-referenced preallocated temp allocations
in the final VA-remap set while profiling other eligible non-constant
allocations during the first three successful executions. Stable profile
candidates are combined with the temp allocations so subsequent command-buffer
executions can treat the final set as persistent and avoid unnecessary address
updates.

- Adds `DebugOptions::SKIP_PROFILED` and routes it through the GPU executable
  buffer-allocation policy.
- Adds a per-executor profile state machine to
  `GpuExecutableVaRemapAllocator`: `kProfiling` transitions exactly once to
  `kActive` when an automatic temp or safe profile candidate exists, or to
  `kDisabled` when the final remap set is empty.
- Automatically selects preallocated temp buffers and profiles eligible
  parameter and live-out candidates while filtering unstable addresses,
  non-VMM-backed parameters, and incompatible shared parameter buffers.
- Counts only successful executions toward the profiling window and passes
  `std::nullopt` for persistent allocation indices while profiling, followed by
  one absent-to-present transition to the final persistent set.
- Uses the parameter-remapping support from #45595 and output-buffer remapping
  support from #45596. Parameters and outputs remain externally owned in the
  const owning allocation table; finalized external addresses are mapped into
  an execution-only allocation-table copy immediately before execution.
- Preserves the base allocator fallback when VMM remapping is unavailable or
  there are neither automatic temp buffers nor stable profile candidates.

This PR is stacked on #45596, **[XLA:GPU] Support VA remapping of output buffer
allocations**. It is part 3 of 3 and, together with #45595 and #45596,
supersedes the closed pre-split PR #45568.

## 🎯 Justification

`SKIP_TEMP` can stabilize preallocated temporary buffers, but parameters and
live-out buffers may also keep stable physical addresses across repeated
executions. Updating command-buffer node addresses for those stable allocations
adds avoidable runtime overhead.

`SKIP_PROFILED` discovers address stability for each executable/executor pair
instead of assuming every dynamic allocation is stable. Workloads with
repeatable allocation behavior can therefore stabilize more
command-buffer-visible addresses, while unstable or unsupported workloads retain
the existing fallback behavior. Mapping finalized owning addresses immediately
before execution also keeps donation and copy-protection changes compatible with
the const owning-allocation API introduced by the lower PRs.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

No benchmark result is claimed in this draft. The PR adds the policy and
correctness coverage; a public-HLO performance benchmark can be added separately
when reporting workload-specific speedups.

## 🧪 Unit Tests:

- Extended `GpuExecutableTest.CommandBufferAllocationPolicy` expectations for
  automatic temp-buffer selection in `SKIP_PROFILED`.
- Added host-only fake-VMM coverage for profiling, transition, active remapping,
  and disabled fallback states.
- Verifies a temp-only executable transitions to an active persistent temp set,
  and an unstable parameter cannot disqualify an automatically selected temp.
- Verifies stable parameter and output selection, including non-granularity-sized
  externally owned outputs.
- Verifies failed executions do not advance the three-successful-execution
  profiling window.
- Verifies fallback for unstable and non-VMM-backed parameters and incompatible
  shared parameter buffers.
- Verifies copy protection both during profiling and after transition: the
  owning table retains the finalized external address while execution uses the
  stable reservation address.

## 🧪 Execution Tests:

- Added one-device `VmmTest` coverage using a two-GEMM command-buffer workload
  across profiling, transition, and remapped executions, checking numerical
  correctness on every run.
- Verifies one-time profile selection, automatic temp inclusion, VA-reservation
  creation and reuse, and successful deallocation of externally returned output
  buffers.
- Added a non-VMM allocator execution test verifying the workload remains
  correct and creates no VA reservation when remapping is unavailable.

Copybara import of the project:

--
8b5c4d7 by Shawn Wang <[email protected]>:

[XLA:GPU] Add SKIP_PROFILED command buffer update mode

Add SKIP_PROFILED to profile command-buffer-referenced non-constant allocation addresses for three successful executions, then VA-remap the stable candidates so recorded command buffers can skip updating them.

The per-executor profile transitions once from profiling to active or disabled, preserving the required absent-to-present persistent-index transition. Active executions allocate selected temps at reservation addresses and map finalized selected input/output addresses into an execution-only BufferAllocations copy, leaving the owning table unchanged.

Filter unstable candidates and parameter addresses that are not VMM-backed or are shared by multiple parameter allocations. Add flag parsing, allocation-policy coverage, GPU end-to-end coverage, and fallback behavior without a VMM allocator.

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

Add host-only unit tests for SKIP_PROFILED

Cover the SKIP_PROFILED profile state machine with the fake-VMM harness through a test-only profile-candidate seam.

Verify stable parameters and outputs transition to execution-only reservation addresses while owning buffers remain external, failed executions do not advance profiling, and unstable, non-VMM-backed, or shared parameter buffers disable remapping.

Verify copy-protected addresses are observed after output finalization and can be remapped both during a stable profile and when copy protection first appears after the transition.

--
153df49 by Shawn Wang <[email protected]>:

Remove profile candidate count accessor

--
16473a0 by Shawn Wang <[email protected]>:

[XLA:GPU] Automatically remap temp buffers in SKIP_PROFILED

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

[XLA:GPU] Add direct VA remap allocator includes

Merging this change closes #45597

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45597 from shawnwang18:shawnw/cmd_buffer_skip_profiled_split 7600262
PiperOrigin-RevId: 950773148

Imported from GitHub PR #45597

## 📝 Summary of Changes

This PR adds the `SKIP_PROFILED` command-buffer update mode. The mode
automatically includes command-buffer-referenced preallocated temp allocations
in the final VA-remap set while profiling other eligible non-constant
allocations during the first three successful executions. Stable profile
candidates are combined with the temp allocations so subsequent command-buffer
executions can treat the final set as persistent and avoid unnecessary address
updates.

- Adds `DebugOptions::SKIP_PROFILED` and routes it through the GPU executable
  buffer-allocation policy.
- Adds a per-executor profile state machine to
  `GpuExecutableVaRemapAllocator`: `kProfiling` transitions exactly once to
  `kActive` when an automatic temp or safe profile candidate exists, or to
  `kDisabled` when the final remap set is empty.
- Automatically selects preallocated temp buffers and profiles eligible
  parameter and live-out candidates while filtering unstable addresses,
  non-VMM-backed parameters, and incompatible shared parameter buffers.
- Counts only successful executions toward the profiling window and passes
  `std::nullopt` for persistent allocation indices while profiling, followed by
  one absent-to-present transition to the final persistent set.
- Uses the parameter-remapping support from #45595 and output-buffer remapping
  support from #45596. Parameters and outputs remain externally owned in the
  const owning allocation table; finalized external addresses are mapped into
  an execution-only allocation-table copy immediately before execution.
- Preserves the base allocator fallback when VMM remapping is unavailable or
  there are neither automatic temp buffers nor stable profile candidates.

This PR is stacked on #45596, **[XLA:GPU] Support VA remapping of output buffer
allocations**. It is part 3 of 3 and, together with #45595 and #45596,
supersedes the closed pre-split PR #45568.

## 🎯 Justification

`SKIP_TEMP` can stabilize preallocated temporary buffers, but parameters and
live-out buffers may also keep stable physical addresses across repeated
executions. Updating command-buffer node addresses for those stable allocations
adds avoidable runtime overhead.

`SKIP_PROFILED` discovers address stability for each executable/executor pair
instead of assuming every dynamic allocation is stable. Workloads with
repeatable allocation behavior can therefore stabilize more
command-buffer-visible addresses, while unstable or unsupported workloads retain
the existing fallback behavior. Mapping finalized owning addresses immediately
before execution also keeps donation and copy-protection changes compatible with
the const owning-allocation API introduced by the lower PRs.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

No benchmark result is claimed in this draft. The PR adds the policy and
correctness coverage; a public-HLO performance benchmark can be added separately
when reporting workload-specific speedups.

## 🧪 Unit Tests:

- Extended `GpuExecutableTest.CommandBufferAllocationPolicy` expectations for
  automatic temp-buffer selection in `SKIP_PROFILED`.
- Added host-only fake-VMM coverage for profiling, transition, active remapping,
  and disabled fallback states.
- Verifies a temp-only executable transitions to an active persistent temp set,
  and an unstable parameter cannot disqualify an automatically selected temp.
- Verifies stable parameter and output selection, including non-granularity-sized
  externally owned outputs.
- Verifies failed executions do not advance the three-successful-execution
  profiling window.
- Verifies fallback for unstable and non-VMM-backed parameters and incompatible
  shared parameter buffers.
- Verifies copy protection both during profiling and after transition: the
  owning table retains the finalized external address while execution uses the
  stable reservation address.

## 🧪 Execution Tests:

- Added one-device `VmmTest` coverage using a two-GEMM command-buffer workload
  across profiling, transition, and remapped executions, checking numerical
  correctness on every run.
- Verifies one-time profile selection, automatic temp inclusion, VA-reservation
  creation and reuse, and successful deallocation of externally returned output
  buffers.
- Added a non-VMM allocator execution test verifying the workload remains
  correct and creates no VA reservation when remapping is unavailable.

Copybara import of the project:

--
8b5c4d7 by Shawn Wang <[email protected]>:

[XLA:GPU] Add SKIP_PROFILED command buffer update mode

Add SKIP_PROFILED to profile command-buffer-referenced non-constant allocation addresses for three successful executions, then VA-remap the stable candidates so recorded command buffers can skip updating them.

The per-executor profile transitions once from profiling to active or disabled, preserving the required absent-to-present persistent-index transition. Active executions allocate selected temps at reservation addresses and map finalized selected input/output addresses into an execution-only BufferAllocations copy, leaving the owning table unchanged.

Filter unstable candidates and parameter addresses that are not VMM-backed or are shared by multiple parameter allocations. Add flag parsing, allocation-policy coverage, GPU end-to-end coverage, and fallback behavior without a VMM allocator.

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

Add host-only unit tests for SKIP_PROFILED

Cover the SKIP_PROFILED profile state machine with the fake-VMM harness through a test-only profile-candidate seam.

Verify stable parameters and outputs transition to execution-only reservation addresses while owning buffers remain external, failed executions do not advance profiling, and unstable, non-VMM-backed, or shared parameter buffers disable remapping.

Verify copy-protected addresses are observed after output finalization and can be remapped both during a stable profile and when copy protection first appears after the transition.

--
153df49 by Shawn Wang <[email protected]>:

Remove profile candidate count accessor

--
16473a0 by Shawn Wang <[email protected]>:

[XLA:GPU] Automatically remap temp buffers in SKIP_PROFILED

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

[XLA:GPU] Add direct VA remap allocator includes

Merging this change closes #45597

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45597 from shawnwang18:shawnw/cmd_buffer_skip_profiled_split 7600262
PiperOrigin-RevId: 950773148

Imported from GitHub PR #45597

## 📝 Summary of Changes

This PR adds the `SKIP_PROFILED` command-buffer update mode. The mode
automatically includes command-buffer-referenced preallocated temp allocations
in the final VA-remap set while profiling other eligible non-constant
allocations during the first three successful executions. Stable profile
candidates are combined with the temp allocations so subsequent command-buffer
executions can treat the final set as persistent and avoid unnecessary address
updates.

- Adds `DebugOptions::SKIP_PROFILED` and routes it through the GPU executable
  buffer-allocation policy.
- Adds a per-executor profile state machine to
  `GpuExecutableVaRemapAllocator`: `kProfiling` transitions exactly once to
  `kActive` when an automatic temp or safe profile candidate exists, or to
  `kDisabled` when the final remap set is empty.
- Automatically selects preallocated temp buffers and profiles eligible
  parameter and live-out candidates while filtering unstable addresses,
  non-VMM-backed parameters, and incompatible shared parameter buffers.
- Counts only successful executions toward the profiling window and passes
  `std::nullopt` for persistent allocation indices while profiling, followed by
  one absent-to-present transition to the final persistent set.
- Uses the parameter-remapping support from #45595 and output-buffer remapping
  support from #45596. Parameters and outputs remain externally owned in the
  const owning allocation table; finalized external addresses are mapped into
  an execution-only allocation-table copy immediately before execution.
- Preserves the base allocator fallback when VMM remapping is unavailable or
  there are neither automatic temp buffers nor stable profile candidates.

This PR is stacked on #45596, **[XLA:GPU] Support VA remapping of output buffer
allocations**. It is part 3 of 3 and, together with #45595 and #45596,
supersedes the closed pre-split PR #45568.

## 🎯 Justification

`SKIP_TEMP` can stabilize preallocated temporary buffers, but parameters and
live-out buffers may also keep stable physical addresses across repeated
executions. Updating command-buffer node addresses for those stable allocations
adds avoidable runtime overhead.

`SKIP_PROFILED` discovers address stability for each executable/executor pair
instead of assuming every dynamic allocation is stable. Workloads with
repeatable allocation behavior can therefore stabilize more
command-buffer-visible addresses, while unstable or unsupported workloads retain
the existing fallback behavior. Mapping finalized owning addresses immediately
before execution also keeps donation and copy-protection changes compatible with
the const owning-allocation API introduced by the lower PRs.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

No benchmark result is claimed in this draft. The PR adds the policy and
correctness coverage; a public-HLO performance benchmark can be added separately
when reporting workload-specific speedups.

## 🧪 Unit Tests:

- Extended `GpuExecutableTest.CommandBufferAllocationPolicy` expectations for
  automatic temp-buffer selection in `SKIP_PROFILED`.
- Added host-only fake-VMM coverage for profiling, transition, active remapping,
  and disabled fallback states.
- Verifies a temp-only executable transitions to an active persistent temp set,
  and an unstable parameter cannot disqualify an automatically selected temp.
- Verifies stable parameter and output selection, including non-granularity-sized
  externally owned outputs.
- Verifies failed executions do not advance the three-successful-execution
  profiling window.
- Verifies fallback for unstable and non-VMM-backed parameters and incompatible
  shared parameter buffers.
- Verifies copy protection both during profiling and after transition: the
  owning table retains the finalized external address while execution uses the
  stable reservation address.

## 🧪 Execution Tests:

- Added one-device `VmmTest` coverage using a two-GEMM command-buffer workload
  across profiling, transition, and remapped executions, checking numerical
  correctness on every run.
- Verifies one-time profile selection, automatic temp inclusion, VA-reservation
  creation and reuse, and successful deallocation of externally returned output
  buffers.
- Added a non-VMM allocator execution test verifying the workload remains
  correct and creates no VA reservation when remapping is unavailable.

Copybara import of the project:

--
8b5c4d7 by Shawn Wang <[email protected]>:

[XLA:GPU] Add SKIP_PROFILED command buffer update mode

Add SKIP_PROFILED to profile command-buffer-referenced non-constant allocation addresses for three successful executions, then VA-remap the stable candidates so recorded command buffers can skip updating them.

The per-executor profile transitions once from profiling to active or disabled, preserving the required absent-to-present persistent-index transition. Active executions allocate selected temps at reservation addresses and map finalized selected input/output addresses into an execution-only BufferAllocations copy, leaving the owning table unchanged.

Filter unstable candidates and parameter addresses that are not VMM-backed or are shared by multiple parameter allocations. Add flag parsing, allocation-policy coverage, GPU end-to-end coverage, and fallback behavior without a VMM allocator.

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

Add host-only unit tests for SKIP_PROFILED

Cover the SKIP_PROFILED profile state machine with the fake-VMM harness through a test-only profile-candidate seam.

Verify stable parameters and outputs transition to execution-only reservation addresses while owning buffers remain external, failed executions do not advance profiling, and unstable, non-VMM-backed, or shared parameter buffers disable remapping.

Verify copy-protected addresses are observed after output finalization and can be remapped both during a stable profile and when copy protection first appears after the transition.

--
153df49 by Shawn Wang <[email protected]>:

Remove profile candidate count accessor

--
16473a0 by Shawn Wang <[email protected]>:

[XLA:GPU] Automatically remap temp buffers in SKIP_PROFILED

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

[XLA:GPU] Add direct VA remap allocator includes

Merging this change closes #45597

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45597 from shawnwang18:shawnw/cmd_buffer_skip_profiled_split 7600262
PiperOrigin-RevId: 950773148

Imported from GitHub PR #45597

## 📝 Summary of Changes

This PR adds the `SKIP_PROFILED` command-buffer update mode. The mode
automatically includes command-buffer-referenced preallocated temp allocations
in the final VA-remap set while profiling other eligible non-constant
allocations during the first three successful executions. Stable profile
candidates are combined with the temp allocations so subsequent command-buffer
executions can treat the final set as persistent and avoid unnecessary address
updates.

- Adds `DebugOptions::SKIP_PROFILED` and routes it through the GPU executable
  buffer-allocation policy.
- Adds a per-executor profile state machine to
  `GpuExecutableVaRemapAllocator`: `kProfiling` transitions exactly once to
  `kActive` when an automatic temp or safe profile candidate exists, or to
  `kDisabled` when the final remap set is empty.
- Automatically selects preallocated temp buffers and profiles eligible
  parameter and live-out candidates while filtering unstable addresses,
  non-VMM-backed parameters, and incompatible shared parameter buffers.
- Counts only successful executions toward the profiling window and passes
  `std::nullopt` for persistent allocation indices while profiling, followed by
  one absent-to-present transition to the final persistent set.
- Uses the parameter-remapping support from #45595 and output-buffer remapping
  support from #45596. Parameters and outputs remain externally owned in the
  const owning allocation table; finalized external addresses are mapped into
  an execution-only allocation-table copy immediately before execution.
- Preserves the base allocator fallback when VMM remapping is unavailable or
  there are neither automatic temp buffers nor stable profile candidates.

This PR is stacked on #45596, **[XLA:GPU] Support VA remapping of output buffer
allocations**. It is part 3 of 3 and, together with #45595 and #45596,
supersedes the closed pre-split PR #45568.

## 🎯 Justification

`SKIP_TEMP` can stabilize preallocated temporary buffers, but parameters and
live-out buffers may also keep stable physical addresses across repeated
executions. Updating command-buffer node addresses for those stable allocations
adds avoidable runtime overhead.

`SKIP_PROFILED` discovers address stability for each executable/executor pair
instead of assuming every dynamic allocation is stable. Workloads with
repeatable allocation behavior can therefore stabilize more
command-buffer-visible addresses, while unstable or unsupported workloads retain
the existing fallback behavior. Mapping finalized owning addresses immediately
before execution also keeps donation and copy-protection changes compatible with
the const owning-allocation API introduced by the lower PRs.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

No benchmark result is claimed in this draft. The PR adds the policy and
correctness coverage; a public-HLO performance benchmark can be added separately
when reporting workload-specific speedups.

## 🧪 Unit Tests:

- Extended `GpuExecutableTest.CommandBufferAllocationPolicy` expectations for
  automatic temp-buffer selection in `SKIP_PROFILED`.
- Added host-only fake-VMM coverage for profiling, transition, active remapping,
  and disabled fallback states.
- Verifies a temp-only executable transitions to an active persistent temp set,
  and an unstable parameter cannot disqualify an automatically selected temp.
- Verifies stable parameter and output selection, including non-granularity-sized
  externally owned outputs.
- Verifies failed executions do not advance the three-successful-execution
  profiling window.
- Verifies fallback for unstable and non-VMM-backed parameters and incompatible
  shared parameter buffers.
- Verifies copy protection both during profiling and after transition: the
  owning table retains the finalized external address while execution uses the
  stable reservation address.

## 🧪 Execution Tests:

- Added one-device `VmmTest` coverage using a two-GEMM command-buffer workload
  across profiling, transition, and remapped executions, checking numerical
  correctness on every run.
- Verifies one-time profile selection, automatic temp inclusion, VA-reservation
  creation and reuse, and successful deallocation of externally returned output
  buffers.
- Added a non-VMM allocator execution test verifying the workload remains
  correct and creates no VA reservation when remapping is unavailable.

Copybara import of the project:

--
8b5c4d7 by Shawn Wang <[email protected]>:

[XLA:GPU] Add SKIP_PROFILED command buffer update mode

Add SKIP_PROFILED to profile command-buffer-referenced non-constant allocation addresses for three successful executions, then VA-remap the stable candidates so recorded command buffers can skip updating them.

The per-executor profile transitions once from profiling to active or disabled, preserving the required absent-to-present persistent-index transition. Active executions allocate selected temps at reservation addresses and map finalized selected input/output addresses into an execution-only BufferAllocations copy, leaving the owning table unchanged.

Filter unstable candidates and parameter addresses that are not VMM-backed or are shared by multiple parameter allocations. Add flag parsing, allocation-policy coverage, GPU end-to-end coverage, and fallback behavior without a VMM allocator.

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

Add host-only unit tests for SKIP_PROFILED

Cover the SKIP_PROFILED profile state machine with the fake-VMM harness through a test-only profile-candidate seam.

Verify stable parameters and outputs transition to execution-only reservation addresses while owning buffers remain external, failed executions do not advance profiling, and unstable, non-VMM-backed, or shared parameter buffers disable remapping.

Verify copy-protected addresses are observed after output finalization and can be remapped both during a stable profile and when copy protection first appears after the transition.

--
153df49 by Shawn Wang <[email protected]>:

Remove profile candidate count accessor

--
16473a0 by Shawn Wang <[email protected]>:

[XLA:GPU] Automatically remap temp buffers in SKIP_PROFILED

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

[XLA:GPU] Add direct VA remap allocator includes

Merging this change closes #45597

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45597 from shawnwang18:shawnw/cmd_buffer_skip_profiled_split 7600262
PiperOrigin-RevId: 950773148

Imported from GitHub PR #45597

## 📝 Summary of Changes

This PR adds the `SKIP_PROFILED` command-buffer update mode. The mode
automatically includes command-buffer-referenced preallocated temp allocations
in the final VA-remap set while profiling other eligible non-constant
allocations during the first three successful executions. Stable profile
candidates are combined with the temp allocations so subsequent command-buffer
executions can treat the final set as persistent and avoid unnecessary address
updates.

- Adds `DebugOptions::SKIP_PROFILED` and routes it through the GPU executable
  buffer-allocation policy.
- Adds a per-executor profile state machine to
  `GpuExecutableVaRemapAllocator`: `kProfiling` transitions exactly once to
  `kActive` when an automatic temp or safe profile candidate exists, or to
  `kDisabled` when the final remap set is empty.
- Automatically selects preallocated temp buffers and profiles eligible
  parameter and live-out candidates while filtering unstable addresses,
  non-VMM-backed parameters, and incompatible shared parameter buffers.
- Counts only successful executions toward the profiling window and passes
  `std::nullopt` for persistent allocation indices while profiling, followed by
  one absent-to-present transition to the final persistent set.
- Uses the parameter-remapping support from #45595 and output-buffer remapping
  support from #45596. Parameters and outputs remain externally owned in the
  const owning allocation table; finalized external addresses are mapped into
  an execution-only allocation-table copy immediately before execution.
- Preserves the base allocator fallback when VMM remapping is unavailable or
  there are neither automatic temp buffers nor stable profile candidates.

This PR is stacked on #45596, **[XLA:GPU] Support VA remapping of output buffer
allocations**. It is part 3 of 3 and, together with #45595 and #45596,
supersedes the closed pre-split PR #45568.

## 🎯 Justification

`SKIP_TEMP` can stabilize preallocated temporary buffers, but parameters and
live-out buffers may also keep stable physical addresses across repeated
executions. Updating command-buffer node addresses for those stable allocations
adds avoidable runtime overhead.

`SKIP_PROFILED` discovers address stability for each executable/executor pair
instead of assuming every dynamic allocation is stable. Workloads with
repeatable allocation behavior can therefore stabilize more
command-buffer-visible addresses, while unstable or unsupported workloads retain
the existing fallback behavior. Mapping finalized owning addresses immediately
before execution also keeps donation and copy-protection changes compatible with
the const owning-allocation API introduced by the lower PRs.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

No benchmark result is claimed in this draft. The PR adds the policy and
correctness coverage; a public-HLO performance benchmark can be added separately
when reporting workload-specific speedups.

## 🧪 Unit Tests:

- Extended `GpuExecutableTest.CommandBufferAllocationPolicy` expectations for
  automatic temp-buffer selection in `SKIP_PROFILED`.
- Added host-only fake-VMM coverage for profiling, transition, active remapping,
  and disabled fallback states.
- Verifies a temp-only executable transitions to an active persistent temp set,
  and an unstable parameter cannot disqualify an automatically selected temp.
- Verifies stable parameter and output selection, including non-granularity-sized
  externally owned outputs.
- Verifies failed executions do not advance the three-successful-execution
  profiling window.
- Verifies fallback for unstable and non-VMM-backed parameters and incompatible
  shared parameter buffers.
- Verifies copy protection both during profiling and after transition: the
  owning table retains the finalized external address while execution uses the
  stable reservation address.

## 🧪 Execution Tests:

- Added one-device `VmmTest` coverage using a two-GEMM command-buffer workload
  across profiling, transition, and remapped executions, checking numerical
  correctness on every run.
- Verifies one-time profile selection, automatic temp inclusion, VA-reservation
  creation and reuse, and successful deallocation of externally returned output
  buffers.
- Added a non-VMM allocator execution test verifying the workload remains
  correct and creates no VA reservation when remapping is unavailable.

Copybara import of the project:

--
8b5c4d7 by Shawn Wang <[email protected]>:

[XLA:GPU] Add SKIP_PROFILED command buffer update mode

Add SKIP_PROFILED to profile command-buffer-referenced non-constant allocation addresses for three successful executions, then VA-remap the stable candidates so recorded command buffers can skip updating them.

The per-executor profile transitions once from profiling to active or disabled, preserving the required absent-to-present persistent-index transition. Active executions allocate selected temps at reservation addresses and map finalized selected input/output addresses into an execution-only BufferAllocations copy, leaving the owning table unchanged.

Filter unstable candidates and parameter addresses that are not VMM-backed or are shared by multiple parameter allocations. Add flag parsing, allocation-policy coverage, GPU end-to-end coverage, and fallback behavior without a VMM allocator.

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

Add host-only unit tests for SKIP_PROFILED

Cover the SKIP_PROFILED profile state machine with the fake-VMM harness through a test-only profile-candidate seam.

Verify stable parameters and outputs transition to execution-only reservation addresses while owning buffers remain external, failed executions do not advance profiling, and unstable, non-VMM-backed, or shared parameter buffers disable remapping.

Verify copy-protected addresses are observed after output finalization and can be remapped both during a stable profile and when copy protection first appears after the transition.

--
153df49 by Shawn Wang <[email protected]>:

Remove profile candidate count accessor

--
16473a0 by Shawn Wang <[email protected]>:

[XLA:GPU] Automatically remap temp buffers in SKIP_PROFILED

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

[XLA:GPU] Add direct VA remap allocator includes

Merging this change closes #45597

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45597 from shawnwang18:shawnw/cmd_buffer_skip_profiled_split 7600262
PiperOrigin-RevId: 950773148

…ntinuous_profiling is enabled to bypass the 2GB protobuf limit.

Reverts changelist 949495726

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#45596 from shawnwang18:shawnw/cmd_buffer_va_remap_outputs b7b79af2174271e65f2c7d8e8bac527a72e97af9
PiperOrigin-RevId: 948107319

Some targets relied on transitive includes and/or transitive dependencies. This
was found when trying to remove unused dependencies.

Reverts changelist 949495726

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#45596 from shawnwang18:shawnw/cmd_buffer_va_remap_outputs b7b79af2174271e65f2c7d8e8bac527a72e97af9
PiperOrigin-RevId: 938385295

Imported from GitHub PR #45597

## 📝 Summary of Changes

This PR adds the `SKIP_PROFILED` command-buffer update mode. The mode
automatically includes command-buffer-referenced preallocated temp allocations
in the final VA-remap set while profiling other eligible non-constant
allocations during the first three successful executions. Stable profile
candidates are combined with the temp allocations so subsequent command-buffer
executions can treat the final set as persistent and avoid unnecessary address
updates.

- Adds `DebugOptions::SKIP_PROFILED` and routes it through the GPU executable
  buffer-allocation policy.
- Adds a per-executor profile state machine to
  `GpuExecutableVaRemapAllocator`: `kProfiling` transitions exactly once to
  `kActive` when an automatic temp or safe profile candidate exists, or to
  `kDisabled` when the final remap set is empty.
- Automatically selects preallocated temp buffers and profiles eligible
  parameter and live-out candidates while filtering unstable addresses,
  non-VMM-backed parameters, and incompatible shared parameter buffers.
- Counts only successful executions toward the profiling window and passes
  `std::nullopt` for persistent allocation indices while profiling, followed by
  one absent-to-present transition to the final persistent set.
- Uses the parameter-remapping support from #45595 and output-buffer remapping
  support from #45596. Parameters and outputs remain externally owned in the
  const owning allocation table; finalized external addresses are mapped into
  an execution-only allocation-table copy immediately before execution.
- Preserves the base allocator fallback when VMM remapping is unavailable or
  there are neither automatic temp buffers nor stable profile candidates.

This PR is stacked on #45596, **[XLA:GPU] Support VA remapping of output buffer
allocations**. It is part 3 of 3 and, together with #45595 and #45596,
supersedes the closed pre-split PR #45568.

## 🎯 Justification

`SKIP_TEMP` can stabilize preallocated temporary buffers, but parameters and
live-out buffers may also keep stable physical addresses across repeated
executions. Updating command-buffer node addresses for those stable allocations
adds avoidable runtime overhead.

`SKIP_PROFILED` discovers address stability for each executable/executor pair
instead of assuming every dynamic allocation is stable. Workloads with
repeatable allocation behavior can therefore stabilize more
command-buffer-visible addresses, while unstable or unsupported workloads retain
the existing fallback behavior. Mapping finalized owning addresses immediately
before execution also keeps donation and copy-protection changes compatible with
the const owning-allocation API introduced by the lower PRs.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

No benchmark result is claimed in this draft. The PR adds the policy and
correctness coverage; a public-HLO performance benchmark can be added separately
when reporting workload-specific speedups.

## 🧪 Unit Tests:

- Extended `GpuExecutableTest.CommandBufferAllocationPolicy` expectations for
  automatic temp-buffer selection in `SKIP_PROFILED`.
- Added host-only fake-VMM coverage for profiling, transition, active remapping,
  and disabled fallback states.
- Verifies a temp-only executable transitions to an active persistent temp set,
  and an unstable parameter cannot disqualify an automatically selected temp.
- Verifies stable parameter and output selection, including non-granularity-sized
  externally owned outputs.
- Verifies failed executions do not advance the three-successful-execution
  profiling window.
- Verifies fallback for unstable and non-VMM-backed parameters and incompatible
  shared parameter buffers.
- Verifies copy protection both during profiling and after transition: the
  owning table retains the finalized external address while execution uses the
  stable reservation address.

## 🧪 Execution Tests:

- Added one-device `VmmTest` coverage using a two-GEMM command-buffer workload
  across profiling, transition, and remapped executions, checking numerical
  correctness on every run.
- Verifies one-time profile selection, automatic temp inclusion, VA-reservation
  creation and reuse, and successful deallocation of externally returned output
  buffers.
- Added a non-VMM allocator execution test verifying the workload remains
  correct and creates no VA reservation when remapping is unavailable.

Copybara import of the project:

--
8b5c4d7 by Shawn Wang <[email protected]>:

[XLA:GPU] Add SKIP_PROFILED command buffer update mode

Add SKIP_PROFILED to profile command-buffer-referenced non-constant allocation addresses for three successful executions, then VA-remap the stable candidates so recorded command buffers can skip updating them.

The per-executor profile transitions once from profiling to active or disabled, preserving the required absent-to-present persistent-index transition. Active executions allocate selected temps at reservation addresses and map finalized selected input/output addresses into an execution-only BufferAllocations copy, leaving the owning table unchanged.

Filter unstable candidates and parameter addresses that are not VMM-backed or are shared by multiple parameter allocations. Add flag parsing, allocation-policy coverage, GPU end-to-end coverage, and fallback behavior without a VMM allocator.

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

Add host-only unit tests for SKIP_PROFILED

Cover the SKIP_PROFILED profile state machine with the fake-VMM harness through a test-only profile-candidate seam.

Verify stable parameters and outputs transition to execution-only reservation addresses while owning buffers remain external, failed executions do not advance profiling, and unstable, non-VMM-backed, or shared parameter buffers disable remapping.

Verify copy-protected addresses are observed after output finalization and can be remapped both during a stable profile and when copy protection first appears after the transition.

--
153df49 by Shawn Wang <[email protected]>:

Remove profile candidate count accessor

--
16473a0 by Shawn Wang <[email protected]>:

[XLA:GPU] Automatically remap temp buffers in SKIP_PROFILED

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

[XLA:GPU] Add direct VA remap allocator includes

Merging this change closes #45597

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45597 from shawnwang18:shawnw/cmd_buffer_skip_profiled_split 7600262
PiperOrigin-RevId: 950773148

…ntinuous_profiling is enabled to bypass the 2GB protobuf limit.

Reverts changelist 949495726

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#45596 from shawnwang18:shawnw/cmd_buffer_va_remap_outputs b7b79af2174271e65f2c7d8e8bac527a72e97af9
PiperOrigin-RevId: 948107319

Imported from GitHub PR #45597

## 📝 Summary of Changes

This PR adds the `SKIP_PROFILED` command-buffer update mode. The mode
automatically includes command-buffer-referenced preallocated temp allocations
in the final VA-remap set while profiling other eligible non-constant
allocations during the first three successful executions. Stable profile
candidates are combined with the temp allocations so subsequent command-buffer
executions can treat the final set as persistent and avoid unnecessary address
updates.

- Adds `DebugOptions::SKIP_PROFILED` and routes it through the GPU executable
  buffer-allocation policy.
- Adds a per-executor profile state machine to
  `GpuExecutableVaRemapAllocator`: `kProfiling` transitions exactly once to
  `kActive` when an automatic temp or safe profile candidate exists, or to
  `kDisabled` when the final remap set is empty.
- Automatically selects preallocated temp buffers and profiles eligible
  parameter and live-out candidates while filtering unstable addresses,
  non-VMM-backed parameters, and incompatible shared parameter buffers.
- Counts only successful executions toward the profiling window and passes
  `std::nullopt` for persistent allocation indices while profiling, followed by
  one absent-to-present transition to the final persistent set.
- Uses the parameter-remapping support from #45595 and output-buffer remapping
  support from #45596. Parameters and outputs remain externally owned in the
  const owning allocation table; finalized external addresses are mapped into
  an execution-only allocation-table copy immediately before execution.
- Preserves the base allocator fallback when VMM remapping is unavailable or
  there are neither automatic temp buffers nor stable profile candidates.

This PR is stacked on #45596, **[XLA:GPU] Support VA remapping of output buffer
allocations**. It is part 3 of 3 and, together with #45595 and #45596,
supersedes the closed pre-split PR #45568.

## 🎯 Justification

`SKIP_TEMP` can stabilize preallocated temporary buffers, but parameters and
live-out buffers may also keep stable physical addresses across repeated
executions. Updating command-buffer node addresses for those stable allocations
adds avoidable runtime overhead.

`SKIP_PROFILED` discovers address stability for each executable/executor pair
instead of assuming every dynamic allocation is stable. Workloads with
repeatable allocation behavior can therefore stabilize more
command-buffer-visible addresses, while unstable or unsupported workloads retain
the existing fallback behavior. Mapping finalized owning addresses immediately
before execution also keeps donation and copy-protection changes compatible with
the const owning-allocation API introduced by the lower PRs.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

No benchmark result is claimed in this draft. The PR adds the policy and
correctness coverage; a public-HLO performance benchmark can be added separately
when reporting workload-specific speedups.

## 🧪 Unit Tests:

- Extended `GpuExecutableTest.CommandBufferAllocationPolicy` expectations for
  automatic temp-buffer selection in `SKIP_PROFILED`.
- Added host-only fake-VMM coverage for profiling, transition, active remapping,
  and disabled fallback states.
- Verifies a temp-only executable transitions to an active persistent temp set,
  and an unstable parameter cannot disqualify an automatically selected temp.
- Verifies stable parameter and output selection, including non-granularity-sized
  externally owned outputs.
- Verifies failed executions do not advance the three-successful-execution
  profiling window.
- Verifies fallback for unstable and non-VMM-backed parameters and incompatible
  shared parameter buffers.
- Verifies copy protection both during profiling and after transition: the
  owning table retains the finalized external address while execution uses the
  stable reservation address.

## 🧪 Execution Tests:

- Added one-device `VmmTest` coverage using a two-GEMM command-buffer workload
  across profiling, transition, and remapped executions, checking numerical
  correctness on every run.
- Verifies one-time profile selection, automatic temp inclusion, VA-reservation
  creation and reuse, and successful deallocation of externally returned output
  buffers.
- Added a non-VMM allocator execution test verifying the workload remains
  correct and creates no VA reservation when remapping is unavailable.

Copybara import of the project:

--
8b5c4d7 by Shawn Wang <[email protected]>:

[XLA:GPU] Add SKIP_PROFILED command buffer update mode

Add SKIP_PROFILED to profile command-buffer-referenced non-constant allocation addresses for three successful executions, then VA-remap the stable candidates so recorded command buffers can skip updating them.

The per-executor profile transitions once from profiling to active or disabled, preserving the required absent-to-present persistent-index transition. Active executions allocate selected temps at reservation addresses and map finalized selected input/output addresses into an execution-only BufferAllocations copy, leaving the owning table unchanged.

Filter unstable candidates and parameter addresses that are not VMM-backed or are shared by multiple parameter allocations. Add flag parsing, allocation-policy coverage, GPU end-to-end coverage, and fallback behavior without a VMM allocator.

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

Add host-only unit tests for SKIP_PROFILED

Cover the SKIP_PROFILED profile state machine with the fake-VMM harness through a test-only profile-candidate seam.

Verify stable parameters and outputs transition to execution-only reservation addresses while owning buffers remain external, failed executions do not advance profiling, and unstable, non-VMM-backed, or shared parameter buffers disable remapping.

Verify copy-protected addresses are observed after output finalization and can be remapped both during a stable profile and when copy protection first appears after the transition.

--
153df49 by Shawn Wang <[email protected]>:

Remove profile candidate count accessor

--
16473a0 by Shawn Wang <[email protected]>:

[XLA:GPU] Automatically remap temp buffers in SKIP_PROFILED

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

[XLA:GPU] Add direct VA remap allocator includes

Merging this change closes #45597

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45597 from shawnwang18:shawnw/cmd_buffer_skip_profiled_split 7600262
PiperOrigin-RevId: 950773148

Imported from GitHub PR #45597

## 📝 Summary of Changes

This PR adds the `SKIP_PROFILED` command-buffer update mode. The mode
automatically includes command-buffer-referenced preallocated temp allocations
in the final VA-remap set while profiling other eligible non-constant
allocations during the first three successful executions. Stable profile
candidates are combined with the temp allocations so subsequent command-buffer
executions can treat the final set as persistent and avoid unnecessary address
updates.

- Adds `DebugOptions::SKIP_PROFILED` and routes it through the GPU executable
  buffer-allocation policy.
- Adds a per-executor profile state machine to
  `GpuExecutableVaRemapAllocator`: `kProfiling` transitions exactly once to
  `kActive` when an automatic temp or safe profile candidate exists, or to
  `kDisabled` when the final remap set is empty.
- Automatically selects preallocated temp buffers and profiles eligible
  parameter and live-out candidates while filtering unstable addresses,
  non-VMM-backed parameters, and incompatible shared parameter buffers.
- Counts only successful executions toward the profiling window and passes
  `std::nullopt` for persistent allocation indices while profiling, followed by
  one absent-to-present transition to the final persistent set.
- Uses the parameter-remapping support from #45595 and output-buffer remapping
  support from #45596. Parameters and outputs remain externally owned in the
  const owning allocation table; finalized external addresses are mapped into
  an execution-only allocation-table copy immediately before execution.
- Preserves the base allocator fallback when VMM remapping is unavailable or
  there are neither automatic temp buffers nor stable profile candidates.

This PR is stacked on #45596, **[XLA:GPU] Support VA remapping of output buffer
allocations**. It is part 3 of 3 and, together with #45595 and #45596,
supersedes the closed pre-split PR #45568.

## 🎯 Justification

`SKIP_TEMP` can stabilize preallocated temporary buffers, but parameters and
live-out buffers may also keep stable physical addresses across repeated
executions. Updating command-buffer node addresses for those stable allocations
adds avoidable runtime overhead.

`SKIP_PROFILED` discovers address stability for each executable/executor pair
instead of assuming every dynamic allocation is stable. Workloads with
repeatable allocation behavior can therefore stabilize more
command-buffer-visible addresses, while unstable or unsupported workloads retain
the existing fallback behavior. Mapping finalized owning addresses immediately
before execution also keeps donation and copy-protection changes compatible with
the const owning-allocation API introduced by the lower PRs.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

No benchmark result is claimed in this draft. The PR adds the policy and
correctness coverage; a public-HLO performance benchmark can be added separately
when reporting workload-specific speedups.

## 🧪 Unit Tests:

- Extended `GpuExecutableTest.CommandBufferAllocationPolicy` expectations for
  automatic temp-buffer selection in `SKIP_PROFILED`.
- Added host-only fake-VMM coverage for profiling, transition, active remapping,
  and disabled fallback states.
- Verifies a temp-only executable transitions to an active persistent temp set,
  and an unstable parameter cannot disqualify an automatically selected temp.
- Verifies stable parameter and output selection, including non-granularity-sized
  externally owned outputs.
- Verifies failed executions do not advance the three-successful-execution
  profiling window.
- Verifies fallback for unstable and non-VMM-backed parameters and incompatible
  shared parameter buffers.
- Verifies copy protection both during profiling and after transition: the
  owning table retains the finalized external address while execution uses the
  stable reservation address.

## 🧪 Execution Tests:

- Added one-device `VmmTest` coverage using a two-GEMM command-buffer workload
  across profiling, transition, and remapped executions, checking numerical
  correctness on every run.
- Verifies one-time profile selection, automatic temp inclusion, VA-reservation
  creation and reuse, and successful deallocation of externally returned output
  buffers.
- Added a non-VMM allocator execution test verifying the workload remains
  correct and creates no VA reservation when remapping is unavailable.

Copybara import of the project:

--
8b5c4d7 by Shawn Wang <[email protected]>:

[XLA:GPU] Add SKIP_PROFILED command buffer update mode

Add SKIP_PROFILED to profile command-buffer-referenced non-constant allocation addresses for three successful executions, then VA-remap the stable candidates so recorded command buffers can skip updating them.

The per-executor profile transitions once from profiling to active or disabled, preserving the required absent-to-present persistent-index transition. Active executions allocate selected temps at reservation addresses and map finalized selected input/output addresses into an execution-only BufferAllocations copy, leaving the owning table unchanged.

Filter unstable candidates and parameter addresses that are not VMM-backed or are shared by multiple parameter allocations. Add flag parsing, allocation-policy coverage, GPU end-to-end coverage, and fallback behavior without a VMM allocator.

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

Add host-only unit tests for SKIP_PROFILED

Cover the SKIP_PROFILED profile state machine with the fake-VMM harness through a test-only profile-candidate seam.

Verify stable parameters and outputs transition to execution-only reservation addresses while owning buffers remain external, failed executions do not advance profiling, and unstable, non-VMM-backed, or shared parameter buffers disable remapping.

Verify copy-protected addresses are observed after output finalization and can be remapped both during a stable profile and when copy protection first appears after the transition.

--
153df49 by Shawn Wang <[email protected]>:

Remove profile candidate count accessor

--
16473a0 by Shawn Wang <[email protected]>:

[XLA:GPU] Automatically remap temp buffers in SKIP_PROFILED

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

[XLA:GPU] Add direct VA remap allocator includes

Merging this change closes #45597

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45597 from shawnwang18:shawnw/cmd_buffer_skip_profiled_split 7600262
PiperOrigin-RevId: 950773148

Imported from GitHub PR #45597

## 📝 Summary of Changes

This PR adds the `SKIP_PROFILED` command-buffer update mode. The mode
automatically includes command-buffer-referenced preallocated temp allocations
in the final VA-remap set while profiling other eligible non-constant
allocations during the first three successful executions. Stable profile
candidates are combined with the temp allocations so subsequent command-buffer
executions can treat the final set as persistent and avoid unnecessary address
updates.

- Adds `DebugOptions::SKIP_PROFILED` and routes it through the GPU executable
  buffer-allocation policy.
- Adds a per-executor profile state machine to
  `GpuExecutableVaRemapAllocator`: `kProfiling` transitions exactly once to
  `kActive` when an automatic temp or safe profile candidate exists, or to
  `kDisabled` when the final remap set is empty.
- Automatically selects preallocated temp buffers and profiles eligible
  parameter and live-out candidates while filtering unstable addresses,
  non-VMM-backed parameters, and incompatible shared parameter buffers.
- Counts only successful executions toward the profiling window and passes
  `std::nullopt` for persistent allocation indices while profiling, followed by
  one absent-to-present transition to the final persistent set.
- Uses the parameter-remapping support from #45595 and output-buffer remapping
  support from #45596. Parameters and outputs remain externally owned in the
  const owning allocation table; finalized external addresses are mapped into
  an execution-only allocation-table copy immediately before execution.
- Preserves the base allocator fallback when VMM remapping is unavailable or
  there are neither automatic temp buffers nor stable profile candidates.

This PR is stacked on #45596, **[XLA:GPU] Support VA remapping of output buffer
allocations**. It is part 3 of 3 and, together with #45595 and #45596,
supersedes the closed pre-split PR #45568.

## 🎯 Justification

`SKIP_TEMP` can stabilize preallocated temporary buffers, but parameters and
live-out buffers may also keep stable physical addresses across repeated
executions. Updating command-buffer node addresses for those stable allocations
adds avoidable runtime overhead.

`SKIP_PROFILED` discovers address stability for each executable/executor pair
instead of assuming every dynamic allocation is stable. Workloads with
repeatable allocation behavior can therefore stabilize more
command-buffer-visible addresses, while unstable or unsupported workloads retain
the existing fallback behavior. Mapping finalized owning addresses immediately
before execution also keeps donation and copy-protection changes compatible with
the const owning-allocation API introduced by the lower PRs.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

No benchmark result is claimed in this draft. The PR adds the policy and
correctness coverage; a public-HLO performance benchmark can be added separately
when reporting workload-specific speedups.

## 🧪 Unit Tests:

- Extended `GpuExecutableTest.CommandBufferAllocationPolicy` expectations for
  automatic temp-buffer selection in `SKIP_PROFILED`.
- Added host-only fake-VMM coverage for profiling, transition, active remapping,
  and disabled fallback states.
- Verifies a temp-only executable transitions to an active persistent temp set,
  and an unstable parameter cannot disqualify an automatically selected temp.
- Verifies stable parameter and output selection, including non-granularity-sized
  externally owned outputs.
- Verifies failed executions do not advance the three-successful-execution
  profiling window.
- Verifies fallback for unstable and non-VMM-backed parameters and incompatible
  shared parameter buffers.
- Verifies copy protection both during profiling and after transition: the
  owning table retains the finalized external address while execution uses the
  stable reservation address.

## 🧪 Execution Tests:

- Added one-device `VmmTest` coverage using a two-GEMM command-buffer workload
  across profiling, transition, and remapped executions, checking numerical
  correctness on every run.
- Verifies one-time profile selection, automatic temp inclusion, VA-reservation
  creation and reuse, and successful deallocation of externally returned output
  buffers.
- Added a non-VMM allocator execution test verifying the workload remains
  correct and creates no VA reservation when remapping is unavailable.

Copybara import of the project:

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

[XLA:GPU] Add SKIP_PROFILED command buffer update mode

Add SKIP_PROFILED to profile command-buffer-referenced non-constant allocation addresses for three successful executions, then VA-remap the stable candidates so recorded command buffers can skip updating them.

The per-executor profile transitions once from profiling to active or disabled, preserving the required absent-to-present persistent-index transition. Active executions allocate selected temps at reservation addresses and map finalized selected input/output addresses into an execution-only BufferAllocations copy, leaving the owning table unchanged.

Filter unstable candidates and parameter addresses that are not VMM-backed or are shared by multiple parameter allocations. Add flag parsing, allocation-policy coverage, GPU end-to-end coverage, and fallback behavior without a VMM allocator.

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

Add host-only unit tests for SKIP_PROFILED

Cover the SKIP_PROFILED profile state machine with the fake-VMM harness through a test-only profile-candidate seam.

Verify stable parameters and outputs transition to execution-only reservation addresses while owning buffers remain external, failed executions do not advance profiling, and unstable, non-VMM-backed, or shared parameter buffers disable remapping.

Verify copy-protected addresses are observed after output finalization and can be remapped both during a stable profile and when copy protection first appears after the transition.

--
89d3dd2 by Shawn Wang <[email protected]>:

Remove profile candidate count accessor

--
51adad4 by Shawn Wang <[email protected]>:

[XLA:GPU] Automatically remap temp buffers in SKIP_PROFILED

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

[XLA:GPU] Add direct VA remap allocator includes

Merging this change closes #45597

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45597 from shawnwang18:shawnw/cmd_buffer_skip_profiled_split e32b96b
PiperOrigin-RevId: 950773148

Imported from GitHub PR #45597

## 📝 Summary of Changes

This PR adds the `SKIP_PROFILED` command-buffer update mode. The mode
automatically includes command-buffer-referenced preallocated temp allocations
in the final VA-remap set while profiling other eligible non-constant
allocations during the first three successful executions. Stable profile
candidates are combined with the temp allocations so subsequent command-buffer
executions can treat the final set as persistent and avoid unnecessary address
updates.

- Adds `DebugOptions::SKIP_PROFILED` and routes it through the GPU executable
  buffer-allocation policy.
- Adds a per-executor profile state machine to
  `GpuExecutableVaRemapAllocator`: `kProfiling` transitions exactly once to
  `kActive` when an automatic temp or safe profile candidate exists, or to
  `kDisabled` when the final remap set is empty.
- Automatically selects preallocated temp buffers and profiles eligible
  parameter and live-out candidates while filtering unstable addresses,
  non-VMM-backed parameters, and incompatible shared parameter buffers.
- Counts only successful executions toward the profiling window and passes
  `std::nullopt` for persistent allocation indices while profiling, followed by
  one absent-to-present transition to the final persistent set.
- Uses the parameter-remapping support from #45595 and output-buffer remapping
  support from #45596. Parameters and outputs remain externally owned in the
  const owning allocation table; finalized external addresses are mapped into
  an execution-only allocation-table copy immediately before execution.
- Preserves the base allocator fallback when VMM remapping is unavailable or
  there are neither automatic temp buffers nor stable profile candidates.

This PR is stacked on #45596, **[XLA:GPU] Support VA remapping of output buffer
allocations**. It is part 3 of 3 and, together with #45595 and #45596,
supersedes the closed pre-split PR #45568.

## 🎯 Justification

`SKIP_TEMP` can stabilize preallocated temporary buffers, but parameters and
live-out buffers may also keep stable physical addresses across repeated
executions. Updating command-buffer node addresses for those stable allocations
adds avoidable runtime overhead.

`SKIP_PROFILED` discovers address stability for each executable/executor pair
instead of assuming every dynamic allocation is stable. Workloads with
repeatable allocation behavior can therefore stabilize more
command-buffer-visible addresses, while unstable or unsupported workloads retain
the existing fallback behavior. Mapping finalized owning addresses immediately
before execution also keeps donation and copy-protection changes compatible with
the const owning-allocation API introduced by the lower PRs.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

No benchmark result is claimed in this draft. The PR adds the policy and
correctness coverage; a public-HLO performance benchmark can be added separately
when reporting workload-specific speedups.

## 🧪 Unit Tests:

- Extended `GpuExecutableTest.CommandBufferAllocationPolicy` expectations for
  automatic temp-buffer selection in `SKIP_PROFILED`.
- Added host-only fake-VMM coverage for profiling, transition, active remapping,
  and disabled fallback states.
- Verifies a temp-only executable transitions to an active persistent temp set,
  and an unstable parameter cannot disqualify an automatically selected temp.
- Verifies stable parameter and output selection, including non-granularity-sized
  externally owned outputs.
- Verifies failed executions do not advance the three-successful-execution
  profiling window.
- Verifies fallback for unstable and non-VMM-backed parameters and incompatible
  shared parameter buffers.
- Verifies copy protection both during profiling and after transition: the
  owning table retains the finalized external address while execution uses the
  stable reservation address.

## 🧪 Execution Tests:

- Added one-device `VmmTest` coverage using a two-GEMM command-buffer workload
  across profiling, transition, and remapped executions, checking numerical
  correctness on every run.
- Verifies one-time profile selection, automatic temp inclusion, VA-reservation
  creation and reuse, and successful deallocation of externally returned output
  buffers.
- Added a non-VMM allocator execution test verifying the workload remains
  correct and creates no VA reservation when remapping is unavailable.

Copybara import of the project:

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

[XLA:GPU] Add SKIP_PROFILED command buffer update mode

Add SKIP_PROFILED to profile command-buffer-referenced non-constant allocation addresses for three successful executions, then VA-remap the stable candidates so recorded command buffers can skip updating them.

The per-executor profile transitions once from profiling to active or disabled, preserving the required absent-to-present persistent-index transition. Active executions allocate selected temps at reservation addresses and map finalized selected input/output addresses into an execution-only BufferAllocations copy, leaving the owning table unchanged.

Filter unstable candidates and parameter addresses that are not VMM-backed or are shared by multiple parameter allocations. Add flag parsing, allocation-policy coverage, GPU end-to-end coverage, and fallback behavior without a VMM allocator.

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

Add host-only unit tests for SKIP_PROFILED

Cover the SKIP_PROFILED profile state machine with the fake-VMM harness through a test-only profile-candidate seam.

Verify stable parameters and outputs transition to execution-only reservation addresses while owning buffers remain external, failed executions do not advance profiling, and unstable, non-VMM-backed, or shared parameter buffers disable remapping.

Verify copy-protected addresses are observed after output finalization and can be remapped both during a stable profile and when copy protection first appears after the transition.

--
89d3dd2 by Shawn Wang <[email protected]>:

Remove profile candidate count accessor

--
51adad4 by Shawn Wang <[email protected]>:

[XLA:GPU] Automatically remap temp buffers in SKIP_PROFILED

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

[XLA:GPU] Add direct VA remap allocator includes

Merging this change closes #45597

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45597 from shawnwang18:shawnw/cmd_buffer_skip_profiled_split e32b96b
PiperOrigin-RevId: 950773148

Imported from GitHub PR #45597

## 📝 Summary of Changes

This PR adds the `SKIP_PROFILED` command-buffer update mode. The mode
automatically includes command-buffer-referenced preallocated temp allocations
in the final VA-remap set while profiling other eligible non-constant
allocations during the first three successful executions. Stable profile
candidates are combined with the temp allocations so subsequent command-buffer
executions can treat the final set as persistent and avoid unnecessary address
updates.

- Adds `DebugOptions::SKIP_PROFILED` and routes it through the GPU executable
  buffer-allocation policy.
- Adds a per-executor profile state machine to
  `GpuExecutableVaRemapAllocator`: `kProfiling` transitions exactly once to
  `kActive` when an automatic temp or safe profile candidate exists, or to
  `kDisabled` when the final remap set is empty.
- Automatically selects preallocated temp buffers and profiles eligible
  parameter and live-out candidates while filtering unstable addresses,
  non-VMM-backed parameters, and incompatible shared parameter buffers.
- Counts only successful executions toward the profiling window and passes
  `std::nullopt` for persistent allocation indices while profiling, followed by
  one absent-to-present transition to the final persistent set.
- Uses the parameter-remapping support from #45595 and output-buffer remapping
  support from #45596. Parameters and outputs remain externally owned in the
  const owning allocation table; finalized external addresses are mapped into
  an execution-only allocation-table copy immediately before execution.
- Preserves the base allocator fallback when VMM remapping is unavailable or
  there are neither automatic temp buffers nor stable profile candidates.

This PR is stacked on #45596, **[XLA:GPU] Support VA remapping of output buffer
allocations**. It is part 3 of 3 and, together with #45595 and #45596,
supersedes the closed pre-split PR #45568.

## 🎯 Justification

`SKIP_TEMP` can stabilize preallocated temporary buffers, but parameters and
live-out buffers may also keep stable physical addresses across repeated
executions. Updating command-buffer node addresses for those stable allocations
adds avoidable runtime overhead.

`SKIP_PROFILED` discovers address stability for each executable/executor pair
instead of assuming every dynamic allocation is stable. Workloads with
repeatable allocation behavior can therefore stabilize more
command-buffer-visible addresses, while unstable or unsupported workloads retain
the existing fallback behavior. Mapping finalized owning addresses immediately
before execution also keeps donation and copy-protection changes compatible with
the const owning-allocation API introduced by the lower PRs.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

No benchmark result is claimed in this draft. The PR adds the policy and
correctness coverage; a public-HLO performance benchmark can be added separately
when reporting workload-specific speedups.

## 🧪 Unit Tests:

- Extended `GpuExecutableTest.CommandBufferAllocationPolicy` expectations for
  automatic temp-buffer selection in `SKIP_PROFILED`.
- Added host-only fake-VMM coverage for profiling, transition, active remapping,
  and disabled fallback states.
- Verifies a temp-only executable transitions to an active persistent temp set,
  and an unstable parameter cannot disqualify an automatically selected temp.
- Verifies stable parameter and output selection, including non-granularity-sized
  externally owned outputs.
- Verifies failed executions do not advance the three-successful-execution
  profiling window.
- Verifies fallback for unstable and non-VMM-backed parameters and incompatible
  shared parameter buffers.
- Verifies copy protection both during profiling and after transition: the
  owning table retains the finalized external address while execution uses the
  stable reservation address.

## 🧪 Execution Tests:

- Added one-device `VmmTest` coverage using a two-GEMM command-buffer workload
  across profiling, transition, and remapped executions, checking numerical
  correctness on every run.
- Verifies one-time profile selection, automatic temp inclusion, VA-reservation
  creation and reuse, and successful deallocation of externally returned output
  buffers.
- Added a non-VMM allocator execution test verifying the workload remains
  correct and creates no VA reservation when remapping is unavailable.

Copybara import of the project:

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

[XLA:GPU] Add SKIP_PROFILED command buffer update mode

Add SKIP_PROFILED to profile command-buffer-referenced non-constant allocation addresses for three successful executions, then VA-remap the stable candidates so recorded command buffers can skip updating them.

The per-executor profile transitions once from profiling to active or disabled, preserving the required absent-to-present persistent-index transition. Active executions allocate selected temps at reservation addresses and map finalized selected input/output addresses into an execution-only BufferAllocations copy, leaving the owning table unchanged.

Filter unstable candidates and parameter addresses that are not VMM-backed or are shared by multiple parameter allocations. Add flag parsing, allocation-policy coverage, GPU end-to-end coverage, and fallback behavior without a VMM allocator.

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

Add host-only unit tests for SKIP_PROFILED

Cover the SKIP_PROFILED profile state machine with the fake-VMM harness through a test-only profile-candidate seam.

Verify stable parameters and outputs transition to execution-only reservation addresses while owning buffers remain external, failed executions do not advance profiling, and unstable, non-VMM-backed, or shared parameter buffers disable remapping.

Verify copy-protected addresses are observed after output finalization and can be remapped both during a stable profile and when copy protection first appears after the transition.

--
89d3dd2 by Shawn Wang <[email protected]>:

Remove profile candidate count accessor

--
51adad4 by Shawn Wang <[email protected]>:

[XLA:GPU] Automatically remap temp buffers in SKIP_PROFILED

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

[XLA:GPU] Add direct VA remap allocator includes

Merging this change closes #45597

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45597 from shawnwang18:shawnw/cmd_buffer_skip_profiled_split e32b96b
PiperOrigin-RevId: 950773148

Imported from GitHub PR #45597

## 📝 Summary of Changes

This PR adds the `SKIP_PROFILED` command-buffer update mode. The mode
automatically includes command-buffer-referenced preallocated temp allocations
in the final VA-remap set while profiling other eligible non-constant
allocations during the first three successful executions. Stable profile
candidates are combined with the temp allocations so subsequent command-buffer
executions can treat the final set as persistent and avoid unnecessary address
updates.

- Adds `DebugOptions::SKIP_PROFILED` and routes it through the GPU executable
  buffer-allocation policy.
- Adds a per-executor profile state machine to
  `GpuExecutableVaRemapAllocator`: `kProfiling` transitions exactly once to
  `kActive` when an automatic temp or safe profile candidate exists, or to
  `kDisabled` when the final remap set is empty.
- Automatically selects preallocated temp buffers and profiles eligible
  parameter and live-out candidates while filtering unstable addresses,
  non-VMM-backed parameters, and incompatible shared parameter buffers.
- Counts only successful executions toward the profiling window and passes
  `std::nullopt` for persistent allocation indices while profiling, followed by
  one absent-to-present transition to the final persistent set.
- Uses the parameter-remapping support from #45595 and output-buffer remapping
  support from #45596. Parameters and outputs remain externally owned in the
  const owning allocation table; finalized external addresses are mapped into
  an execution-only allocation-table copy immediately before execution.
- Preserves the base allocator fallback when VMM remapping is unavailable or
  there are neither automatic temp buffers nor stable profile candidates.

This PR is stacked on #45596, **[XLA:GPU] Support VA remapping of output buffer
allocations**. It is part 3 of 3 and, together with #45595 and #45596,
supersedes the closed pre-split PR #45568.

## 🎯 Justification

`SKIP_TEMP` can stabilize preallocated temporary buffers, but parameters and
live-out buffers may also keep stable physical addresses across repeated
executions. Updating command-buffer node addresses for those stable allocations
adds avoidable runtime overhead.

`SKIP_PROFILED` discovers address stability for each executable/executor pair
instead of assuming every dynamic allocation is stable. Workloads with
repeatable allocation behavior can therefore stabilize more
command-buffer-visible addresses, while unstable or unsupported workloads retain
the existing fallback behavior. Mapping finalized owning addresses immediately
before execution also keeps donation and copy-protection changes compatible with
the const owning-allocation API introduced by the lower PRs.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

No benchmark result is claimed in this draft. The PR adds the policy and
correctness coverage; a public-HLO performance benchmark can be added separately
when reporting workload-specific speedups.

## 🧪 Unit Tests:

- Extended `GpuExecutableTest.CommandBufferAllocationPolicy` expectations for
  automatic temp-buffer selection in `SKIP_PROFILED`.
- Added host-only fake-VMM coverage for profiling, transition, active remapping,
  and disabled fallback states.
- Verifies a temp-only executable transitions to an active persistent temp set,
  and an unstable parameter cannot disqualify an automatically selected temp.
- Verifies stable parameter and output selection, including non-granularity-sized
  externally owned outputs.
- Verifies failed executions do not advance the three-successful-execution
  profiling window.
- Verifies fallback for unstable and non-VMM-backed parameters and incompatible
  shared parameter buffers.
- Verifies copy protection both during profiling and after transition: the
  owning table retains the finalized external address while execution uses the
  stable reservation address.

## 🧪 Execution Tests:

- Added one-device `VmmTest` coverage using a two-GEMM command-buffer workload
  across profiling, transition, and remapped executions, checking numerical
  correctness on every run.
- Verifies one-time profile selection, automatic temp inclusion, VA-reservation
  creation and reuse, and successful deallocation of externally returned output
  buffers.
- Added a non-VMM allocator execution test verifying the workload remains
  correct and creates no VA reservation when remapping is unavailable.

Copybara import of the project:

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

[XLA:GPU] Add SKIP_PROFILED command buffer update mode

Add SKIP_PROFILED to profile command-buffer-referenced non-constant allocation addresses for three successful executions, then VA-remap the stable candidates so recorded command buffers can skip updating them.

The per-executor profile transitions once from profiling to active or disabled, preserving the required absent-to-present persistent-index transition. Active executions allocate selected temps at reservation addresses and map finalized selected input/output addresses into an execution-only BufferAllocations copy, leaving the owning table unchanged.

Filter unstable candidates and parameter addresses that are not VMM-backed or are shared by multiple parameter allocations. Add flag parsing, allocation-policy coverage, GPU end-to-end coverage, and fallback behavior without a VMM allocator.

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

Add host-only unit tests for SKIP_PROFILED

Cover the SKIP_PROFILED profile state machine with the fake-VMM harness through a test-only profile-candidate seam.

Verify stable parameters and outputs transition to execution-only reservation addresses while owning buffers remain external, failed executions do not advance profiling, and unstable, non-VMM-backed, or shared parameter buffers disable remapping.

Verify copy-protected addresses are observed after output finalization and can be remapped both during a stable profile and when copy protection first appears after the transition.

--
89d3dd2 by Shawn Wang <[email protected]>:

Remove profile candidate count accessor

--
51adad4 by Shawn Wang <[email protected]>:

[XLA:GPU] Automatically remap temp buffers in SKIP_PROFILED

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

[XLA:GPU] Add direct VA remap allocator includes

Merging this change closes #45597

FUTURE_COPYBARA_INTEGRATE_REVIEW=#45597 from shawnwang18:shawnw/cmd_buffer_skip_profiled_split e32b96b
PiperOrigin-RevId: 950773148

Imported from GitHub PR #45597

## 📝 Summary of Changes

This PR adds the `SKIP_PROFILED` command-buffer update mode. The mode
automatically includes command-buffer-referenced preallocated temp allocations
in the final VA-remap set while profiling other eligible non-constant
allocations during the first three successful executions. Stable profile
candidates are combined with the temp allocations so subsequent command-buffer
executions can treat the final set as persistent and avoid unnecessary address
updates.

- Adds `DebugOptions::SKIP_PROFILED` and routes it through the GPU executable
  buffer-allocation policy.
- Adds a per-executor profile state machine to
  `GpuExecutableVaRemapAllocator`: `kProfiling` transitions exactly once to
  `kActive` when an automatic temp or safe profile candidate exists, or to
  `kDisabled` when the final remap set is empty.
- Automatically selects preallocated temp buffers and profiles eligible
  parameter and live-out candidates while filtering unstable addresses,
  non-VMM-backed parameters, and incompatible shared parameter buffers.
- Counts only successful executions toward the profiling window and passes
  `std::nullopt` for persistent allocation indices while profiling, followed by
  one absent-to-present transition to the final persistent set.
- Uses the parameter-remapping support from #45595 and output-buffer remapping
  support from #45596. Parameters and outputs remain externally owned in the
  const owning allocation table; finalized external addresses are mapped into
  an execution-only allocation-table copy immediately before execution.
- Preserves the base allocator fallback when VMM remapping is unavailable or
  there are neither automatic temp buffers nor stable profile candidates.

This PR is stacked on #45596, **[XLA:GPU] Support VA remapping of output buffer
allocations**. It is part 3 of 3 and, together with #45595 and #45596,
supersedes the closed pre-split PR #45568.

## 🎯 Justification

`SKIP_TEMP` can stabilize preallocated temporary buffers, but parameters and
live-out buffers may also keep stable physical addresses across repeated
executions. Updating command-buffer node addresses for those stable allocations
adds avoidable runtime overhead.

`SKIP_PROFILED` discovers address stability for each executable/executor pair
instead of assuming every dynamic allocation is stable. Workloads with
repeatable allocation behavior can therefore stabilize more
command-buffer-visible addresses, while unstable or unsupported workloads retain
the existing fallback behavior. Mapping finalized owning addresses immediately
before execution also keeps donation and copy-protection changes compatible with
the const owning-allocation API introduced by the lower PRs.

## 🚀 Kind of Contribution

✨ New Feature

## 📊 Benchmark (for Performance Improvements)

No benchmark result is claimed in this draft. The PR adds the policy and
correctness coverage; a public-HLO performance benchmark can be added separately
when reporting workload-specific speedups.

## 🧪 Unit Tests:

- Extended `GpuExecutableTest.CommandBufferAllocationPolicy` expectations for
  automatic temp-buffer selection in `SKIP_PROFILED`.
- Added host-only fake-VMM coverage for profiling, transition, active remapping,
  and disabled fallback states.
- Verifies a temp-only executable transitions to an active persistent temp set,
  and an unstable parameter cannot disqualify an automatically selected temp.
- Verifies stable parameter and output selection, including non-granularity-sized
  externally owned outputs.
- Verifies failed executions do not advance the three-successful-execution
  profiling window.
- Verifies fallback for unstable and non-VMM-backed parameters and incompatible
  shared parameter buffers.
- Verifies copy protection both during profiling and after transition: the
  owning table retains the finalized external address while execution uses the
  stable reservation address.

## 🧪 Execution Tests:

- Added one-device `VmmTest` coverage using a two-GEMM command-buffer workload
  across profiling, transition, and remapped executions, checking numerical
  correctness on every run.
- Verifies one-time profile selection, automatic temp inclusion, VA-reservation
  creation and reuse, and successful deallocation of externally returned output
  buffers.
- Added a non-VMM allocator execution test verifying the workload remains
  correct and creates no VA reservation when remapping is unavailable.

Copybara import of the project:

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

[XLA:GPU] Add SKIP_PROFILED command buffer update mode

Add SKIP_PROFILED to profile command-buffer-referenced non-constant allocation addresses for three successful executions, then VA-remap the stable candidates so recorded command buffers can skip updating them.

The per-executor profile transitions once from profiling to active or disabled, preserving the required absent-to-present persistent-index transition. Active executions allocate selected temps at reservation addresses and map finalized selected input/output addresses into an execution-only BufferAllocations copy, leaving the owning table unchanged.

Filter unstable candidates and parameter addresses that are not VMM-backed or are shared by multiple parameter allocations. Add flag parsing, allocation-policy coverage, GPU end-to-end coverage, and fallback behavior without a VMM allocator.

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

Add host-only unit tests for SKIP_PROFILED

Cover the SKIP_PROFILED profile state machine with the fake-VMM harness through a test-only profile-candidate seam.

Verify stable parameters and outputs transition to execution-only reservation addresses while owning buffers remain external, failed executions do not advance profiling, and unstable, non-VMM-backed, or shared parameter buffers disable remapping.

Verify copy-protected addresses are observed after output finalization and can be remapped both during a stable profile and when copy protection first appears after the transition.

--
89d3dd2 by Shawn Wang <[email protected]>:

Remove profile candidate count accessor

--
51adad4 by Shawn Wang <[email protected]>:

[XLA:GPU] Automatically remap temp buffers in SKIP_PROFILED

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

[XLA:GPU] Add direct VA remap allocator includes

Merging this change closes #45597

COPYBARA_INTEGRATE_REVIEW=#45597 from shawnwang18:shawnw/cmd_buffer_skip_profiled_split e32b96b
PiperOrigin-RevId: 952017749