Skip to content

r2.4: Update TFLite CMake build rules - #45130

Closed
terryheo wants to merge 3 commits into
tensorflow:r2.4from
terryheo:fix-r2.4-cmake-opencl
Closed

terryheo wants to merge 3 commits into
tensorflow:r2.4from
terryheo:fix-r2.4-cmake-opencl

Conversation

@terryheo

Copy link
Copy Markdown
Member

No description provided.

$ cmake ../tensorflow/lite/
$ cmake --build . -t minimal

PiperOrigin-RevId: 339585522
Change-Id: I58511caa6b0115ea0eebd6b9c51b9f9076766f2a
Created a CMakeLists.txt under the lite/examples/minimal

Usage)

$ git clone https://github.com/tensorflow/tensorflow.git tensorflow_src
$ mkdir minimal_build
$ cd miniaml_build
$ cmake ../tensorflow_src/tensorflow/lite/examples/minimal
$ cmake --build . -j

PiperOrigin-RevId: 339977906
Change-Id: I85cada5430407e72354d0a31a17584c48407db2f
Release build uses "-O3 -DNDEBUG" flags.
Also updated tools/cmake/README.md

PiperOrigin-RevId: 343415907
Change-Id: If8bc6b15febc79d7d982e972e9fc14ab8f8dbbff
@google-cla google-cla Bot added the cla: yes label Nov 24, 2020
@terryheo terryheo closed this Nov 24, 2020
@terryheo terryheo reopened this Nov 24, 2020
@terryheo terryheo closed this Nov 24, 2020
@terryheo
terryheo deleted the fix-r2.4-cmake-opencl branch November 24, 2020 01:12
@terryheo
terryheo restored the fix-r2.4-cmake-opencl branch November 24, 2020 01:20
copybara-service Bot pushed a commit that referenced this pull request Jul 14, 2026
…llocations set changed.

Imported from GitHub PR openxla/xla#45130

For command buffer run, add the required processing for allowing `persistent_alloc_indices` to transition only from absent to present across execution steps and must remain unchanged once present.

We may require this transition because the chosen of buffer allocations whose address are in `persistent_alloc_indices` might not be available from the first step, e.g, if persistent allocations are only available after several steps of profiling.  So we need to make sure that command buffer is correctly updated when `persistent_alloc_indices` is changed across steps.

📝 Summary of Changes

- Document that `persistent_alloc_indices` may transition only from absent to
  present across execution steps and must remain unchanged once present.
- Track only whether persistent-allocation information was valid when each
  executor's command buffer was recorded, avoiding an owned vector copy and
  repeated value comparisons.
- Force a full command-buffer update in `CommandBufferThunk::Initialize` when
  persistent allocation information first becomes valid.
- In `ExecuteOnStream`, recover from a recorded/current information-validity
  mismatch by fully updating the command buffer. This covers a
  profiling-skipped update and interleaving with an older execution initialized
  before policy activation.
- Use the public `ShapedSlice` header and its direct BUILD dependency.

🎯 Justification

The persistent allocation policy controls which recorded addresses a command
buffer can safely retain. When the policy first becomes available, an
allocation-filtered update can skip a command whose allocations are now all
persistent, leaving stale addresses in the command buffer. A full update fixes
that transition.

Although the policy transition is logically one-way, executions initialized on
opposite sides of it can interleave on the shared per-executor command buffer.
An Execute-time validity check therefore forces a full update for either
recorded/current mismatch without comparing policy vectors on every step.

🚀 Kind of Contribution

🐛 Bug Fix

📊 Benchmark (for Performance Improvements)

N/A. This is a correctness fix.

🧪 Unit Tests:

Extended `PersistentAllocIndicesBecomingAvailableTriggersUpdate` to verify:

- absent-to-present activation during Initialize triggers exactly one update;
- Execute recovers when the initialization update was skipped;
- an older absent-policy execution and the following present-policy execution
  each restore the correct recorded addresses when they interleave.

Also ran `GpuExecutableTest.CommandBufferAllocationPolicy`.

🧪 Execution Tests:

Ran the focused command-buffer regression and allocation-policy test with
Bazel. All seven configured GPU architecture variants of
`command_buffer_thunk_test` passed, and `gpu_executable_test` passed.

