[ROCm] Fix cross-stream race between H2D transfers and compute kernels

[mminutoli]

Summary

Hunting a flaky CI failure in tests/batching_test.py::BatchingTest::testConvGeneralDilated on the ROCm 4-GPU MI350X runner, where per_example_direct (the Python loop reference) was all zeros while the vmap result was correct.

What Could Happen

In the op-by-op execution path (StreamExecutorGpuClient::RunAsync, used under pytest-xdist with single-GPU worker isolation), input buffers produced by H2D transfers have their definition event tied to the compute-stream sync_point counter — not to host_to_device_stream_. When the compute stream advances past that sync_point and evicts it from the sliding window, GetDefinitionEvent(nullptr_if_past=true) returns nullptr, and RunAsync skips waiting before dispatching the kernel. The kernel then reads the input before the H2D transfer completes, getting zeros.

The invariant broken is subtle: AllocatedRawSEDeviceMemory always used the compute_stream counter to represent "buffer is ready", but H2D-produced buffers are written by host_to_device_stream_, an independent HIP stream. "Compute stream past sync_point" does not imply "H2D transfer complete."

Fix

1. xla/pjrt/tracked_device_buffer.cc/.h — Add SetProducerEvent() / producer_event_ to AllocatedRawSEDeviceMemory. GetDefinitionEvent returns this event (on the real producer stream) when set, falling back to the existing compute-stream sync_point mechanism for compute outputs.

2. xla/pjrt/se_raw_buffer.cc — In CopyRawHostToDeviceAndReturnEvent, pre-create the H2D completion event and call SetProducerEvent synchronously before scheduling the async task. Any consumer calling GetDefinitionEvent after this point sees the host_to_device_stream_ event, not a compute-stream sync_point.

3. xla/pjrt/gpu/se_gpu_pjrt_client.cc — In RunAsync, call WaitForAllocation for each input flat_argument before ExecuteThunks. For H2D-produced inputs, this enqueues hipStreamWaitEvent(compute_stream, h2d_event), preventing the kernel from starting before the transfer completes. For already-complete or compute-stream buffers, the call is a no-op.

Design notes

• No race on producer_event_: the event handle is pre-created and stored in the buffer before the async H2D task is scheduled, and before the buffer is returned to the caller. The mutex in GetDefinitionEvent guards concurrent reads from RunAsync against the (rare) case where SetProducerEvent is called concurrently.
• Compute-stream buffers unaffected: sync_point_ fallback is unchanged; no performance regression on the common path.
• D2D transfers (ScheduleCopyTo, IntraClientCopyToWithDependencies) have the same cross-stream pattern but are not implicated in the known failure — left as follow-on work.

[claude]

Bug (medium severity): SlicedRawSEDeviceMemory inherits this no-op, so calling SetProducerEvent on a slice silently drops the event. More importantly, SlicedRawSEDeviceMemory does not override GetDefinitionEvent either — so even if the base buffer has a producer event set, a slice of it will not see it (GetDefinitionEvent on the slice falls through to the uninitialized sync_point_ path and returns an empty event).

If a buffer produced by CopyRawHostToDeviceAndReturnEvent is later sliced and then used as an argument in RunAsync, WaitForAllocation on the slice will not wait for the H2D transfer, reintroducing the race this PR aims to fix.

SlicedRawSEDeviceMemory should delegate both SetProducerEvent and GetDefinitionEvent to its base_ object.

[claude]

Nit: SetProducerEvent unconditionally overwrites producer_event_. The comment says "Must be called before the buffer is visible to consumers," but nothing enforces single-call semantics. If a second call ever happens, the earlier event is silently dropped — any consumer that obtained the first event reference would still wait on the right thing, but the buffer would lose its ordering guarantee with respect to the first transfer.

Consider adding a DCHECK to catch violations during development:

Suggested change

	void SetProducerEvent(BufferSequencingEventRef event) override {
	absl::MutexLock l(&producer_event_mu_);
	producer_event_ = std::move(event);
	}
	void SetProducerEvent(BufferSequencingEventRef event) override {
	absl::MutexLock l(&producer_event_mu_);
	DCHECK(!producer_event_) << "SetProducerEvent called more than once";
	producer_event_ = std::move(event);
	}

[claude]

Potential issue: IsComplete() documentation says it "blocks the calling thread until the event has been recorded" if SetSequencingEvent has not yet been called. The device_event is created via BufferSequencingEvent::Create() and the sequencing event is only set later inside the async task via AllocateAndRecordEvent.

There is a window between SetProducerEvent (before the async task is scheduled) and when AllocateAndRecordEvent records the GPU event, during which IsComplete() could block the thread calling GetDefinitionEvent. If the async work runner is starved or queued behind other work, this blocks the RunAsync caller. With nullptr_if_past = true, the intent is a non-blocking optimization — the blocking behavior here contradicts that.

Consider guarding this with a non-blocking check (e.g., checking whether the underlying event has been allocated before calling IsComplete()), or document that this path may block.

[claude]

Note (reference cycle): device_event.AndThen(...) on line 109 captures a copy of device_buffer_ (an AsyncValueRef<RawSEDeviceMemory>). Then SetProducerEvent on line 115 stores device_event into device_buffer_'s producer_event_ field. This creates a reference cycle: device_buffer_ → producer_event_ → (via AndThen callback) → device_buffer_.

The cycle should be broken when the event completes and the AndThen callback runs and releases its captured reference. But if the event never completes (e.g., GPU error causes the stream to fail without firing the callback), this could leak.

Verify that SetEventAsError (line 161) reliably triggers AndThen callbacks on the BufferSequencingEvent to break the cycle in all error paths.

[claude]

Claude Review Summary

This PR fixes a cross-stream race on ROCm where H2D transfers on host_to_device_stream could complete after a compute kernel starts reading the buffer. The approach — binding a producer_event_ to AllocatedRawSEDeviceMemory and inserting WaitForAllocation in RunAsync — is sound and well-targeted.

Key finding: SlicedRawSEDeviceMemory inherits the base-class no-op SetProducerEvent and does not override GetDefinitionEvent, so sliced buffers lose the producer event and would not be waited on in RunAsync. This could reintroduce the race for sliced H2D buffers.

Other findings (see inline comments):

• SetProducerEvent silently allows double-calls — a DCHECK would catch violations
• IsComplete() can block when called before the event is recorded, contradicting the nullptr_if_past non-blocking intent
• Potential reference cycle between device_buffer_ and producer_event_ via AndThen callback — safe if error paths trigger AndThen
• No tests included for this synchronization fix

[mfrancepillois]

Could you please ensure that this PR only contains actual changes? It appears that many lines have been changed unnecessarily, which make the PR review and code tracking more complex.

[mfrancepillois]

I don't think IsComplete() is the correct method to use here, as it is a blocking method that wait until the event is completed before returning, which is not the goal of nullptr_if_past which aims to speed-up the execution. You should instead consider implementing a function like the following to perform a non blocking check:

// Non-blocking: returns nullopt if not yet definable
std::optional<bool> TryIsComplete(BufferSequencingEvent& ev) {
  if (!ev.IsDefined()) {
    return std::nullopt;  // SetSequencingEvent not called yet — state unknown
  }
  // event_ is now safely readable without blocking
  if (ev.IsPredeterminedError()) {
    return true;
  }
  return ev.event()->event.event()->PollForStatus() == se::Event::Status::kComplete;
}

[magaonka-amd]

@mfrancepillois I dont think this PR is fully valid now. we have found the main issue thats causing this race

I have opened up draft solution ( which is still not final )
#980