OpenXLA PR #49428: Why an Unrolled Loop Updated the Same GPU Command Twice

OpenXLA PR #49428 fixes a command-buffer update bug that appeared when one GPU custom-call thunk was recorded more than once into the same unrolled loop.

Each loop iteration created a distinct set of command-buffer nodes with its own operand and result addresses.

The custom call’s record state kept only one sequence of FFI command handles.

Every new RecordCreate replaced the previous sequence. When XLA later updated the command buffer with new buffer allocations, every loop iteration was given the command sequence from the last recorded iteration.

The final iteration was updated repeatedly. Earlier iterations retained stale buffer pointers.

The patch replaces the single stored sequence with a map:

recorded sink command
→ full FFI command sequence for that recording

An update can now recover the commands belonging to the exact loop iteration being updated.

Command-buffer creation and update have different jobs

XLA’s GPU command-buffer path has two relevant phases.

During creation, a command records new nodes and establishes dependencies:

RecordCreate
→ create command-buffer nodes
→ return a command handle representing the recorded work

During update, the graph structure is already fixed:

RecordUpdate
→ receive the previously returned command handle
→ update kernel arguments or buffer pointers
→ keep the dependency graph unchanged

The returned command handle is therefore more than a completion token. It identifies the concrete recording that a later update must modify.

This distinction is especially important for FFI custom calls. One record handler may create one command or a small sequence of commands. XLA must retain those opaque FFI command handles so the same handler can update their arguments on a later execution.

Loop unrolling records one source thunk several times

A while loop with a known trip count can be represented as an unrolled command sequence:

condition iteration 0
body iteration 0
condition iteration 1
body iteration 1
...

OpenXLA’s WhileThunk records the condition and body executors once for each iteration.

Each iteration receives a separate RecordId, but the same underlying command objects are reused.

A custom call inside the body can therefore follow this pattern:

same CustomCallThunk object
        ↓
RecordCreate for iteration 0
        ↓
RecordCreate for iteration 1
        ↓
RecordCreate for iteration 2

The command-buffer nodes are different even though the source thunk object is the same.

That is a valid use of the command API. The state attached to the thunk must preserve each recorded instance separately.

The external state key was intentionally broader than an iteration

CommandStateManager owns mutable state for a command-buffer lifetime.

Its key includes:

source Command object
command buffer
state type

It does not include the unrolled loop’s RecordId.

The source documentation explicitly allows a command to be recorded multiple times for different record IDs.

That means a state object can be shared by several recordings of the same command in one command buffer. Any state whose cardinality is per recording must add another identity level inside that state.

CustomCallRecordState did not.

Before the patch, it contained:

std::vector<const XLA_FFI_Command*> commands;

There was only one slot for the command sequence, regardless of how many times the custom call had been recorded.

The second create erased the first recording’s identity

Consider two unrolled iterations whose FFI handlers each create one GPU launch.

The command-buffer creation phase produces:

iteration 0:
outer sink = cmd_0
FFI sequence = [ffi_0]

iteration 1:
outer sink = cmd_1
FFI sequence = [ffi_1]

The old state evolved like this:

after iteration 0 create:
state.commands = [ffi_0]

after iteration 1 create:
state.commands = [ffi_1]

The first sequence was no longer reachable through the state.

The command-buffer graph still contained both launches. What disappeared was the association between the first outer command handle and the first FFI command handle.

Both updates targeted the last iteration

On the next execution, XLA can receive new buffer allocations for both loop iterations.

The intended updates are:

RecordUpdate(cmd_0)
→ update ffi_0 with iteration 0 buffers

RecordUpdate(cmd_1)
→ update ffi_1 with iteration 1 buffers

The old implementation ignored the command passed in RecordUpdate.

It copied the one stored sequence into the FFI record frame for every update.

After two creates, that sequence belonged to iteration 1:

RecordUpdate(cmd_0)
→ supplies [ffi_1]

RecordUpdate(cmd_1)
→ supplies [ffi_1] again

The second launch received both updates in sequence and ended with the second iteration’s newest pointers.

The first launch was never rebound. It continued to reference the buffers supplied when it was created.