Copybara import of the project:

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

Update command buffers when allocation policy changes

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

Restrict persistent allocation policy transitions

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

Fix persistent allocation policy update fallback

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

Clarify persistent allocation validity names

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

Remove debug-only allocator mode checks

--
87cc9025b69cced36607307bb0301b094395f27e by Shawn Wang <[email protected]>:

Reuse counting memset test helper

--
24283bd9b477aad1835a9cd58c0889f5abfe9b7d by Shawn Wang <[email protected]>:

Revert "Reuse counting memset test helper"

This reverts commit 87cc9025b69cced36607307bb0301b094395f27e.

Merging this change closes #45130

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#45130 from shawnwang18:shawnw/trigers_update_on_update_policy_change 24283bd9b477aad1835a9cd58c0889f5abfe9b7d
PiperOrigin-RevId: 947620409
copybara-service Bot pushed a commit that referenced this pull request Jul 14, 2026
…llocations set changed.

Imported from GitHub PR openxla/xla#45130

For command buffer run, add the required processing for allowing `persistent_alloc_indices` to transition only from absent to present across execution steps and must remain unchanged once present.

We may require this transition because the chosen of buffer allocations whose address are in `persistent_alloc_indices` might not be available from the first step, e.g, if persistent allocations are only available after several steps of profiling.  So we need to make sure that command buffer is correctly updated when `persistent_alloc_indices` is changed across steps.

📝 Summary of Changes

- Document that `persistent_alloc_indices` may transition only from absent to
  present across execution steps and must remain unchanged once present.
- Track only whether persistent-allocation information was valid when each
  executor's command buffer was recorded, avoiding an owned vector copy and
  repeated value comparisons.
- Force a full command-buffer update in `CommandBufferThunk::Initialize` when
  persistent allocation information first becomes valid.
- In `ExecuteOnStream`, recover from a recorded/current information-validity
  mismatch by fully updating the command buffer. This covers a
  profiling-skipped update and interleaving with an older execution initialized
  before policy activation.
- Use the public `ShapedSlice` header and its direct BUILD dependency.

🎯 Justification

The persistent allocation policy controls which recorded addresses a command
buffer can safely retain. When the policy first becomes available, an
allocation-filtered update can skip a command whose allocations are now all
persistent, leaving stale addresses in the command buffer. A full update fixes
that transition.

Although the policy transition is logically one-way, executions initialized on
opposite sides of it can interleave on the shared per-executor command buffer.
An Execute-time validity check therefore forces a full update for either
recorded/current mismatch without comparing policy vectors on every step.

🚀 Kind of Contribution

🐛 Bug Fix

📊 Benchmark (for Performance Improvements)

N/A. This is a correctness fix.

🧪 Unit Tests:

Extended `PersistentAllocIndicesBecomingAvailableTriggersUpdate` to verify:

- absent-to-present activation during Initialize triggers exactly one update;
- Execute recovers when the initialization update was skipped;
- an older absent-policy execution and the following present-policy execution
  each restore the correct recorded addresses when they interleave.

Also ran `GpuExecutableTest.CommandBufferAllocationPolicy`.

🧪 Execution Tests:

Ran the focused command-buffer regression and allocation-policy test with
Bazel. All seven configured GPU architecture variants of
`command_buffer_thunk_test` passed, and `gpu_executable_test` passed.

Copybara import of the project:

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

Update command buffers when allocation policy changes

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

Restrict persistent allocation policy transitions

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

Fix persistent allocation policy update fallback

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

Clarify persistent allocation validity names

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

Remove debug-only allocator mode checks

--
87cc9025b69cced36607307bb0301b094395f27e by Shawn Wang <[email protected]>:

Reuse counting memset test helper

--
24283bd9b477aad1835a9cd58c0889f5abfe9b7d by Shawn Wang <[email protected]>:

Revert "Reuse counting memset test helper"

This reverts commit 87cc9025b69cced36607307bb0301b094395f27e.

Merging this change closes #45130

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#45130 from shawnwang18:shawnw/trigers_update_on_update_policy_change 24283bd9b477aad1835a9cd58c0889f5abfe9b7d
PiperOrigin-RevId: 947620409
copybara-service Bot pushed a commit that referenced this pull request Jul 14, 2026
…llocations set changed.

