Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

IFRT Proxy: Support fetching an array with a sub-byte dtype #23050

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

IFRT Proxy: Support fetching an array with a sub-byte dtype #23050

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 13 additions & 6 deletions xla/python/ifrt_proxy/common/array_util.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@

#include "xla/python/ifrt_proxy/common/array_util.h"

#include <cstddef>
#include <memory>
#include <string>
#include <vector>
Expand Down Expand Up @@ -63,7 +64,13 @@ absl::StatusOr<std::vector<int64_t>> DefaultByteStrides(const DType dtype,
absl::StatusOr<ArrayMemRegion> ArrayMemRegion::FromZerothElementPointer(
const void* zeroth_element, const DType dtype, const Shape& shape,
ByteStrides byte_strides) {
if (!dtype.byte_size().has_value()) {
int byte_size;
if (dtype.byte_size().has_value()) {
byte_size = *dtype.byte_size();
} else if (dtype.bit_size().has_value() && *dtype.bit_size() < 8) {
// IFRT uses 1 byte per element for S4 and S2.
byte_size = 1;
} else {
return absl::InvalidArgumentError(
absl::StrCat("Unsupported data type to construct ArrayMemRegion: ",
dtype.DebugString()));
Expand All @@ -74,8 +81,7 @@ absl::StatusOr<ArrayMemRegion> ArrayMemRegion::FromZerothElementPointer(

if (!byte_strides.has_value() ||
(byte_strides->empty() && shape.dims().empty())) {
return ArrayMemRegion(mem_region_start,
dtype.byte_size().value() * shape.num_elements());
return ArrayMemRegion(mem_region_start, byte_size * shape.num_elements());
}
if (shape.num_elements() == 0) {
return ArrayMemRegion(mem_region_start, 0);
Expand Down Expand Up @@ -109,7 +115,7 @@ absl::StatusOr<ArrayMemRegion> ArrayMemRegion::FromZerothElementPointer(
return absl::UnimplementedError(
absl::StrCat("Negative or zero strides are not fully supported: ",
StridesAsStr(byte_strides)));
} else if (stride % dtype.byte_size().value() != 0) {
} else if (stride % byte_size != 0) {
return absl::UnimplementedError(absl::StrCat(
"byte_stride[", i, "] is not a multiple of the data-type's size: ",
StridesAsStr(byte_strides), ", dtype=", dtype.DebugString()));
Expand All @@ -120,8 +126,7 @@ absl::StatusOr<ArrayMemRegion> ArrayMemRegion::FromZerothElementPointer(
last_element_byte_offset += (stride * (shape.dims()[i] - 1));
}
}
return ArrayMemRegion(mem_region_start,
last_element_byte_offset + dtype.byte_size().value());
return ArrayMemRegion(mem_region_start, last_element_byte_offset + byte_size);
}

absl::StatusOr<ArrayMemRegion> ArrayMemRegion::FromMinimalMemRegion(
Expand Down Expand Up @@ -155,6 +160,8 @@ void* ArrayMemRegion::zeroth_element() const {
return mem_region_start_;
}

size_t ArrayMemRegion::nbytes() const { return nbytes_; }

absl::StatusOr<std::unique_ptr<std::string>> SerializeStringHostBuffer(
absl::Span<const absl::Cord> cords) {
proto::StringArrayContents string_array_proto;
Expand Down
4 changes: 4 additions & 0 deletions xla/python/ifrt_proxy/common/array_util.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -70,6 +70,10 @@ class ArrayMemRegion {
// the) Array.
void* zeroth_element() const;

// Returns the number of bytes necessary for the (in-host representation of
// the) Array.
size_t nbytes() const;

private:
ArrayMemRegion(void* mem_region_start, size_t nbytes)
: mem_region_start_(mem_region_start), nbytes_(nbytes) {}
Expand Down
1 change: 0 additions & 1 deletion xla/python/ifrt_proxy/server/BUILD
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -125,7 +125,6 @@ cc_library(
":host_callback",
":version",
"//xla:shape_util",
"//xla:status_macros",
"//xla:xla_data_proto_cc",
"//xla/pjrt:pjrt_layout",
"//xla/python/ifrt",
Expand Down
29 changes: 14 additions & 15 deletions xla/python/ifrt_proxy/server/ifrt_backend.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,7 +74,6 @@
#include "xla/python/ifrt_proxy/server/host_callback.h"
#include "xla/python/ifrt_proxy/server/version.h"
#include "xla/python/pjrt_ifrt/xla_compiler.h"
#include "xla/status_macros.h"
#include "xla/tsl/concurrency/ref_count.h"
#include "xla/tsl/platform/env.h"
#include "xla/tsl/platform/errors.h"
Expand Down Expand Up @@ -879,26 +878,26 @@ Future<BackendInterface::Response> IfrtBackend::HandleCopyToHostBufferRequest(
return HandleCopyToStringHostBufferRequest(std::move(request));
}

// Determine the size and allocate the host buffer.
// TODO(b/282757875): We may need to redo this to account for byte_strides,
// padding, and alignment requirements.
std::optional<int> element_size = (*array)->dtype().byte_size();
if (element_size == std::nullopt) {
return Future<Response>(
absl::InternalError("Array element size is unknown."));
}
int64_t host_buffer_size =
(*array)->shape().num_elements() * element_size.value();
// Use `std::unique_ptr<std::string>` for pointer stability.
auto host_buffer = std::make_unique<std::string>();
host_buffer->resize(host_buffer_size);

const auto byte_strides = [&]() -> std::optional<std::vector<int64_t>> {
if (!copy_to_host.has_byte_strides()) {
return std::nullopt;
}
return FromByteStridesProto(copy_to_host.byte_strides());
}();

// Use `ArrayMemRegion`'s factory functions to determine the size necessary
// for the host buffer.
const auto pseudo_mem_region = ArrayMemRegion::FromZerothElementPointer(
/*zeroth_element=*/nullptr, (*array)->dtype(), (*array)->shape(),
byte_strides);
if (!pseudo_mem_region.ok()) {
return Future<Response>(pseudo_mem_region.status());
}

// Use `std::unique_ptr<std::string>` for pointer stability.
auto host_buffer = std::make_unique<std::string>();
host_buffer->resize(pseudo_mem_region->nbytes());

const auto mem_region = ArrayMemRegion::FromMinimalMemRegion(
absl::string_view(*host_buffer), (*array)->dtype(), (*array)->shape(),
byte_strides);
Expand Down