This is not a graph-dependency error. The correct nodes exist in the command buffer. The update layer has lost the mapping from a recorded node to the opaque commands it owns.

The regression makes the stale update visible

The new test records the same CustomCallThunk twice into one command buffer using one CommandStateManager.

The FFI handler launches a kernel that adds two 32-bit integers.

The two creates return different sink commands:

cmd_0 != cmd_1

The test then updates the command buffer with two new sets of allocations.

Iteration 0 receives:

40 + 50
→ expected output 90

Iteration 1 receives:

100 + 200
→ expected output 300

After replay, the test copies both updated result buffers back to the host and requires:

output 0 = 90
output 1 = 300

Under the old single-vector state, the update for cmd_0 would also target the FFI command belonging to cmd_1. The first updated result buffer would not receive the required 90.

The test directly protects the missing association rather than relying on a larger model to expose stale pointers indirectly.

The sink command is the right per-recording key

An FFI record handler can create more than one command.

CustomCallThunk::Record() returns the last command in that sequence as the sink used for dependency tracking.

If several recorded commands are independent, the FFI client is responsible for adding a dummy node so the complete sequence still has one sink.

That sink is then passed back in:

RecordUpdate{command}

The final patch uses this existing round trip as the identity key.

The state becomes:

flat_hash_map<
    sink command pointer,
    full FFI command sequence
>

During creation:

FFI handler records its commands
        ↓
CustomCallThunk obtains the sink
        ↓
state.commands[sink] = recorded sequence
        ↓
sink is returned to the command executor

During update:

RecordUpdate supplies the sink
        ↓
state.commands.find(sink)
        ↓
the matching FFI sequence is copied into the update frame
        ↓
the handler updates the correct launches and memory operations

The identity used by the outer command-buffer API and the identity used by the FFI update state are now the same.

The full sequence still matters

Keying only one FFI command would be insufficient.

A custom call can record a chain such as:

device-to-device copy
        ↓
kernel launch
        ↓
barrier or sink node

On update, the handler may need the original handles for every operation whose parameters can change.

The map value is therefore an InlinedVector of FFI commands, not just the sink.

The inline capacity is one, which matches the common single-command case used by the regression while retaining support for multi-command handlers.

No performance claim is attached to that container choice in the public change.

The map’s lifetime remains tied to the command buffer

The patch does not create a process-wide command cache.

CustomCallRecordState is still owned through CommandStateManager, whose lifetime is tied to the command buffer.

The map grows only with distinct recordings of that custom call inside that command buffer and disappears when the command-buffer state is destroyed.

This preserves the runtime’s existing ownership model:

source command object:
reusable and stateless

command-buffer state:
owns the recorded-instance mappings

recorded sink:
identifies one concrete command sequence

The fix changes the identity inside the state, not the lifetime of the state itself.

The test is focused rather than end to end

The regression reproduces the important structure directly:

same thunk
same command buffer
same state manager
two RecordCreate calls
two distinct RecordUpdate calls
different buffer allocations

It does not construct an HLO while loop, run WhileThunk, or prove that every unrolled control-flow configuration is covered.

That focus makes the state bug deterministic and keeps unrelated loop lowering out of the test.

The test skips ROCm because its AddI32 helper uses PTX. The production fix is in shared GPU runtime code, but this committed regression does not directly execute the ROCm path.

What the public evidence does not show

The pull request does not identify a public model, user workload, or wrong numerical output that originally exposed the defect.

It publishes no benchmark.

It also contains no public review discussion because the GitHub change was imported without visible review submissions or comments.

This review did not build OpenXLA or run the GPU test.

The source and regression establish the stale-pointer mechanism and the corrected identity mapping. They do not quantify how often unrolled FFI custom calls appear in production programs.

The state contract restored by the patch

A source command and a recorded command are not the same identity.

The corrected lifecycle is:

record the same source command for iteration 0
→ retain iteration 0’s sink-to-sequence mapping

record it again for iteration 1
→ retain a separate mapping

update iteration 0
→ recover iteration 0’s FFI commands

update iteration 1
→ recover iteration 1’s FFI commands