Skip to content

Commit dea8836

Browse files
pschuhtensorflower-gardener
pschuh
authored and
tensorflower-gardener
committed
Add optional dependencies to CopyRawDeviceToHostAndReturnEvent and
CopyRawHostToDeviceAndReturnEvent. PiperOrigin-RevId: 934642516
1 parent 90d5c74 commit dea8836

10 files changed

Lines changed: 160 additions & 58 deletions

‎third_party/xla/xla/pjrt/c/pjrt_c_api_raw_buffer_extension.h‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,8 @@ typedef struct PJRT_RawBuffer_FunctionTable {
5050
// this method for specific alignment requirements.
5151
PJRT_Error* (*copy_raw_host_to_device_and_return_event)(
5252
PJRT_RawBuffer* raw_buffer, const void* src, int64_t offset,
53-
int64_t transfer_size, PJRT_DeviceEvent* event);
53+
int64_t transfer_size, PJRT_DeviceEventVector* dependencies,
54+
PJRT_DeviceEvent* event);
5455
// Transfers a sub-range of the on-device representation of the buffer.
5556
// offset+transfer_size must be less than get_on_device_size_in_bytes. The
5657
// returned event transitions to ready on error, or after the transfer has
@@ -61,7 +62,8 @@ typedef struct PJRT_RawBuffer_FunctionTable {
6162
// this method for specific alignment requirements.
6263
PJRT_Error* (*copy_raw_device_to_host_and_return_event)(
6364
PJRT_RawBuffer* raw_buffer, void* dst, int64_t offset,
64-
int64_t transfer_size, PJRT_DeviceEvent* event);
65+
int64_t transfer_size, PJRT_DeviceEventVector* dependencies,
66+
PJRT_DeviceEvent* event);
6567
// Return opaque device memory pointer to the underlying memory.
6668
void* (*opaque_device_memory_data_pointer)(const PJRT_RawBuffer* raw_buffer);
6769
// Fill `event` with the event that signals when the buffer allocation is

‎third_party/xla/xla/pjrt/c/pjrt_c_api_raw_buffer_external.cc‎

Lines changed: 10 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -181,18 +181,21 @@ Future<> PjRtCApiRawBuffer::CopyRawDeviceToHost(void* dst, int64_t offset,
181181
}
182182

