public:

explicit TensorToPlasmaOp(tf::OpKernelConstruction* context) : tf::AsyncOpKernel(context) {

OP_REQUIRES_OK(context, context->GetAttr("plasma_store_socket_name",

&plasma_store_socket_name_));

tf::mutex_lock lock(mu_);

if (!connected_) {

VLOG(1) << "Connecting to Plasma...";

ARROW_CHECK_OK(client_.Connect(plasma_store_socket_name_));

VLOG(1) << "Connected!";

connected_ = true;

}

}

~TensorToPlasmaOp() override {

{

tf::mutex_lock lock(mu_);

ARROW_CHECK_OK(client_.Disconnect());

connected_ = false;

}

{

tf::mutex_lock lock(d2h_stream_mu);

if (d2h_stream != nullptr) {

delete d2h_stream;

}

}

}

void ComputeAsync(tf::OpKernelContext* context, DoneCallback done) override {

const int num_inputs = context->num_inputs();

OP_REQUIRES_ASYNC(

context, num_inputs >= 2,

tf::errors::InvalidArgument("Input should have at least 1 tensor and 1 object_id"),

done);

const int num_tensors = num_inputs - 1;

// Check that all tensors have the same dtype

tf::DataType tf_dtype = context->input(0).dtype();

for (int i = 1; i < num_inputs - 1; i++) {

if (tf_dtype != context->input(i).dtype()) {

ARROW_CHECK_OK(arrow::Status(arrow::StatusCode::TypeError,

"All input tensors must have the same data type"));

}

}

std::shared_ptr<arrow::DataType> arrow_dtype;

ARROW_CHECK_OK(arrow::adapters::tensorflow::GetArrowType(tf_dtype, &arrow_dtype));

int64_t byte_width = get_byte_width(*arrow_dtype);

std::vector<size_t> offsets;

offsets.reserve(num_tensors + 1);

offsets.push_back(0);

int64_t total_bytes = 0;

for (int i = 0; i < num_tensors; ++i) {

const size_t s = context->input(i).TotalBytes();

CHECK_EQ(s, context->input(i).NumElements() * byte_width);

CHECK_GT(s, 0);

total_bytes += s;

offsets.push_back(total_bytes);

}

const tf::Tensor& plasma_object_id = context->input(num_inputs - 1);

CHECK_EQ(plasma_object_id.NumElements(), 1);

const std::string& plasma_object_id_str = plasma_object_id.flat<std::string>()(0);

VLOG(1) << "plasma_object_id_str: '" << plasma_object_id_str << "'";

const plasma::ObjectID object_id =

plasma::ObjectID::from_binary(plasma_object_id_str);

std::vector<int64_t> shape = {total_bytes / byte_width};

arrow::io::MockOutputStream mock;

ARROW_CHECK_OK(arrow::py::WriteNdarrayHeader(arrow_dtype, shape, 0, &mock));

int64_t header_size = mock.GetExtentBytesWritten();

std::shared_ptr<Buffer> data_buffer;

{

tf::mutex_lock lock(mu_);

ARROW_CHECK_OK(client_.Create(object_id, header_size + total_bytes,

/*metadata=*/nullptr, 0, &data_buffer));

}

int64_t offset;

arrow::io::FixedSizeBufferWriter buf(data_buffer);

ARROW_CHECK_OK(arrow::py::WriteNdarrayHeader(arrow_dtype, shape, total_bytes, &buf));

ARROW_CHECK_OK(buf.Tell(&offset));

uint8_t* data = reinterpret_cast<uint8_t*>(data_buffer->mutable_data() + offset);

auto wrapped_callback = [this, context, done, data_buffer, data, object_id]() {

{

tf::mutex_lock lock(mu_);

ARROW_CHECK_OK(client_.Seal(object_id));

ARROW_CHECK_OK(client_.Release(object_id));

#ifdef GOOGLE_CUDA

auto orig_stream = context->op_device_context()->stream();

auto stream_executor = orig_stream->parent();

CHECK(stream_executor->HostMemoryUnregister(static_cast<void*>(data)));

#endif

}

context->SetStatus(tensorflow::Status::OK());

done();

};

if (std::is_same<Device, CPUDevice>::value) {

for (int i = 0; i < num_tensors; ++i) {

const auto& input_tensor = context->input(i);

std::memcpy(static_cast<void*>(data + offsets[i]),

input_tensor.tensor_data().data(),

static_cast<tf::uint64>(offsets[i + 1] - offsets[i]));

}

wrapped_callback();

} else {

#ifdef GOOGLE_CUDA

auto orig_stream = context->op_device_context()->stream();

OP_REQUIRES_ASYNC(context, orig_stream != nullptr,

tf::errors::Internal("No GPU stream available."), done);

auto stream_executor = orig_stream->parent();

// NOTE(zongheng): this is critical of getting good performance out of D2H

// async memcpy. Under the hood it performs cuMemHostRegister(), see:

// http://docs.nvidia.com/cuda/cuda-driver-api/group__CUDA__MEM.html#group__CUDA__MEM_1gf0a9fe11544326dabd743b7aa6b54223

CHECK(stream_executor->HostMemoryRegister(static_cast<void*>(data),

static_cast<tf::uint64>(total_bytes)));

{

tf::mutex_lock l(d2h_stream_mu);

if (d2h_stream == nullptr) {

d2h_stream = new Stream(stream_executor);

CHECK(d2h_stream->Init().ok());

}

}

// Needed to make sure the input buffers have been computed.

// NOTE(ekl): this is unnecessary when the op is behind a NCCL allreduce already

CHECK(d2h_stream->ThenWaitFor(orig_stream).ok());

for (int i = 0; i < num_tensors; ++i) {

const auto& input_tensor = context->input(i);

auto input_buffer = const_cast<char*>(input_tensor.tensor_data().data());

perftools::gputools::DeviceMemoryBase wrapped_src(

static_cast<void*>(input_buffer));

const bool success =

d2h_stream

->ThenMemcpy(static_cast<void*>(data + offsets[i]), wrapped_src,

static_cast<tf::uint64>(offsets[i + 1] - offsets[i]))

.ok();

OP_REQUIRES_ASYNC(context, success,

tf::errors::Internal("D2H memcpy failed to be enqueued."), done);

}

context->device()->tensorflow_gpu_device_info()->event_mgr->ThenExecute(

d2h_stream, std::move(wrapped_callback));

#endif

}

}

