Skip to content

Commit 89dcc46

Browse files
pschuhtensorflower-gardener
pschuh
authored and
tensorflower-gardener
committed
Switch GPU to DelinearizeAsync and remove the CopyToLiteralAsync API from CommonPjRtRawBuffer.
PiperOrigin-RevId: 927568231
1 parent b14c7a8 commit 89dcc46

8 files changed

Lines changed: 74 additions & 153 deletions

File tree

‎third_party/xla/xla/pjrt/BUILD‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -839,6 +839,7 @@ cc_library(
839839
":pjrt_stream_executor_device_description",
840840
":raw_buffer",
841841
":semaphore",
842+
":staging_buffer",
842843
":stream_executor_executable",
843844
":stream_executor_pjrt_abi_version",
844845
":thread_pool_async_work_runner",

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

Lines changed: 10 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -654,21 +654,6 @@ Future<> CommonPjRtRawBufferImpl::CopyRawDeviceToHost(void* dst, int64_t offset,
654654
"CommonPjRtRawBuffer", "CopyRawDeviceToHost");
655655
}
656656

657-
void CommonPjRtRawBufferImpl::CopyToLiteralAsync(
658-
Promise<> promise, PjRtDeviceEventPromiseRef device_promise,
659-
MutableLiteralBase* literal, xla::Shape shape) {
660-
auto* client = absl::down_cast<CommonPjRtClient*>(memory_space()->client());
661-
662-
auto staging_buffer = ToStagingBuffer(
663-
tsl::FormRef(this), device_promise,
664-
[&](size_t size, PjRtMemorySpace* memory_space) {
665-
return client->AllocateForDelinearizationAsync(size, memory_space);
666-
});
667-
668-
client->DelinearizeAsync(std::move(staging_buffer), memory_space(), shape,
669-
literal, std::move(promise));
670-
}
671-
672657
void CommonPjRtBufferImpl::CopyToRemoteDevice(
673658
Future<std::string> serialized_descriptor, RemoteSendCallback on_done) {
674659
auto* common_client = absl::down_cast<CommonPjRtClient*>(client());
@@ -2605,8 +2590,16 @@ Future<> CommonPjRtBufferImpl::ToLiteralImpl(
26052590
launcher.AddDependency(device_promise.event());
26062591
return;
26072592
}
2608-
raw_buffer->CopyToLiteralAsync(std::move(promise), device_promise,
2609-
literal, std::move(shape));
2593+
PjRtMemorySpace* memory_space = raw_buffer->memory_space();
2594+
auto staging_buffer = ToStagingBuffer(
2595+
std::move(raw_buffer), device_promise,
2596+
[common_client](size_t size, PjRtMemorySpace* memory_space) {
2597+
return common_client->AllocateForDelinearizationAsync(
2598+
size, memory_space);
2599+
});
2600+
common_client->DelinearizeAsync(std::move(staging_buffer),
2601+
memory_space, shape, literal,
2602+
std::move(promise));
26102603
};
26112604