183183
absl::StatusOr<PjRtDeviceEventRef>
184-
PjRtCApiRawBuffer::CopyRawHostToDeviceAndReturnEvent(const void* src,
185-
int64_t offset,
186-
int64_t transfer_size) {
184+
PjRtCApiRawBuffer::CopyRawHostToDeviceAndReturnEvent(
185+
const void* src, int64_t offset, int64_t transfer_size,
186+
PjRtDeviceEventRefVector dependencies) {
187187
return static_cast<PjRtRawBufferInterface*>(c_buffer_)
188-
->CopyRawHostToDeviceAndReturnEvent(src, offset, transfer_size);
188+
->CopyRawHostToDeviceAndReturnEvent(src, offset, transfer_size,
189+
std::move(dependencies));
189190
}
190191

191192
absl::StatusOr<PjRtDeviceEventRef>
192-
PjRtCApiRawBuffer::CopyRawDeviceToHostAndReturnEvent(void* dst, int64_t offset,
193-
int64_t transfer_size) {
193+
PjRtCApiRawBuffer::CopyRawDeviceToHostAndReturnEvent(
194+
void* dst, int64_t offset, int64_t transfer_size,
195+
PjRtDeviceEventRefVector dependencies) {
194196
return static_cast<PjRtRawBufferInterface*>(c_buffer_)
195-
->CopyRawDeviceToHostAndReturnEvent(dst, offset, transfer_size);
197+
->CopyRawDeviceToHostAndReturnEvent(dst, offset, transfer_size,
198+
std::move(dependencies));
196199
}
197200

198201
void* PjRtCApiRawBuffer::OpaqueDeviceMemoryDataPointer() const {

‎third_party/xla/xla/pjrt/c/pjrt_c_api_raw_buffer_external.h‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,9 +76,11 @@ class PjRtCApiRawBuffer : public PjRtRawBuffer {
7676
Future<> CopyRawDeviceToHost(void* dst, int64_t offset,
7777
int64_t transfer_size) override;
7878
absl::StatusOr<PjRtDeviceEventRef> CopyRawHostToDeviceAndReturnEvent(
79-
const void* src, int64_t offset, int64_t transfer_size) override;
79+
const void* src, int64_t offset, int64_t transfer_size,
80+
PjRtDeviceEventRefVector dependencies) override;
8081
absl::StatusOr<PjRtDeviceEventRef> CopyRawDeviceToHostAndReturnEvent(
81-
void* dst, int64_t offset, int64_t transfer_size) override;
82+
void* dst, int64_t offset, int64_t transfer_size,
83+
PjRtDeviceEventRefVector dependencies) override;
8284
void* OpaqueDeviceMemoryDataPointer() const override;
8385

8486
absl::StatusOr<PjRtRawBufferRef> Slice(int64_t offset, int64_t size) override;

‎third_party/xla/xla/pjrt/common_pjrt_client.cc‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -634,7 +634,8 @@ CommonPjRtClient::GetTransposePlan(const TransposePlan::Options& options) {
634634
Future<> CommonPjRtRawBufferImpl::CopyRawHostToDevice(const void* src,
635635
int64_t offset,
636636
int64_t transfer_size) {
637-
auto event = CopyRawHostToDeviceAndReturnEvent(src, offset, transfer_size);
637+
auto event =
638+
CopyRawHostToDeviceAndReturnEvent(src, offset, transfer_size, {});
638639
if (!event.ok()) {
639640
return Future<>(event.status());
640641
}
@@ -645,7 +646,8 @@ Future<> CommonPjRtRawBufferImpl::CopyRawHostToDevice(const void* src,
645646

646647
Future<> CommonPjRtRawBufferImpl::CopyRawDeviceToHost(void* dst, int64_t offset,
647648
int64_t transfer_size) {
648-
auto event = CopyRawDeviceToHostAndReturnEvent(dst, offset, transfer_size);
649+
auto event =
650+
CopyRawDeviceToHostAndReturnEvent(dst, offset, transfer_size, {});
649651
if (!event.ok()) {
650652
return Future<>(event.status());
651653
}

‎third_party/xla/xla/pjrt/cpu/raw_buffer.cc‎

Lines changed: 59 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -47,6 +47,7 @@ limitations under the License.
4747
#include "xla/pjrt/cpu/cpu_event.h"
4848
#include "xla/pjrt/cpu/tracked_cpu_device_buffer.h"
4949
#include "xla/pjrt/device_event.h"
50+
#include "xla/pjrt/device_event_utils.h"
5051
#include "xla/pjrt/pjrt_client.h"
5152
#include "xla/pjrt/raw_buffer.h"
5253
#include "xla/pjrt/transpose.h"
@@ -149,21 +150,69 @@ absl::StatusOr<PjRtRawBufferRef> CpuRawBuffer::Slice(int64_t offset,
149150
}
150151

151152
absl::StatusOr<PjRtDeviceEventRef>
152-
CpuRawBuffer::CopyRawHostToDeviceAndReturnEvent(const void* src, int64_t offset,
153-
int64_t transfer_size) {
153+
CpuRawBuffer::CopyRawHostToDeviceAndReturnEvent(
154+
const void* src, int64_t offset, int64_t transfer_size,
155+
PjRtDeviceEventRefVector dependencies) {
154156
RETURN_IF_ERROR(ValidateSlice(offset, transfer_size));
155-
std::memcpy(static_cast<uint8_t*>(GetHostPointer()) + offset, src,
156-
transfer_size);
157-
return PjRtDeviceEventRef(tsl::MakeAvailableAsyncValueRef<CpuEvent>());
157+
158+
if (dependencies.empty()) {
159+
std::memcpy(static_cast<uint8_t*>(GetHostPointer()) + offset, src,
160+
transfer_size);
161+
return PjRtDeviceEventRef(tsl::MakeAvailableAsyncValueRef<CpuEvent>());
162+
}
163+
164+
auto event = tsl::MakeConstructedAsyncValueRef<CpuEvent>();
165+
166+
auto run_copy = [buf = tsl::FormRef(this), src, offset, transfer_size,
167+
dependencies, event]() mutable {
168+
absl::Status dep_status = GetErrors(dependencies);
169+
if (!dep_status.ok()) {
170+
event.SetError(dep_status);
171+
return;
172+
}
173+
std::memcpy(static_cast<uint8_t*>(buf->GetHostPointer()) + offset, src,
174+
transfer_size);
175+
event.SetStateConcrete();
176+
};
177+
178+
auto* client = absl::down_cast<CommonPjRtClient*>(memory_space()->client());
179+
ExecuteWhenReady(dependencies, client->async_work_runner(),
180+
std::move(run_copy));
181+
182+
return PjRtDeviceEventRef(std::move(event));
158183
}
159184

160185
absl::StatusOr<PjRtDeviceEventRef>
161-
CpuRawBuffer::CopyRawDeviceToHostAndReturnEvent(void* dst, int64_t offset,
162-
int64_t transfer_size) {
186+
CpuRawBuffer::CopyRawDeviceToHostAndReturnEvent(
187+
void* dst, int64_t offset, int64_t transfer_size,
188+
PjRtDeviceEventRefVector dependencies) {
163189
RETURN_IF_ERROR(ValidateSlice(offset, transfer_size));
164-
std::memcpy(dst, static_cast<uint8_t*>(GetHostPointer()) + offset,
165-
transfer_size);
166-
return PjRtDeviceEventRef(tsl::MakeAvailableAsyncValueRef<CpuEvent>());
190+
191+
if (dependencies.empty()) {
192+
std::memcpy(dst, static_cast<uint8_t*>(GetHostPointer()) + offset,
193+
transfer_size);
194+
return PjRtDeviceEventRef(tsl::MakeAvailableAsyncValueRef<CpuEvent>());
195+
}
196+
197+
auto event = tsl::MakeConstructedAsyncValueRef<CpuEvent>();
198+
199+
auto run_copy = [buf = tsl::FormRef(this), dst, offset, transfer_size,
200+
dependencies, event]() mutable {
201+
absl::Status dep_status = GetErrors(dependencies);
202+
if (!dep_status.ok()) {
203+
event.SetError(dep_status);
204+
return;
205+
}
206+
std::memcpy(dst, static_cast<uint8_t*>(buf->GetHostPointer()) + offset,
207+
transfer_size);
208+
event.SetStateConcrete();
209+
};
210+
211+
auto* client = absl::down_cast<CommonPjRtClient*>(memory_space()->client());
212+
ExecuteWhenReady(dependencies, client->async_work_runner(),
213+
std::move(run_copy));
214+
215+
return PjRtDeviceEventRef(std::move(event));
167216
}
168217

169218
absl::StatusOr<PjRtDeviceEventRef> CpuRawBuffer::CopyFromLiteral(

‎third_party/xla/xla/pjrt/cpu/raw_buffer.h‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -116,10 +116,12 @@ class CpuRawBuffer : public CommonPjRtRawBufferImpl {
116116
absl::StatusOr<PjRtRawBufferRef> Slice(int64_t offset, int64_t size) override;
117117

118118
absl::StatusOr<PjRtDeviceEventRef> CopyRawHostToDeviceAndReturnEvent(
119-
const void* src, int64_t offset, int64_t transfer_size) override;
119+
const void* src, int64_t offset, int64_t transfer_size,
120+
PjRtDeviceEventRefVector dependencies) override;
120121

121122
absl::StatusOr<PjRtDeviceEventRef> CopyRawDeviceToHostAndReturnEvent(
122-
void* dst, int64_t offset, int64_t transfer_size) override;
123+
void* dst, int64_t offset, int64_t transfer_size,
124+
PjRtDeviceEventRefVector dependencies) override;
123125

124126
absl::StatusOr<PjRtDeviceEventRef> CopyFromLiteral(
125127
const LiteralSlice& literal, const xla::Layout& layout,

‎third_party/xla/xla/pjrt/raw_buffer.cc‎

Lines changed: 24 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -224,10 +224,15 @@ const PJRT_RawBuffer_FunctionTable PjRtRawBuffer::kRawBufferVtable = {
224224
},
225225
/*copy_raw_host_to_device_and_return_event=*/
226226
+[](PJRT_RawBuffer* raw_buffer, const void* src, int64_t offset,
227-
int64_t transfer_size, PJRT_DeviceEvent* event) -> PJRT_Error* {
228-
auto result =
229-
static_cast<PjRtRawBuffer*>(raw_buffer)
230-
->CopyRawHostToDeviceAndReturnEvent(src, offset, transfer_size);
227+
int64_t transfer_size, PJRT_DeviceEventVector* dependencies,
228+
PJRT_DeviceEvent* event) -> PJRT_Error* {
229+
PjRtDeviceEventRefVector cpp_deps;
230+
if (dependencies != nullptr) {
231+
cpp_deps = PjRtDeviceEventRefVector::MoveFromC(dependencies);
232+
}
233+
auto result = static_cast<PjRtRawBuffer*>(raw_buffer)
234+
->CopyRawHostToDeviceAndReturnEvent(
235+
src, offset, transfer_size, std::move(cpp_deps));
231236
if (!result.ok()) {
232237
return pjrt::StatusToPjRtError(result.status());
233238
}
@@ -236,10 +241,15 @@ const PJRT_RawBuffer_FunctionTable PjRtRawBuffer::kRawBufferVtable = {
236241
},
237242
/*copy_raw_device_to_host_and_return_event=*/
238243
+[](PJRT_RawBuffer* raw_buffer, void* dst, int64_t offset,
239-
int64_t transfer_size, PJRT_DeviceEvent* event) -> PJRT_Error* {
240-
auto result =
241-
static_cast<PjRtRawBuffer*>(raw_buffer)
242-
->CopyRawDeviceToHostAndReturnEvent(dst, offset, transfer_size);
244+
int64_t transfer_size, PJRT_DeviceEventVector* dependencies,
245+
PJRT_DeviceEvent* event) -> PJRT_Error* {
246+
PjRtDeviceEventRefVector cpp_deps;
247+
if (dependencies != nullptr) {
248+
cpp_deps = PjRtDeviceEventRefVector::MoveFromC(dependencies);
249+
}
250+
auto result = static_cast<PjRtRawBuffer*>(raw_buffer)
251+
->CopyRawDeviceToHostAndReturnEvent(
252+
dst, offset, transfer_size, std::move(cpp_deps));
243253
if (!result.ok()) {
244254
return pjrt::StatusToPjRtError(result.status());
245255
}
@@ -327,10 +337,11 @@ PjRtRawBuffer::PjRtRawBuffer() { vtable = &kRawBufferVtable; }
327337

328338
absl::StatusOr<PjRtDeviceEventRef>
329339
PjRtRawBufferInterface::CopyRawHostToDeviceAndReturnEvent(
330-
const void* src, int64_t offset, int64_t transfer_size) {
340+
const void* src, int64_t offset, int64_t transfer_size,
341+
PjRtDeviceEventRefVector dependencies) {
331342
PJRT_DeviceEvent device_event;
332343
PJRT_Error* error = vtable->copy_raw_host_to_device_and_return_event(
333-
this, src, offset, transfer_size, &device_event);
344+
this, src, offset, transfer_size, &dependencies.ToC(), &device_event);
334345
if (error != nullptr) {
335346
return pjrt::PjrtErrorToStatus(error);
336347
}
@@ -339,10 +350,11 @@ PjRtRawBufferInterface::CopyRawHostToDeviceAndReturnEvent(
339350

340351
absl::StatusOr<PjRtDeviceEventRef>
341352
PjRtRawBufferInterface::CopyRawDeviceToHostAndReturnEvent(
342-
void* dst, int64_t offset, int64_t transfer_size) {
353+
void* dst, int64_t offset, int64_t transfer_size,
354+
PjRtDeviceEventRefVector dependencies) {
343355
PJRT_DeviceEvent device_event;
344356
PJRT_Error* error = vtable->copy_raw_device_to_host_and_return_event(
345-
this, dst, offset, transfer_size, &device_event);
357+
this, dst, offset, transfer_size, &dependencies.ToC(), &device_event);
346358
if (error != nullptr) {
347359
return pjrt::PjrtErrorToStatus(error);
348360
}

‎third_party/xla/xla/pjrt/raw_buffer.h‎

Lines changed: 8 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -63,9 +63,11 @@ class PjRtRawBufferInterface : public PJRT_RawBuffer {
6363
int64_t transfer_size);
6464

6565
absl::StatusOr<PjRtDeviceEventRef> CopyRawHostToDeviceAndReturnEvent(
66-
const void* src, int64_t offset, int64_t transfer_size);
66+
const void* src, int64_t offset, int64_t transfer_size,
67+
PjRtDeviceEventRefVector dependencies = {});
6768
absl::StatusOr<PjRtDeviceEventRef> CopyRawDeviceToHostAndReturnEvent(
68-
void* dst, int64_t offset, int64_t transfer_size);
69+
void* dst, int64_t offset, int64_t transfer_size,
70+
PjRtDeviceEventRefVector dependencies = {});
6971

7072
void* OpaqueDeviceMemoryDataPointer() const;
7173

@@ -160,7 +162,8 @@ class PjRtRawBuffer : public PjRtRawBufferInterface,
160162
// on the alignment of `src` and `offset` as well. Look at implementations of
161163
// this method for specific alignment requirements.
162164
virtual absl::StatusOr<PjRtDeviceEventRef> CopyRawHostToDeviceAndReturnEvent(
163-
const void* src, int64_t offset, int64_t transfer_size) = 0;
165+
const void* src, int64_t offset, int64_t transfer_size,
166+
PjRtDeviceEventRefVector dependencies) = 0;
164167

165168
// Transfers a sub-range of the on-device representation of the buffer.
166169
// offset+transfer_size must be less than GetOnDeviceSizeInBytes. The
@@ -171,7 +174,8 @@ class PjRtRawBuffer : public PjRtRawBufferInterface,
171174
// on the alignment of `dst` and `offset` as well. Look at implementations of
172175
// this method for specific alignment requirements.
173176
virtual absl::StatusOr<PjRtDeviceEventRef> CopyRawDeviceToHostAndReturnEvent(
174-
void* dst, int64_t offset, int64_t transfer_size) = 0;
177+
void* dst, int64_t offset, int64_t transfer_size,
178+
PjRtDeviceEventRefVector dependencies) = 0;
175179

176180
// A sliced buffer is a view into the offset and range of this buffer.
177181
//

‎third_party/xla/xla/pjrt/se_raw_buffer.cc‎

Lines changed: 39 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -103,15 +103,23 @@ void PjRtStreamExecutorDeviceEventPromise::SetReady() {
103103

104104
absl::StatusOr<PjRtDeviceEventRef>
105105
PjRtStreamExecutorRawBuffer::CopyRawHostToDeviceAndReturnEvent(
106-
const void* src, int64_t offset, int64_t transfer_size) {
106+
const void* src, int64_t offset, int64_t transfer_size,
107+
PjRtDeviceEventRefVector dependencies) {
107108
se::Stream* stream = local_device_->host_to_device_stream();
108109
auto device_event =
109110
BufferSequencingEvent::Create(client_->async_work_runner());
110111
device_event.AndThen([device_buffer = device_buffer_]() {});
111-
client_->async_work_runner()->Schedule([client = client_, device_event,
112-
local_device = local_device_, stream,
113-
src, offset, transfer_size,
114-
buf = tsl::FormRef(this)]() mutable {
112+
113+
auto run_transfer = [client = client_, device_event,
114+
local_device = local_device_, stream, src, offset,
115+
transfer_size, buf = tsl::FormRef(this),
116+
dependencies = dependencies]() mutable {
117+
absl::Status dep_status = GetErrors(dependencies);
118+
if (!dep_status.ok()) {
119+
client->SetEventAsError(device_event, dep_status);
120+
return;
121+
}
122+
115123
se::DeviceAddressBase sub_buffer = buf->device_buffer_->mem();
116124
if (transfer_size < sub_buffer.size()) {
117125
sub_buffer = sub_buffer.GetByteSlice(offset, transfer_size);
@@ -155,21 +163,33 @@ PjRtStreamExecutorRawBuffer::CopyRawHostToDeviceAndReturnEvent(
155163
if (!status.ok()) {
156164
client->SetEventAsError(device_event, status);
157165
}
158-
});
166+
};
167+
168+
ExecuteWhenReady(dependencies, client_->async_work_runner(),
169+
std::move(run_transfer));
170+
159171
return PjRtDeviceEventRef(std::move(device_event));
160172
}
161173

162174
absl::StatusOr<PjRtDeviceEventRef>
163175
PjRtStreamExecutorRawBuffer::CopyRawDeviceToHostAndReturnEvent(
164-
void* dst, int64_t offset, int64_t transfer_size) {
176+
void* dst, int64_t offset, int64_t transfer_size,
177+
PjRtDeviceEventRefVector dependencies) {
165178
se::Stream* stream = local_device_->GetDeviceToHostStream();
166179
auto device_event =
167180
BufferSequencingEvent::Create(client_->async_work_runner());
168181
device_event.AndThen([device_buffer = device_buffer_]() {});
169-
client_->async_work_runner()->Schedule([client = client_, device_event,
170-
local_device = local_device_, stream,
171-
dst, offset, transfer_size,
172-
buf = tsl::FormRef(this)]() mutable {
182+
183+
auto run_transfer = [client = client_, device_event,
184+
local_device = local_device_, stream, dst, offset,
185+
transfer_size, buf = tsl::FormRef(this),
186+
dependencies = dependencies]() mutable {
187+
absl::Status dep_status = GetErrors(dependencies);
188+
if (!dep_status.ok()) {
189+
client->SetEventAsError(device_event, dep_status);
190+
return;
191+
}
192+
173193
se::DeviceAddressBase sub_buffer = buf->device_buffer_->mem();
174194
if (transfer_size < sub_buffer.size()) {
175195
sub_buffer = sub_buffer.GetByteSlice(offset, transfer_size);
@@ -214,7 +234,11 @@ PjRtStreamExecutorRawBuffer::CopyRawDeviceToHostAndReturnEvent(
214234
if (!status.ok()) {
215235
client->SetEventAsError(device_event, status);
216236
}
217-
});
237+
};
238+
239+
ExecuteWhenReady(dependencies, client_->async_work_runner(),
240+
std::move(run_transfer));
241+
218242
return PjRtDeviceEventRef(std::move(device_event));
219243
}
220244

@@ -304,8 +328,8 @@ void PjRtStreamExecutorRawBuffer::CopyTo(
304328
definition_event_promise.Set(*std::move(h2d_event));
305329
return;
306330
} else if (auto* dst_ptr = dst_raw_buffer->GetHostPointer()) {
307-
auto d2h_event =
308-
CopyRawDeviceToHostAndReturnEvent(dst_ptr, 0, GetOnDeviceSizeInBytes());
331+
auto d2h_event = CopyRawDeviceToHostAndReturnEvent(
332+
dst_ptr, 0, GetOnDeviceSizeInBytes(), {});
309333
if (!d2h_event.ok()) {
310334
definition_event_promise.SetError(d2h_event.status());
311335
src_usage_event_promise.SetError(d2h_event.status());
@@ -332,7 +356,7 @@ void PjRtStreamExecutorRawBuffer::CopyTo(
332356
client_->GetHostMemoryAllocator()->Allocate(GetOnDeviceSizeInBytes(),
333357
alloc_opts);
334358
auto d2h_event = CopyRawDeviceToHostAndReturnEvent(
335-
staging_buffer.get(), 0, GetOnDeviceSizeInBytes());
359+
staging_buffer.get(), 0, GetOnDeviceSizeInBytes(), {});
336360
if (!d2h_event.ok()) {
337361
definition_event_promise.SetError(d2h_event.status());
338362
src_usage_event_promise.SetError(d2h_event.status());

0 commit comments

Comments
 (0)