PR #45597: [XLA:GPU] Add SKIP_PROFILED command buffer update mode - #123628

Imported from GitHub PR openxla/xla#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:

--
f53c11fbb9238a010cdfb8effe8df509a0fb6f31 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.

--
d76dd3608493868568470916e2f90c67f8683bb2 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.

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

Remove profile candidate count accessor

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

[XLA:GPU] Automatically remap temp buffers in SKIP_PROFILED

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

[XLA:GPU] Add direct VA remap allocator includes

Merging this change closes #45597

PiperOrigin-RevId: 952017749