Add CrossHostTransferBuffers; modify CrossHostSendBuffers to call into it

Author: rao-ashish

xla/pjrt/gpu/se_gpu_pjrt_client.cc

@@ -663,10 +663,7 @@ StreamExecutorGpuClient::CrossHostSendBuffers(
   std::vector<tsl::RCReference<PjRtRawBuffer>> raw_buffers;
   raw_buffers.reserve(buffers.size());
 
-  std::vector<std::vector<tsl::RCReference<tsl::AsyncValue>>>
-      transfer_dependency_avs;
-  transfer_dependency_avs.reserve(buffers.size());
-
+  std::vector<tsl::RCReference<tsl::AsyncValue>> transfer_dependency_avs;
   std::vector<tsl::RCReference<PjRtDeviceEventPromise>> usage_event_promises;
   usage_event_promises.reserve(buffers.size());
 
@@ -693,67 +690,145 @@ StreamExecutorGpuClient::CrossHostSendBuffers(
                 [&](tsl::RCReference<CommonPjRtRawBuffer> buf_raw_buffer,
                     std::vector<tsl::RCReference<tsl::AsyncValue>>
                         buf_definition_events) mutable
-                    -> absl::StatusOr<PjRtDeviceEventRef> {
+                -> absl::StatusOr<PjRtDeviceEventRef> {
                   // Keep raw_buffer alive until the usage_event completes,
                   // preventing the allocation from being freed while the
                   // send is in-flight.
                   usage_event.AndThen([buf_raw_buffer]() {});
                   raw_buffers.push_back(std::move(buf_raw_buffer));
-                  transfer_dependency_avs.push_back(
-                      std::move(buf_definition_events));
+                  for (tsl::RCReference<tsl::AsyncValue>& definition_event :
+                       buf_definition_events) {
+                    transfer_dependency_avs.push_back(
+                        std::move(definition_event));
+                  }
                   return PjRtDeviceEventRef(usage_event);
                 },
                 "CrossHostSendBuffers"));
   }
 
-  // Group the sends by local device.
-  absl::flat_hash_map<PjRtDevice*, std::vector<int>> sends_by_device;
+  // Build the CrossHostTransferSpec for each buffer.
+  std::vector<CrossHostTransferSpec> transfer_specs;
+  transfer_specs.reserve(buffers.size());
   for (int i = 0; i < buffers.size(); ++i) {
-    sends_by_device[buffers[i]->device()].push_back(i);
+    transfer_specs.push_back(CrossHostTransferSpec{
+        /*src_global_device_id=*/buffers[i]->device()->global_device_id(),
+        dst_global_device_ids[i], std::move(raw_buffers[i])});
   }
 
   // Schedule sends.
