first commit

.venv/Lib/site-packages/pyarrow/tensorflow/plasma_op.cc

@@ -0,0 +1,391 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+#include "tensorflow/core/framework/device_base.h"
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/platform/logging.h"
+#include "tensorflow/core/platform/mutex.h"
+#include "tensorflow/stream_executor/device_memory.h"
+#include "tensorflow/stream_executor/event.h"
+#include "tensorflow/stream_executor/stream.h"
+
+#ifdef GOOGLE_CUDA
+#include "tensorflow/core/common_runtime/gpu/gpu_event_mgr.h"
+#include "tensorflow/core/platform/stream_executor.h"
+#endif
+
+#include "arrow/adapters/tensorflow/convert.h"
+#include "arrow/api.h"
+#include "arrow/io/api.h"
+#include "arrow/util/logging.h"
+
+// These headers do not include Python.h
+#include "arrow/python/deserialize.h"
+#include "arrow/python/serialize.h"
+
+#include "plasma/client.h"
+
+namespace tf = tensorflow;
+
+using ArrowStatus = arrow::Status;
+using CPUDevice = Eigen::ThreadPoolDevice;
+using GPUDevice = Eigen::GpuDevice;
+
+using Event = perftools::gputools::Event;
+using Stream = perftools::gputools::Stream;
+
+// NOTE(zongheng): for some reason using unique_ptr or shared_ptr results in
+// CUDA_ERROR_DEINITIALIZED on program exit.  I suspect this is because the
+// static object's dtor gets called *after* TensorFlow's own CUDA cleanup.
+// Instead, we use a raw pointer here and manually clean up in the Ops' dtors.
+static Stream* d2h_stream = nullptr;
+static tf::mutex d2h_stream_mu;
+
+// TODO(zongheng): CPU kernels' std::memcpy might be able to be sped up by
+// parallelization.
+
+int64_t get_byte_width(const arrow::DataType& dtype) {
+  return arrow::internal::checked_cast<const arrow::FixedWidthType&>(dtype)
+      .bit_width() / CHAR_BIT;
+}
+
+// Put:  tf.Tensor -> plasma.
+template <typename Device>
+class TensorToPlasmaOp : public tf::AsyncOpKernel {
+ 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_);
+};
+
+static Stream* h2d_stream = nullptr;
+static tf::mutex h2d_stream_mu;
+
+// Get:  plasma -> tf.Tensor.
+template <typename Device>
+class PlasmaToTensorOp : public tf::AsyncOpKernel {
+ 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_);
+};
+
+REGISTER_OP("TensorToPlasma")
+    .Input("input_tensor: dtypes")
+    .Input("plasma_object_id: string")
+    .Attr("dtypes: list(type)")
+    .Attr("plasma_store_socket_name: string");
+
+REGISTER_KERNEL_BUILDER(Name("TensorToPlasma").Device(tf::DEVICE_CPU),
+                        TensorToPlasmaOp<CPUDevice>);
+#ifdef GOOGLE_CUDA
+REGISTER_KERNEL_BUILDER(Name("TensorToPlasma").Device(tf::DEVICE_GPU),
+                        TensorToPlasmaOp<GPUDevice>);
+#endif
+
+REGISTER_OP("PlasmaToTensor")
+    .Input("plasma_object_id: string")
+    .Output("tensor: dtype")
+    .Attr("dtype: type")
+    .Attr("plasma_store_socket_name: string");
+
+REGISTER_KERNEL_BUILDER(Name("PlasmaToTensor").Device(tf::DEVICE_CPU),
+                        PlasmaToTensorOp<CPUDevice>);
+#ifdef GOOGLE_CUDA
+REGISTER_KERNEL_BUILDER(Name("PlasmaToTensor").Device(tf::DEVICE_GPU),
+                        PlasmaToTensorOp<GPUDevice>);
+#endif