26122605
if (literal != nullptr) {

‎third_party/xla/xla/pjrt/common_pjrt_client.h‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -753,11 +753,6 @@ class CommonPjRtRawBufferImpl : public CommonPjRtRawBuffer {
753753

754754
Future<> CopyRawDeviceToHost(void* dst, int64_t offset,
755755
int64_t transfer_size) override;
756-
757-
void CopyToLiteralAsync(Promise<> promise,
758-
PjRtDeviceEventPromiseRef device_promise,
759-
MutableLiteralBase* literal,
760-
xla::Shape shape) override;
761756
};
762757

763758
// TODO(parkers): Merge everything here into CommonPjRtBuffer.

‎third_party/xla/xla/pjrt/pjrt_stream_executor_client.cc‎

Lines changed: 56 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -107,6 +107,7 @@ limitations under the License.
107107
#include "xla/hlo/builder/xla_computation.h"
108108
#include "xla/hlo/ir/hlo_module.h"
109109
#include "xla/layout.h"
110+
#include "xla/layout_util.h"
110111
#include "xla/literal.h"
111112
#include "xla/mlir_hlo/mhlo/transforms/passes.h"
112113
#include "xla/pjrt/buffer_sequencing_event.h"
@@ -135,6 +136,7 @@ limitations under the License.
135136
#include "xla/pjrt/raw_buffer.h"
136137
#include "xla/pjrt/se_raw_buffer.h"
137138
#include "xla/pjrt/semaphore.h"
139+
#include "xla/pjrt/staging_buffer.h"
138140
#include "xla/pjrt/stream_executor_executable.h"
139141
#include "xla/pjrt/stream_executor_pjrt_abi_version.h"
140142
#include "xla/pjrt/thread_pool_async_work_runner.h"
@@ -178,6 +180,27 @@ limitations under the License.
178180

179181
namespace xla {
180182

183+
namespace {
184+
class HostStagingBuffer : public PjRtStagingBuffer {
185+
public:
186+
HostStagingBuffer(std::shared_ptr<void> buffer, size_t size)
187+
: buffer_(std::move(buffer)), size_(size) {}
188+
189+
absl::Span<uint8_t> data() override {
190+
return absl::MakeSpan(static_cast<uint8_t*>(buffer_.get()), size_);
191+
}
192+
193+
absl::Span<const uint8_t> const_data() const override {
194+
return absl::MakeConstSpan(static_cast<const uint8_t*>(buffer_.get()),
195+
size_);
196+
}
197+
198+
private:
199+
std::shared_ptr<void> buffer_;
200+
size_t size_;
201+
};
202+
} // namespace
203+
181204
template <typename T>
182205
static std::function<void()> WrapClosureAsCopyable(T cb) {
183206
return [state = std::make_shared<T>(std::move(cb))]() { return (*state)(); };
@@ -2845,4 +2868,37 @@ absl::Status PjRtStreamExecutorClient::WaitOnStream(
28452868
stream);
28462869
}
28472870

2871+
bool PjRtStreamExecutorClient::ShouldDoDirectTransfer(
2872+
const MutableLiteralBase& literal, const Shape& shape,
2873+
PjRtMemorySpace* memory_space) const {
2874+
if (shape.IsTuple()) {
2875+
return false;
2876+
}
2877+
if (primitive_util::IsSubByteNonPredType(shape.element_type())) {
2878+
return false;
2879+
}
2880+
2881+
if (literal.shape().has_layout()) {
2882+
return Layout::Equal().IgnoreMemorySpace()(shape.layout(),
2883+
literal.shape().layout());
2884+
}
2885+
2886+
return LayoutUtil::HasDescendingLayout(shape.layout());
2887+
}
2888+
2889+
tsl::AsyncValueRef<PjRtStagingBuffer>
2890+
PjRtStreamExecutorClient::AllocateForDelinearizationAsync(
2891+
size_t size, PjRtMemorySpace* memory_space) {
2892+
void* ptr = malloc(size);
2893+
if (ptr == nullptr) {
2894+
return tsl::MakeErrorAsyncValueRef(absl::ResourceExhaustedError(
2895+
absl::StrCat("Failed to allocate staging buffer of size ", size)));
2896+
}
2897+
std::shared_ptr<void> buffer(ptr, free);
2898+
2899+
auto staging_buffer = tsl::MakeAvailableAsyncValueRef<HostStagingBuffer>(
2900+
std::move(buffer), size);
2901+
return tsl::AsyncValueRef<PjRtStagingBuffer>(std::move(staging_buffer));
2902+
}
2903+
28482904
} // namespace xla

‎third_party/xla/xla/pjrt/pjrt_stream_executor_client.h‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -440,6 +440,13 @@ class PjRtStreamExecutorClient : public CommonPjRtClient {
440440
PjRtDeviceEventRef event,
441441
std::intptr_t stream) override;
442442

443+
bool ShouldDoDirectTransfer(const MutableLiteralBase& literal,
444+
const Shape& shape,
445+
PjRtMemorySpace* memory_space) const override;
446+
447+
tsl::AsyncValueRef<PjRtStagingBuffer> AllocateForDelinearizationAsync(
448+
size_t size, PjRtMemorySpace* memory_space) override;
449+
443450
absl::Status WaitForAllocation(se::Stream* stream,
444451
const CommonPjRtRawBuffer& raw_buffer);
445452

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

Lines changed: 0 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -148,13 +148,6 @@ class CommonPjRtRawBuffer : public PjRtRawBuffer {
148148
// Creates an event which signals when the allocation is complete.
149149
virtual absl::StatusOr<PjRtDeviceEventRef> MakeAllocationReadyEvent() = 0;
150150

151-
// Interprets buffer contents as having shape and linearizes these contents
152-
// async into the provided literal.
153-
virtual void CopyToLiteralAsync(Promise<> promise,
154-
PjRtDeviceEventPromiseRef device_promise,
155-
MutableLiteralBase* literal,
156-
xla::Shape shape) = 0;
157-
158151
// Copies directly into dst_raw_buffer. Must set definition_event_promise,
159152
// when dst_raw_buffer is ready, allocation_event before using dst_raw_buffer
160153
// and src_usage_event_promise when done using this buffer.

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

Lines changed: 0 additions & 119 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,6 @@ limitations under the License.
4444
#include "xla/pjrt/pjrt_stream_executor_client.h"
4545
#include "xla/pjrt/raw_buffer.h"
4646
#include "xla/pjrt/tracked_device_buffer.h"
47-
#include "xla/pjrt/transpose.h"
4847
#include "xla/primitive_util.h"
4948
#include "xla/service/generic_transfer_manager.h"
5049
#include "xla/shape_util.h"
@@ -250,124 +249,6 @@ absl::StatusOr<PjRtRawBufferRef> PjRtStreamExecutorRawBuffer::Slice(
250249
RawSEDeviceMemory::CreateSlice(device_buffer_, offset, size), size);
251250
}
252251

253-
void PjRtStreamExecutorRawBuffer::CopyToLiteralAsync(
254-
Promise<> promise, PjRtDeviceEventPromiseRef device_promise,
255-
MutableLiteralBase* literal, xla::Shape shape) {
256-
auto usage_event =
257-
BufferSequencingEvent::Create(client_->async_work_runner());
258-
usage_event.AndThen([device_buffer = device_buffer_]() {});
259-
client_->async_work_runner()->Schedule(
260-
[usage_event, local_device = local_device_,
261-
on_device_shape = std::move(shape), promise = std::move(promise),
262-
literal, client = client_, memory_space = memory_space_,
263-
device_buffer = device_buffer_]() mutable {
264-
std::shared_ptr<TransposePlan> transpose;
265-
se::Stream* stream = local_device->GetDeviceToHostStream();
266-
TransferManager* transfer_manager =
267-
client->client()->backend().transfer_manager();
268-
if (on_device_shape.IsArray()) {
269-
xla::Layout literal_layout;
270-
if (literal->shape().has_layout()) {
271-
literal_layout = literal->shape().layout();
272-
} else {
273-
literal_layout = LayoutUtil::MakeDescendingLayout(
274-
on_device_shape.dimensions().size());
275-
}
276-
277-
if (on_device_shape.layout() != literal_layout) {
278-
absl::InlinedVector<int64_t, 4> byte_strides(
279-
on_device_shape.dimensions().size());
280-
absl::Status s = ShapeUtil::UnpackedByteStrides(
281-
on_device_shape, absl::MakeSpan(byte_strides));
282-
if (!s.ok()) {
283-
promise.Set(s);
284-
client->SetEventAsError(usage_event, s);
285-
return;
286-
}
287-
absl::Span<const int64_t> dims = on_device_shape.dimensions();
288-
absl::InlinedVector<int64_t, 4> permutation(dims.size());
289-
absl::c_reverse_copy(literal_layout.minor_to_major(),
290-
permutation.begin());
291-
TransposePlan::Options options;
292-
options.elem_size_in_bytes =
293-
primitive_util::ByteWidth(on_device_shape.element_type());
294-
options.dims = on_device_shape.dimensions();
295-
options.permutation = permutation;
296-
options.input_striding = TransposePlan::Striding{byte_strides};
297-
absl::StatusOr<std::shared_ptr<TransposePlan>> t =
298-
client->GetTransposePlan(options);
299-
if (!t.ok()) {
300-
promise.Set(t.status());
301-
client->SetEventAsError(usage_event, t.status());
302-
return;
303-
}
304-
transpose = *std::move(t);
305-
}
306-
}
307-
308-
absl::StatusOr<EventPool::Handle> event_or =
309-
local_device->event_pool().AllocateEvent(
310-
client->async_work_runner(), stream->parent());
311-
if (!event_or.ok()) {
312-
promise.Set(event_or.status());
313-
client->SetEventAsError(usage_event, event_or.status());
314-
return;
315-
}
316-
317-
ShapedBuffer shaped_buffer = device_buffer->AsShapedBuffer(
318-
memory_space->devices()[0], on_device_shape);
319-
320-
GenericTransferManager::LiteralFromDeviceMetadata transfer_metadata;
321-
// We never call device functions from the `done` callback.
322-
transfer_metadata.callback_is_host_callback_safe = true;
323-
324-
TransferManager::TransferMetadata* transfer_metadata_ptr =
325-
(dynamic_cast<GenericTransferManager*>(transfer_manager) != nullptr)
326-
? &transfer_metadata
327-
: nullptr;
328-
329-
if (transpose) {
330-
// Copy the device buffer to a temporary literal with descending
331-
// layout and transpose to the requested layout.
332-
333-
Shape stage_shape = literal->shape();
334-
*stage_shape.mutable_layout() =
335-
LayoutUtil::MakeDescendingLayout(stage_shape.dimensions().size());
336-
auto staged = std::make_shared<Literal>(stage_shape);
337-
338-
transfer_manager->TransferLiteralFromDevice(
339-
stream, shaped_buffer, staged.get(),
340-
[transpose = std::move(transpose),
341-
promise = std::move(promise).ToShared(), staged, client,
342-
literal = std::move(literal)](absl::Status status) mutable {
343-
if (status.ok()) {
344-
transpose->Execute(staged->untyped_data(),
345-
literal->untyped_data());
346-
}
347-
client->async_work_runner()->Schedule(
348-
[promise = std::move(promise),
349-
status = std::move(status)]() {
350-
promise->Set(std::move(status));
351-
});
352-
},
353-
transfer_metadata_ptr);
354-
} else {
355-
transfer_manager->TransferLiteralFromDevice(
356-
stream, shaped_buffer, literal,
357-
[promise =
358-
std::move(promise).ToShared()](absl::Status status) mutable {
359-
promise->Set(std::move(status));
360-
},
361-
transfer_metadata_ptr);
362-
}
363-
364-
client->ThenRecordEvent(usage_event, local_device,
365-
std::move(event_or).value(), stream);
366-
});
367-
368-
device_promise.Set(PjRtDeviceEventRef(std::move(usage_event)));
369-
}
370-
371252
absl::StatusOr<PjRtDeviceEventRef>
372253
PjRtStreamExecutorRawBuffer::MakeAllocationReadyEvent() {
373254
auto* client =

‎third_party/xla/xla/pjrt/se_raw_buffer.h‎

Lines changed: 0 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -108,11 +108,6 @@ class PjRtStreamExecutorRawBuffer : public CommonPjRtRawBufferImpl {
108108

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

111-
void CopyToLiteralAsync(Promise<> promise,
112-
PjRtDeviceEventPromiseRef device_promise,
113-
MutableLiteralBase* literal,
114-
xla::Shape shape) override;
115-
116111
void CopyTo(PjRtRawBufferRef dst_raw_buffer,
117112
PjRtDeviceEventPromiseRef definition_event_promise,
118113
PjRtDeviceEventPromiseRef src_usage_event_promise,

0 commit comments

Comments
 (0)