-  for (const auto& [device, send_idxs] : sends_by_device) {
-    const GlobalDeviceId src_global_device_id = device->global_device_id();
+  TF_ASSIGN_OR_RETURN(
+      std::vector<PjRtDeviceEventRef> usage_events,
+      CrossHostTransferBuffers(std::move(transfer_dependency_avs),
+                               std::move(transfer_specs)));
+
+  // Populate usage events.
+  for (int i = 0; i < buffers.size(); ++i) {
+    usage_event_promises[i]->Set(usage_events[i]);
+  }
+
+  return futures;
+}
+
+absl::StatusOr<std::vector<PjRtDeviceEventRef>>
+StreamExecutorGpuClient::CrossHostTransferBuffers(
+    std::vector<tsl::RCReference<tsl::AsyncValue>> transfer_dependency_avs,
+    std::vector<CrossHostTransferSpec> transfer_specs) {
+  // Validate arguments.
+  for (int i = 0; i < transfer_specs.size(); ++i) {
+    if (transfer_specs[i].raw_buffer->memory_space()->devices().size() != 1) {
+      return InvalidArgument(
+          "CrossHostTransferBuffers: Received a raw buffer with a memory space "
+          "that is not attached to exactly 1 device.");
+    }
+    PjRtDevice* buffer_device =
+        transfer_specs[i].raw_buffer->memory_space()->devices()[0];
+    if (!buffer_device->IsAddressable()) {
+      return InvalidArgument(
+          "CrossHostTransferBuffers: raw buffer %d is on non-addressable "
+          "device with global device id %d.",
+          i, buffer_device->global_device_id().value());
+    }
+    // Each transfer must be between an addressable and a non-addressable
+    // device. If both devices are addressable, then both a data transfer and a
+    // 'normal' XLA SPMD executable may try to acquire the same GPU clique,
+    // causing issues.
+    GlobalDeviceId remote_id = (transfer_specs[i].src_global_device_id ==
+                                buffer_device->global_device_id())
+                                   ? transfer_specs[i].dst_global_device_id
+                                   : transfer_specs[i].src_global_device_id;
+    TF_ASSIGN_OR_RETURN(PjRtDevice * remote_device, LookupDevice(remote_id));
+    if (remote_device->IsAddressable()) {
+      return InvalidArgument(
+          "CrossHostTransferBuffers: remote device for buffer %d is "
+          "addressable (global device id %d), but cross-host transfers must "
+          "be between an addressable and a non-addressable device.",
+          i, remote_id.value());
+    }
+  }
+
+  // Group the transfers by their buffers' device.
+  absl::flat_hash_map<PjRtDevice*, std::vector<int>> transfers_by_device;
+  for (int i = 0; i < transfer_specs.size(); ++i) {
+    PjRtDevice* buffer_device =
+        transfer_specs[i].raw_buffer->memory_space()->devices()[0];
+    transfers_by_device[buffer_device].push_back(i);
+  }
+
+  // We will register a single transfer event for all transfers to/from the same
+  // device. We will collect the references to those events inside
+  // output_transfer_events. This will eventually be returned to the user.
+  std::vector<PjRtDeviceEventRef> output_transfer_events(transfer_specs.size(),
+                                                         PjRtDeviceEventRef());
+
+  // Schedule transfers.
+  for (const auto& [device, transfer_idxs] : transfers_by_device) {
+    const GlobalDeviceId device_id = device->global_device_id();
 
     // Create a transfer event for transfers on this device.
     tsl::AsyncValueRef<BufferSequencingEvent> transfer_event =
         BufferSequencingEvent::Create(this->async_work_runner());
 
-    // Extract transfer dependencies, form transfer specs, and fulfill
-    // usage_event_promises.
-    std::vector<tsl::RCReference<tsl::AsyncValue>> curr_transfer_dependency_avs;
-    std::vector<CrossHostTransferSpec> transfer_specs;
-    transfer_specs.reserve(send_idxs.size());
+    // Form transfer specs.
+    std::vector<CrossHostTransferSpec> curr_transfer_specs;
+    curr_transfer_specs.reserve(transfer_idxs.size());
 
-    for (int idx : send_idxs) {
-      for (tsl::RCReference<tsl::AsyncValue>& event :
-           transfer_dependency_avs[idx]) {
-        curr_transfer_dependency_avs.push_back(std::move(event));
-      }
-      transfer_specs.push_back(CrossHostTransferSpec{
-          src_global_device_id, dst_global_device_ids[idx],
-          std::move(raw_buffers[idx])});
-
-      usage_event_promises[idx]->Set(PjRtDeviceEventRef(transfer_event));
+    for (int idx : transfer_idxs) {
+      curr_transfer_specs.push_back(std::move(transfer_specs[idx]));
+      output_transfer_events[idx] = PjRtDeviceEventRef(transfer_event);
     }
 
     // Get the local_device_state and use it to schedule transfers. Fail
     // transfers early if we cannot get the local_device_state.
-    absl::StatusOr<LocalDeviceState*> local_device_state =
+    absl::StatusOr<LocalDeviceState*> maybe_local_device_state =
         tensorflow::down_cast<PjRtStreamExecutorDevice*>(device)
             ->GetLocalDeviceState();
-    if (!local_device_state.ok()) {
-      SetEventAsError(transfer_event, local_device_state.status());
+    if (!maybe_local_device_state.ok()) {
+      SetEventAsError(transfer_event, maybe_local_device_state.status());
       continue;
     }
-
-    ScheduleTransfersOnLocalDevice(
-        *local_device_state, src_global_device_id, std::move(transfer_event),
-        std::move(curr_transfer_dependency_avs), std::move(transfer_specs));
+    LocalDeviceState* local_device_state = *maybe_local_device_state;
+
+    // Launch ScheduleTransfersOnLocalDevice on either the async dispatch thread
+    // of the calling thread.
+    if (local_device_state->async_dispatch_thread()) {
+      local_device_state->async_dispatch_thread()->Schedule(
+          tsl::WithCurrentContext(
+              [this, local_device_state, device_id, transfer_dependency_avs,
+               curr_transfer_specs = std::move(curr_transfer_specs),
+               transfer_event = std::move(transfer_event)]() mutable {
+                ScheduleTransfersOnLocalDevice(
+                    local_device_state, device_id, std::move(transfer_event),
+                    std::move(transfer_dependency_avs),
+                    std::move(curr_transfer_specs));
+              }));
+    } else {
+      ScheduleTransfersOnLocalDevice(
+          local_device_state, device_id, std::move(transfer_event),
+          transfer_dependency_avs, std::move(curr_transfer_specs));
+    }
   }
 
-  return futures;
+  return output_transfer_events;
 }
 
 void StreamExecutorGpuClient::ScheduleTransfersOnLocalDevice(

xla/pjrt/gpu/se_gpu_pjrt_client.h

@@ -226,6 +226,10 @@ class StreamExecutorGpuClient : public xla::PjRtStreamExecutorClient {
     tsl::RCReference<PjRtRawBuffer> raw_buffer;
   };
 
+  absl::StatusOr<std::vector<PjRtDeviceEventRef>> CrossHostTransferBuffers(
+      std::vector<tsl::RCReference<tsl::AsyncValue>> transfer_dependency_avs,
+      std::vector<CrossHostTransferSpec> transfer_specs);
+
   void ScheduleTransfersOnLocalDevice(
       LocalDeviceState* local_device_state, GlobalDeviceId device_id,
       tsl::AsyncValueRef<BufferSequencingEvent> transfer_event,