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

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

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