private:

std::string plasma_store_socket_name_;

tf::mutex mu_;

bool connected_ = false;

plasma::PlasmaClient client_ GUARDED_BY(mu_);

public:

explicit PlasmaToTensorOp(tf::OpKernelConstruction* context) : tf::AsyncOpKernel(context) {

OP_REQUIRES_OK(context, context->GetAttr("plasma_store_socket_name",

&plasma_store_socket_name_));

tf::mutex_lock lock(mu_);

if (!connected_) {

VLOG(1) << "Connecting to Plasma...";

ARROW_CHECK_OK(client_.Connect(plasma_store_socket_name_));

VLOG(1) << "Connected!";

connected_ = true;

}

}

~PlasmaToTensorOp() override {

{

tf::mutex_lock lock(mu_);

ARROW_CHECK_OK(client_.Disconnect());

connected_ = false;

}

{

tf::mutex_lock lock(h2d_stream_mu);

if (h2d_stream != nullptr) {

delete h2d_stream;

}

}

}

void ComputeAsync(tf::OpKernelContext* context, DoneCallback done) override {

const tf::Tensor& plasma_object_id = context->input(0);

CHECK_EQ(plasma_object_id.NumElements(), 1);

const std::string& plasma_object_id_str = plasma_object_id.flat<std::string>()(0);

VLOG(1) << "plasma_object_id_str: '" << plasma_object_id_str << "'";

const plasma::ObjectID object_id =

plasma::ObjectID::from_binary(plasma_object_id_str);

plasma::ObjectBuffer object_buffer;

{

tf::mutex_lock lock(mu_);

// NOTE(zongheng): this is a blocking call. We might want to (1) make

// Plasma asynchronous, (2) launch a thread / event here ourselves, or

// something like that...

ARROW_CHECK_OK(client_.Get(&object_id, /*num_objects=*/1,

/*timeout_ms=*/-1, &object_buffer));

}

std::shared_ptr<arrow::Tensor> ndarray;

ARROW_CHECK_OK(arrow::py::NdarrayFromBuffer(object_buffer.data, &ndarray));

int64_t byte_width = get_byte_width(*ndarray->type());

const int64_t size_in_bytes = ndarray->data()->size();

tf::TensorShape shape({static_cast<int64_t>(size_in_bytes / byte_width)});

const float* plasma_data = reinterpret_cast<const float*>(ndarray->raw_data());

tf::Tensor* output_tensor = nullptr;

OP_REQUIRES_OK_ASYNC(context, context->allocate_output(0, shape, &output_tensor),

done);

auto wrapped_callback = [this, context, done, plasma_data, object_id]() {

{

tf::mutex_lock lock(mu_);

ARROW_CHECK_OK(client_.Release(object_id));

#ifdef GOOGLE_CUDA

auto orig_stream = context->op_device_context()->stream();

auto stream_executor = orig_stream->parent();

CHECK(stream_executor->HostMemoryUnregister(

const_cast<void*>(static_cast<const void*>(plasma_data))));

#endif

}

done();

};

if (std::is_same<Device, CPUDevice>::value) {

std::memcpy(

reinterpret_cast<void*>(const_cast<char*>(output_tensor->tensor_data().data())),

plasma_data, size_in_bytes);

wrapped_callback();

} else {

#ifdef GOOGLE_CUDA

auto orig_stream = context->op_device_context()->stream();

OP_REQUIRES_ASYNC(context, orig_stream != nullptr,

tf::errors::Internal("No GPU stream available."), done);

auto stream_executor = orig_stream->parent();

{

tf::mutex_lock l(h2d_stream_mu);

if (h2d_stream == nullptr) {

h2d_stream = new Stream(stream_executor);

CHECK(h2d_stream->Init().ok());

}

}

// Important. See note in T2P op.

CHECK(stream_executor->HostMemoryRegister(

const_cast<void*>(static_cast<const void*>(plasma_data)),

static_cast<tf::uint64>(size_in_bytes)));

perftools::gputools::DeviceMemoryBase wrapped_dst(

reinterpret_cast<void*>(const_cast<char*>(output_tensor->tensor_data().data())));

const bool success =

h2d_stream

->ThenMemcpy(&wrapped_dst, static_cast<const void*>(plasma_data),

static_cast<tf::uint64>(size_in_bytes))

.ok();

OP_REQUIRES_ASYNC(context, success,

tf::errors::Internal("H2D memcpy failed to be enqueued."), done);

// Without this sync the main compute stream might proceed to use the

// Tensor buffer, but its contents might still be in-flight from our

// h2d_stream.

CHECK(orig_stream->ThenWaitFor(h2d_stream).ok());

context->device()->tensorflow_gpu_device_info()->event_mgr->ThenExecute(

h2d_stream, std::move(wrapped_callback));

#endif

}

}

private:

std::string plasma_store_socket_name_;

tf::mutex mu_;

bool connected_ = false;

plasma::PlasmaClient client_ GUARDED_BY(mu_);