Relax the strict 1:1 buffer-device validation in ExecuteReplicated to support OpenXLA's 1:N PjRtMemorySpace mapping

[VNDJ]

🚀 Feature

Motivation

Currently, in torch_xla, we perform a strict 1:1 validation between buffers and devices. This is implemented in the ExecuteReplicated function within xla/torch_xla/csrc/runtime/pjrt_computation_client.cpp.

However, in OpenXLA, the PjRtMemorySpace class is designed to map to multiple PjRtDevice instances (a 1:N relationship). For example, multiple devices (clusters) on the same DSA chip can share a memory space and access each other's buffers without requiring a copy. The current 1:1 constraint in torch_xla prevents users from fully utilizing this upstream capability and causes friction when working with multi-device memory spaces.

Pitch

I @propose updating the validation logic in ExecuteReplicated to relax the strict 1:1 constraint and support the 1:N mapping between PjRtMemorySpace and PjRtDevice. Specifically, we should allow the buffer if its memory space includes the target device.

A proposed code change would look like this:

// Allow the buffer if its memory space includes the target                        
// device. Multiple PjRtDevices sharing a MemorySpace 
// can access each other's buffers without a copy.                                                         
bool same_memory_space = false;                                                    
if (auto* ms = shard->buffer->memory_space()) {                                    
  for (auto* d : ms->devices()) {                                                  
    if (d == pjrt_device) {                                                        
      same_memory_space = true;                                                        
      break;                                                                       
    }                                                                    
  }
XLA_CHECK(shard->buffer->device() == pjrt_device || same_memory_space)  

[VNDJ]

Update Changes:

XLA_CHECK(pjrt_device == pjrt_data->buffer->device() || *pjrt_device->default_memory_space() == pjrt_data->buffer->memory_space())

Check The Default MemorySpace With Buffer MemorySpace