raw_transfer: compile across the XLA PjRtRawBufferInterface->CommonPjRtRawBuffer rename

[lokic233]

Problem

tpu_sync/core/raw_transfer_core.h names xla::PjRtRawBufferInterface directly (the common_raw_buffer field, the tsl::FormRef<...> in Acquire, and the AcquireFromRaw parameter). XLA has since renamed that class to xla::CommonPjRtRawBuffer (see xla/pjrt/raw_buffer.h, where class CommonPjRtRawBuffer : public PjRtRawBuffer and using PjRtRawBufferRef = tsl::RCReference<CommonPjRtRawBuffer>).

Because the concrete class name is hard-coded, this header fails to compile against XLA revisions on the post-rename side:

error: 'PjRtRawBufferInterface' is not a member of 'xla'; did you mean 'PjRtRawBufferRef'?
error: 'struct raiden::RaidenBufferHandle' has no member named 'common_raw_buffer'

(The common_raw_buffer / memory_space cascade errors are all downstream of the single type-name failure.)

This is not hypothetical: when a prebuilt _raw_transfer.so (built against a pre-rename XLA) is loaded against a runtime whose XLA is on the other side of the rename, the mismatch surfaces at runtime as RuntimeError: Not a PjRt compatible array from jax_utils.h (CastToPjRtCompatibleArray returns null across the ABI), and every transfer silently moves 0 bytes because the exception is swallowed inside producer threads. Building the .so version-matched to the target runtime is the real fix — but the source only compiles against one side of the rename at a time.

Fix (version-agnostic, no preprocessor gate)

The xla::PjRtRawBufferRef alias (tsl::RCReference<...>) is stable across the rename — only the pointee class name changed. Derive the concrete type from it via decltype and use that everywhere:

using RaidenRawBuffer =
    std::remove_reference_t<decltype(*std::declval<xla::PjRtRawBufferRef>())>;

RaidenRawBuffer resolves to PjRtRawBufferInterface on pre-rename XLA and CommonPjRtRawBuffer on post-rename XLA, so raw_transfer_core.h builds against both without a version macro. The 4 usages are switched to RaidenRawBuffer; <type_traits> is added for std::remove_reference_t (<utility> for std::declval is already included).

Verified: the decltype idiom compiles under -std=c++17 and a static_assert confirms the alias resolves to the correct concrete type on each side of the rename.

1 file changed, +17/-4.

— via Navi on behalf of Loki Chen

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[juncgu-google]

Thanks, @lokic233.
Could you share the specific XLA version where you observed the renaming?