Imported from GitHub PR openxla/xla#45130

For command buffer run, add the required processing for allowing `persistent_alloc_indices` to transition only from absent to present across execution steps and must remain unchanged once present.

We may require this transition because the chosen of buffer allocations whose address are in `persistent_alloc_indices` might not be available from the first step, e.g, if persistent allocations are only available after several steps of profiling.  So we need to make sure that command buffer is correctly updated when `persistent_alloc_indices` is changed across steps.

📝 Summary of Changes

- Document that `persistent_alloc_indices` may transition only from absent to
  present across execution steps and must remain unchanged once present.
- Track only whether persistent-allocation information was valid when each
  executor's command buffer was recorded, avoiding an owned vector copy and
  repeated value comparisons.
- Force a full command-buffer update in `CommandBufferThunk::Initialize` when
  persistent allocation information first becomes valid.
- In `ExecuteOnStream`, recover from a recorded/current information-validity
  mismatch by fully updating the command buffer. This covers a
  profiling-skipped update and interleaving with an older execution initialized
  before policy activation.
- Use the public `ShapedSlice` header and its direct BUILD dependency.

🎯 Justification

The persistent allocation policy controls which recorded addresses a command
buffer can safely retain. When the policy first becomes available, an
allocation-filtered update can skip a command whose allocations are now all
persistent, leaving stale addresses in the command buffer. A full update fixes
that transition.

Although the policy transition is logically one-way, executions initialized on
opposite sides of it can interleave on the shared per-executor command buffer.
An Execute-time validity check therefore forces a full update for either
recorded/current mismatch without comparing policy vectors on every step.

🚀 Kind of Contribution

🐛 Bug Fix

📊 Benchmark (for Performance Improvements)

N/A. This is a correctness fix.

🧪 Unit Tests:

Extended `PersistentAllocIndicesBecomingAvailableTriggersUpdate` to verify:

- absent-to-present activation during Initialize triggers exactly one update;
- Execute recovers when the initialization update was skipped;
- an older absent-policy execution and the following present-policy execution
  each restore the correct recorded addresses when they interleave.

Also ran `GpuExecutableTest.CommandBufferAllocationPolicy`.

🧪 Execution Tests:

Ran the focused command-buffer regression and allocation-policy test with
Bazel. All seven configured GPU architecture variants of
`command_buffer_thunk_test` passed, and `gpu_executable_test` passed.

Copybara import of the project:

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

Update command buffers when allocation policy changes

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

Restrict persistent allocation policy transitions

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

Fix persistent allocation policy update fallback

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

Clarify persistent allocation validity names

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

Remove debug-only allocator mode checks

--
87cc9025b69cced36607307bb0301b094395f27e by Shawn Wang <[email protected]>:

Reuse counting memset test helper

--
24283bd9b477aad1835a9cd58c0889f5abfe9b7d by Shawn Wang <[email protected]>:

Revert "Reuse counting memset test helper"

This reverts commit 87cc9025b69cced36607307bb0301b094395f27e.

Merging this change closes #45130

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#45130 from shawnwang18:shawnw/trigers_update_on_update_policy_change 24283bd9b477aad1835a9cd58c0889f5abfe9b7d
PiperOrigin-RevId: 947620409
copybara-service Bot pushed a commit that referenced this pull request Jul 14, 2026
…llocations set changed.

Imported from GitHub PR openxla/xla#45130

For command buffer run, add the required processing for allowing `persistent_alloc_indices` to transition only from absent to present across execution steps and must remain unchanged once present.

We may require this transition because the chosen of buffer allocations whose address are in `persistent_alloc_indices` might not be available from the first step, e.g, if persistent allocations are only available after several steps of profiling.  So we need to make sure that command buffer is correctly updated when `persistent_alloc_indices` is changed across steps.

📝 Summary of Changes

- Document that `persistent_alloc_indices` may transition only from absent to
  present across execution steps and must remain unchanged once present.
- Track only whether persistent-allocation information was valid when each
  executor's command buffer was recorded, avoiding an owned vector copy and
  repeated value comparisons.
- Force a full command-buffer update in `CommandBufferThunk::Initialize` when
  persistent allocation information first becomes valid.
