experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[XLA] Change LocalClient::TransferFromOutfeedLocal to write into a caller-provided literal rather than allocating its own literal.

Browse Source 
PiperOrigin-RevId: 353904746
Change-Id: Ia9868bee8c90c5e2a99fbb787eede16dbe75d2d1
This commit is contained in:
Peter Hawkins 2021-01-26 11:04:36 -08:00 committed by TensorFlower Gardener
parent 8e985babf8
commit f5d89fe581
12 changed files with 31 additions and 35 deletions

10
tensorflow/compiler/xla/client/local_client.cc
View File

@ -433,14 +433,12 @@ Status LocalClient::TransferToInfeedLocal(const LiteralSlice& literal,
literal);
}
StatusOr<Literal> LocalClient::TransferFromOutfeedLocal(const Shape& shape,
int device_ordinal) {
Status LocalClient::TransferFromOutfeedLocal(int device_ordinal,
MutableBorrowingLiteral literal) {
TF_ASSIGN_OR_RETURN(se::StreamExecutor * executor,
backend().stream_executor(device_ordinal));
auto literal = Literal::CreateFromShape(shape);
TF_RETURN_IF_ERROR(backend().transfer_manager()->TransferLiteralFromOutfeed(
executor, &literal));
return std::move(literal);
return backend().transfer_manager()->TransferLiteralFromOutfeed(executor,
literal);
}
StatusOr<int> LocalClient::ReplicaNumberToDeviceOrdinal(int replica_number) {

8
tensorflow/compiler/xla/client/local_client.h
View File

@ -172,13 +172,13 @@ class LocalClient : public Client {
// Client::TransferToInfeed.
Status TransferToInfeedLocal(const LiteralSlice& literal, int device_ordinal);
// Transfer and return a value of the given shape from the outfeed of the
// given device.
// Transfer and return a value from the outfeed of the given device. The
// shape of the object to transfer is determined by `literal`'s shape.
// TODO(b/69670845): Remove the 'Local' from the name when LocalClient does
// not inherit from Client and there is no possibility of confusion with
// Client::TransferFromOutfeed.
StatusOr<Literal> TransferFromOutfeedLocal(const Shape& shape,
int device_ordinal);
Status TransferFromOutfeedLocal(int device_ordinal,
MutableBorrowingLiteral literal);
// Returns the device ordinal that corresponds to the given replica number.
//

1
tensorflow/compiler/xla/pjrt/BUILD
View File

@ -132,6 +132,7 @@ cc_library(
hdrs = ["pjrt_client.h"],
visibility = ["//tensorflow/compiler/xla:friends"],
deps = [
"//tensorflow/compiler/xla:literal",
"//tensorflow/compiler/xla:shape_util",
"//tensorflow/compiler/xla:status",
"//tensorflow/compiler/xla:statusor",

3
tensorflow/compiler/xla/pjrt/pjrt_client.h
View File

@ -27,6 +27,7 @@ limitations under the License.
#include "tensorflow/compiler/xla/client/executable_build_options.h"
#include "tensorflow/compiler/xla/client/xla_computation.h"
#include "tensorflow/compiler/xla/layout.h"
#include "tensorflow/compiler/xla/literal.h"
#include "tensorflow/compiler/xla/service/hlo_cost_analysis.h"
#include "tensorflow/compiler/xla/service/hlo_module.h"
#include "tensorflow/compiler/xla/shape.h"
@ -93,7 +94,7 @@ class PjRtDevice {
virtual Status TransferToInfeed(const LiteralSlice& literal) const = 0;
// Transfer and return a value of the given shape from the outfeed queue.
virtual StatusOr<Literal> TransferFromOutfeed(const Shape& shape) const = 0;
virtual Status TransferFromOutfeed(MutableBorrowingLiteral literal) const = 0;
};
// Forward declaration.

6
tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.cc
View File

@ -910,11 +910,11 @@ Status PjRtStreamExecutorDevice::TransferToInfeed(
literal, local_device->device_ordinal());
}
StatusOr<Literal> PjRtStreamExecutorDevice::TransferFromOutfeed(
const Shape& shape) const {
Status PjRtStreamExecutorDevice::TransferFromOutfeed(
MutableBorrowingLiteral literal) const {
TF_ASSIGN_OR_RETURN(LocalDeviceState * local_device, GetLocalDeviceState());
return local_device->client()->TransferFromOutfeedLocal(
shape, local_device->device_ordinal());
local_device->device_ordinal(), literal);
}
StatusOr<PjRtDevice*> PjRtStreamExecutorClient::LookupAddressableDevice(

3
tensorflow/compiler/xla/pjrt/pjrt_stream_executor_client.h
View File

@ -32,6 +32,7 @@ limitations under the License.
#include "tensorflow/compiler/xla/client/local_client.h"
#include "tensorflow/compiler/xla/client/xla_computation.h"
#include "tensorflow/compiler/xla/layout.h"
#include "tensorflow/compiler/xla/literal.h"
#include "tensorflow/compiler/xla/pjrt/local_device_state.h"
#include "tensorflow/compiler/xla/pjrt/pjrt_client.h"
#include "tensorflow/compiler/xla/pjrt/tracked_device_buffer.h"
@ -106,7 +107,7 @@ class PjRtStreamExecutorDevice : public PjRtDevice {
Status TransferToInfeed(const LiteralSlice& literal) const override;
StatusOr<Literal> TransferFromOutfeed(const Shape& shape) const override;
Status TransferFromOutfeed(MutableBorrowingLiteral literal) const override;
private:
const int id_;

8
tensorflow/compiler/xla/python/outfeed_receiver.cc
View File

@ -341,11 +341,9 @@ void OutfeedReceiverImpl::EnqueueReceivedData(
StatusOr<std::unique_ptr<Literal>> OutfeedReceiverImpl::ReceiveRawFromOutfeed(
const PjRtDevice* device, const Shape& shape) {
std::shared_ptr<Literal> literal_shared;
TF_ASSIGN_OR_RETURN(Literal literal, device->TransferFromOutfeed(shape));
return absl::make_unique<Literal>(std::move(literal));
auto literal = std::make_unique<Literal>(shape);
TF_RETURN_IF_ERROR(device->TransferFromOutfeed(literal.get()));
return literal;
}
void OutfeedReceiverImpl::CallbackThreadLoop() {

2
tensorflow/compiler/xla/python/tpu_driver/client/tpu_client.h
View File

@ -69,7 +69,7 @@ class TpuDevice : public PjRtDevice {
return Unimplemented("Infeed not yet implemented via this API");
}
StatusOr<Literal> TransferFromOutfeed(const Shape& shape) const override {
Status TransferFromOutfeed(MutableBorrowingLiteral literal) const override {
return Unimplemented("Outfeed not yet implemented via this API");
}

10
tensorflow/compiler/xla/python/xla.cc
View File

@ -146,7 +146,7 @@ PYBIND11_MODULE(xla_extension, m) {
[](const PjRtDevice& device,
const Shape& shape) -> StatusOr<py::object> {
GlobalPyRefManager()->CollectGarbage();
std::shared_ptr<Literal> literal_shared;
std::shared_ptr<Literal> literal;
{
py::gil_scoped_release gil_release;
Shape shape_with_layout = shape;
@ -156,12 +156,10 @@ PYBIND11_MODULE(xla_extension, m) {
LayoutUtil::SetToDefaultLayout(subshape);
}
});
TF_ASSIGN_OR_RETURN(Literal literal, device.TransferFromOutfeed(
shape_with_layout));
literal_shared = std::make_shared<Literal>(std::move(literal));
literal = std::make_shared<Literal>(shape_with_layout);
TF_RETURN_IF_ERROR(device.TransferFromOutfeed(literal.get()));
}
return LiteralToPython(std::move(literal_shared));
return LiteralToPython(std::move(literal));
});
py::class_<CpuDevice, PjRtDevice, ClientAndPtr<CpuDevice>>(m, "CpuDevice")

6
tensorflow/compiler/xla/tests/local_client_execute_test.cc
View File

@ -937,9 +937,9 @@ XLA_TEST_F(LocalClientExecuteTest, DISABLED_ON_INTERPRETER(InfeedOutfeedTest)) {
LiteralUtil::CreateR1<float>({-5.0, 123.0, 42.0}),
local_client_->default_device_ordinal()));
TF_ASSERT_OK_AND_ASSIGN(Literal result,
local_client_->TransferFromOutfeedLocal(
shape, local_client_->default_device_ordinal()));
Literal result(shape);
ASSERT_IS_OK(local_client_->TransferFromOutfeedLocal(
local_client_->default_device_ordinal(), &result));
LiteralTestUtil::ExpectR1Equal<float>({-4.0, 125.0, 45.0}, result);
}

5
tensorflow/compiler/xla/tests/multiple_devices_on_host_test.cc
View File

@ -92,9 +92,8 @@ void TestWithDeviceCount(const int device_count) {
for (int device_ordinal = 0; device_ordinal < device_count;
device_ordinal++) {
TF_ASSERT_OK_AND_ASSIGN(Literal outfeed,
client->TransferFromOutfeedLocal(
ShapeUtil::MakeShape(S32, {}), device_ordinal));
Literal outfeed(ShapeUtil::MakeShape(S32, {}));
TF_ASSERT_OK(client->TransferFromOutfeedLocal(device_ordinal, &outfeed));
EXPECT_EQ(outfeed, LiteralUtil::CreateR0<int32>(device_ordinal * 100 + 1));
}

4
tensorflow/compiler/xla/tools/replay_computation.cc
View File

@ -282,9 +282,9 @@ StatusOr<Literal> ReplayComputation(const HloSnapshot& module,
outfeed_thread_pool.emplace(tensorflow::Env::Default(), "infeed",
/*num_threads=*/1);
auto consume_outfeed = [client, outfeed_shape] {
Literal outfeed(*outfeed_shape);
TF_CHECK_OK(
client->TransferFromOutfeedLocal(*outfeed_shape, /*device_ordinal=*/0)
.status());
client->TransferFromOutfeedLocal(/*device_ordinal=*/0, &outfeed));
VLOG(1) << "Received outfeed data of shape "
<< ShapeUtil::HumanStringWithLayout(*outfeed_shape);
};