- In `ExecuteOnStream`, recover from a recorded/current information-validity
  mismatch by fully updating the command buffer. This covers a
  profiling-skipped update and interleaving with an older execution initialized
  before policy activation.
- Use the public `ShapedSlice` header and its direct BUILD dependency.

🎯 Justification

The persistent allocation policy controls which recorded addresses a command
buffer can safely retain. When the policy first becomes available, an
allocation-filtered update can skip a command whose allocations are now all
persistent, leaving stale addresses in the command buffer. A full update fixes
that transition.

Although the policy transition is logically one-way, executions initialized on
opposite sides of it can interleave on the shared per-executor command buffer.
An Execute-time validity check therefore forces a full update for either
recorded/current mismatch without comparing policy vectors on every step.

🚀 Kind of Contribution

🐛 Bug Fix

📊 Benchmark (for Performance Improvements)

N/A. This is a correctness fix.

🧪 Unit Tests:

Extended `PersistentAllocIndicesBecomingAvailableTriggersUpdate` to verify:

- absent-to-present activation during Initialize triggers exactly one update;
- Execute recovers when the initialization update was skipped;
- an older absent-policy execution and the following present-policy execution
  each restore the correct recorded addresses when they interleave.

Also ran `GpuExecutableTest.CommandBufferAllocationPolicy`.

🧪 Execution Tests:

Ran the focused command-buffer regression and allocation-policy test with
Bazel. All seven configured GPU architecture variants of
`command_buffer_thunk_test` passed, and `gpu_executable_test` passed.

Copybara import of the project:

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

Update command buffers when allocation policy changes

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

Restrict persistent allocation policy transitions

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

Fix persistent allocation policy update fallback

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

Clarify persistent allocation validity names

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

Remove debug-only allocator mode checks

--
87cc9025b69cced36607307bb0301b094395f27e by Shawn Wang <[email protected]>:

Reuse counting memset test helper

--
24283bd9b477aad1835a9cd58c0889f5abfe9b7d by Shawn Wang <[email protected]>:

Revert "Reuse counting memset test helper"

This reverts commit 87cc9025b69cced36607307bb0301b094395f27e.

Merging this change closes #45130

PiperOrigin-RevId: 947905086
copybara-service Bot pushed a commit that referenced this pull request Jul 15, 2026
…llocator class

Imported from GitHub PR openxla/xla#45567

## 📝 Summary of Changes

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

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

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

## 🎯 Justification

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

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

## 🚀 Kind of Contribution

♻️ Cleanup

## 📊 Benchmark (for Performance Improvements)

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

## 🧪 Unit Tests:

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

## 🧪 Execution Tests:

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

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

Update command buffers when allocation policy changes

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

Restrict persistent allocation policy transitions

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

Fix persistent allocation policy update fallback

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

Clarify persistent allocation validity names

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

Remove debug-only allocator mode checks

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

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

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

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

No functional change intended.

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

Fix VA remap allocator CI failures

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

Restore platform-restricted GPU PJRT deps

Merging this change closes #45567

Reverts 6c3c2e3

FUTURE_COPYBARA_INTEGRATE_REVIEW=openxla/xla#45567 from shawnwang18:shawnw/cmd_buffer_va_remap_allocator 7da1c618865c4a910928df8dc03f60afd7466ab2
PiperOrigin-RevId: 947940077
copybara-service Bot pushed a commit that referenced this pull request Jul 15, 2026
…llocator class

Imported from GitHub PR openxla/xla#45567

## 📝 Summary of Changes

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

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

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

## 🎯 Justification

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

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

## 🚀 Kind of Contribution

♻️ Cleanup

## 📊 Benchmark (for Performance Improvements)

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

## 🧪 Unit Tests:

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

## 🧪 Execution Tests:

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

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

Update command buffers when allocation policy changes

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

Restrict persistent allocation policy transitions

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

Fix persistent allocation policy update fallback

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

Clarify persistent allocation validity names

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

Remove debug-only allocator mode checks

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

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

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

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

No functional change intended.

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

Fix VA remap allocator CI failures

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

Restore platform-restricted GPU PJRT deps

Merging this change closes #45567

PiperOrigin-RevId: 948290643
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant