Add classes and methods for simplifying use of CUDA IPC machinery. No tests yet

wesm · wesm · commit 84e452527cd7 · 2017-08-28T10:38:25.000-04:00
Change-Id: Ib46b646219e83c35828cf19a5e8d3bc8cc096f25
diff --git a/cpp/src/arrow/gpu/CMakeLists.txt b/cpp/src/arrow/gpu/CMakeLists.txt
@@ -45,7 +45,8 @@ ADD_ARROW_LIB(arrow_gpu
 )
 
 install(FILES
-  cuda_common.h
+  cuda_api.h
+  cuda_context.h
   cuda_memory.h
   DESTINATION "${CMAKE_INSTALL_INCLUDEDIR}/arrow/gpu")
 
diff --git a/cpp/src/arrow/gpu/cuda-benchmark.cc b/cpp/src/arrow/gpu/cuda-benchmark.cc
@@ -25,7 +25,7 @@
 #include "arrow/memory_pool.h"
 #include "arrow/test-util.h"
 
-#include "arrow/gpu/cuda_memory.h"
+#include "arrow/gpu/cuda_api.h"
 
 namespace arrow {
 namespace gpu {
@@ -41,7 +41,7 @@ static void CudaBufferWriterBenchmark(benchmark::State& state, const int64_t tot
   ABORT_NOT_OK(manager->GetContext(kGpuNumber, &context));
 
   std::shared_ptr<CudaBuffer> device_buffer;
-  ABORT_NOT_OK(AllocateCudaBuffer(total_bytes, context, &device_buffer));
+  ABORT_NOT_OK(context->Allocate(total_bytes, &device_buffer));
   CudaBufferWriter writer(device_buffer);
 
   if (buffer_size > 0) {
diff --git a/cpp/src/arrow/gpu/cuda-test.cc b/cpp/src/arrow/gpu/cuda-test.cc
@@ -24,7 +24,7 @@
 #include "arrow/status.h"
 #include "arrow/test-util.h"
 
-#include "arrow/gpu/cuda_memory.h"
+#include "arrow/gpu/cuda_api.h"
 
 namespace arrow {
 namespace gpu {
@@ -51,7 +51,7 @@ class TestCudaBuffer : public TestCudaBufferBase {
 TEST_F(TestCudaBuffer, Allocate) {
   const int64_t kSize = 100;
   std::shared_ptr<CudaBuffer> buffer;
-  ASSERT_OK(AllocateCudaBuffer(kSize, context_, &buffer));
+  ASSERT_OK(context_->Allocate(kSize, &buffer));
   ASSERT_EQ(kSize, buffer->size());
 }
 
@@ -66,7 +66,7 @@ void AssertCudaBufferEquals(const CudaBuffer& buffer, const uint8_t* host_data,
 TEST_F(TestCudaBuffer, CopyFromHost) {
   const int64_t kSize = 1000;
   std::shared_ptr<CudaBuffer> device_buffer;
-  ASSERT_OK(AllocateCudaBuffer(kSize, context_, &device_buffer));
+  ASSERT_OK(context_->Allocate(kSize, &device_buffer));
 
   std::shared_ptr<PoolBuffer> host_buffer;
   ASSERT_OK(test::MakeRandomBytePoolBuffer(kSize, default_memory_pool(), &host_buffer));
@@ -82,7 +82,7 @@ class TestCudaBufferWriter : public TestCudaBufferBase {
   void SetUp() { TestCudaBufferBase::SetUp(); }
 
   void Allocate(const int64_t size) {
-    ASSERT_OK(AllocateCudaBuffer(size, context_, &device_buffer_));
+    ASSERT_OK(context_->Allocate(size, &device_buffer_));
     writer_.reset(new CudaBufferWriter(device_buffer_));
   }
 
@@ -189,7 +189,7 @@ TEST_F(TestCudaBufferReader, Basics) {
   std::shared_ptr<CudaBuffer> device_buffer;
 
   const int64_t size = 1000;
-  ASSERT_OK(AllocateCudaBuffer(size, context_, &device_buffer));
+  ASSERT_OK(context_->Allocate(size, &device_buffer));
 
   std::shared_ptr<PoolBuffer> buffer;
   ASSERT_OK(test::MakeRandomBytePoolBuffer(1000, default_memory_pool(), &buffer));
diff --git a/cpp/src/arrow/gpu/cuda_api.h b/cpp/src/arrow/gpu/cuda_api.h
@@ -1,3 +1,4 @@
+<<<<<<< HEAD
 // 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
@@ -18,6 +19,7 @@
 #ifndef ARROW_GPU_CUDA_API_H
 #define ARROW_GPU_CUDA_API_H
 
+#include "arrow/gpu/cuda_context.h"
 #include "arrow/gpu/cuda_memory.h"
 #include "arrow/gpu/cuda_version.h"
 
diff --git a/cpp/src/arrow/gpu/cuda_context.cc b/cpp/src/arrow/gpu/cuda_context.cc
@@ -88,6 +88,17 @@ class CudaContext::CudaContextImpl {
     return Status::OK();
   }
 
+  Status OpenIpcBuffer(const CudaIpcMemHandle& ipc_handle, uint8_t** out) {
+    CU_RETURN_NOT_OK(cuCtxSetCurrent(context_));
+    auto handle = reinterpret_cast<const CUipcMemHandle*>(ipc_handle.handle());
+
+    CUdeviceptr data;
+    CU_RETURN_NOT_OK(
+        cuIpcOpenMemHandle(&data, *handle, CU_IPC_MEM_LAZY_ENABLE_PEER_ACCESS));
+    *out = reinterpret_cast<uint8_t*>(data);
+    return Status::OK();
+  }
+
   const CudaDevice device() const { return device_; }
 
  private:
@@ -203,8 +214,11 @@ CudaContext::CudaContext() { impl_.reset(new CudaContextImpl()); }
 
 CudaContext::~CudaContext() {}
 
-Status CudaContext::Allocate(int64_t nbytes, uint8_t** out) {
-  return impl_->Allocate(nbytes, out);
+Status CudaContext::Allocate(int64_t nbytes, std::shared_ptr<CudaBuffer>* out) {
+  uint8_t* data = nullptr;
+  RETURN_NOT_OK(impl_->Allocate(nbytes, &data));
+  *out = std::make_shared<CudaBuffer>(data, nbytes, this->shared_from_this(), true);
+  return Status::OK();
 }
 
 Status CudaContext::CopyHostToDevice(uint8_t* dst, const uint8_t* src, int64_t nbytes) {
@@ -219,5 +233,20 @@ Status CudaContext::Free(uint8_t* device_ptr, int64_t nbytes) {
   return impl_->Free(device_ptr, nbytes);
 }
 
+Status CudaContext::OpenIpcBuffer(const CudaIpcMemHandle& ipc_handle,
+                                  std::shared_ptr<CudaBuffer>* out) {
+  uint8_t* data = nullptr;
+  RETURN_NOT_OK(impl_->OpenIpcBuffer(ipc_handle, &data));
+
+  // Need to ask the device how big the buffer is
+  size_t allocation_size = 0;
+  CU_RETURN_NOT_OK(cuMemGetAddressRange(nullptr, &allocation_size,
+                                        reinterpret_cast<CUdeviceptr>(data)));
+
+  *out = std::make_shared<CudaBuffer>(data, allocation_size, this->shared_from_this(),
+                                      true, true);
+  return Status::OK();
+}
+
 }  // namespace gpu
 }  // namespace arrow
diff --git a/cpp/src/arrow/gpu/cuda_context.h b/cpp/src/arrow/gpu/cuda_context.h
@@ -24,12 +24,11 @@
 #include "arrow/status.h"
 #include "arrow/util/visibility.h"
 
+#include "arrow/gpu/cuda_memory.h"
+
 namespace arrow {
 namespace gpu {
 
-class CudaBuffer;
-class CudaHostBuffer;
-
 // Forward declaration
 class CudaContext;
 
@@ -41,6 +40,7 @@ class ARROW_EXPORT CudaDeviceManager {
   Status GetContext(int gpu_number, std::shared_ptr<CudaContext>* ctx);
 
   Status AllocateHost(int64_t nbytes, std::shared_ptr<CudaHostBuffer>* buffer);
+
   Status FreeHost(uint8_t* data, int64_t nbytes);
 
   int num_devices() const;
@@ -59,26 +59,40 @@ struct ARROW_EXPORT CudaDeviceInfo {};
 
 /// \class CudaContext
 /// \brief Friendlier interface to the CUDA driver API
-class ARROW_EXPORT CudaContext {
+class ARROW_EXPORT CudaContext : public std::enable_shared_from_this<CudaContext> {
  public:
   ~CudaContext();
 
   Status Destroy();
 
-  Status CopyHostToDevice(uint8_t* dst, const uint8_t* src, int64_t nbytes);
-  Status CopyDeviceToHost(uint8_t* dst, const uint8_t* src, int64_t nbytes);
+  /// \brief Allocate CUDA memory on GPU device for this context
+  /// \param[in] nbytes number of bytes
+  /// \param[out] out the allocated buffer
+  /// \return Status
+  Status Allocate(int64_t nbytes, std::shared_ptr<CudaBuffer>* out);
 
-  Status Allocate(int64_t nbytes, uint8_t** out);
-  Status Free(uint8_t* device_ptr, int64_t nbytes);
+  /// \brief Open existing CUDA IPC memory handle
+  /// \param[in] ipc_handle opaque pointer to CUipcMemHandle (driver API)
+  /// \param[out] buffer a CudaBuffer referencing
+  /// \return Status
+  Status OpenIpcBuffer(const CudaIpcMemHandle& ipc_handle,
+                       std::shared_ptr<CudaBuffer>* buffer);
 
   int64_t bytes_allocated() const;
 
  private:
   CudaContext();
 
+  Status CopyHostToDevice(uint8_t* dst, const uint8_t* src, int64_t nbytes);
+  Status CopyDeviceToHost(uint8_t* dst, const uint8_t* src, int64_t nbytes);
+  Status Free(uint8_t* device_ptr, int64_t nbytes);
+
   class CudaContextImpl;
   std::unique_ptr<CudaContextImpl> impl_;
 
+  friend CudaBuffer;
+  friend CudaBufferReader;
+  friend CudaBufferWriter;
   friend CudaDeviceManager::CudaDeviceManagerImpl;
 };
 
diff --git a/cpp/src/arrow/gpu/cuda_memory.cc b/cpp/src/arrow/gpu/cuda_memory.cc
@@ -19,8 +19,11 @@
 
 #include <algorithm>
 #include <cstdint>
+#include <cstdlib>
 #include <memory>
 
+#include <cuda.h>
+
 #include "arrow/buffer.h"
 #include "arrow/io/memory.h"
 #include "arrow/status.h"
@@ -32,15 +35,69 @@
 namespace arrow {
 namespace gpu {
 
-CudaBuffer::~CudaBuffer() {
+// ----------------------------------------------------------------------
+// CUDA IPC memory handle
+
+struct CudaIpcMemHandle::CudaIpcMemHandleImpl {
+  explicit CudaIpcMemHandleImpl(const void* handle) {
+    memcpy(&ipc_handle, handle, sizeof(CUipcMemHandle));
+  }
+
+  CUipcMemHandle ipc_handle;
+};
+
+CudaIpcMemHandle::CudaIpcMemHandle(const void* handle) {
+  impl_.reset(new CudaIpcMemHandleImpl(handle));
+}
+
+CudaIpcMemHandle::~CudaIpcMemHandle() {}
+
+Status CudaIpcMemHandle::FromBuffer(const void* opaque_handle,
+                                    std::unique_ptr<CudaIpcMemHandle>* handle) {
+  *handle = std::unique_ptr<CudaIpcMemHandle>(new CudaIpcMemHandle(opaque_handle));
+  return Status::OK();
+}
+
+Status CudaIpcMemHandle::Serialize(MemoryPool* pool, std::shared_ptr<Buffer>* out) const {
+  std::shared_ptr<MutableBuffer> buffer;
+  constexpr size_t kHandleSize = sizeof(CUipcMemHandle);
+  RETURN_NOT_OK(AllocateBuffer(pool, static_cast<int64_t>(kHandleSize), &buffer));
+  memcpy(buffer->mutable_data(), &impl_->ipc_handle, kHandleSize);
+  *out = buffer;
+  return Status::OK();
+}
+
+const void* CudaIpcMemHandle::handle() const { return &impl_->ipc_handle; }
+
+// ----------------------------------------------------------------------
+
+CudaBuffer::CudaBuffer(uint8_t* data, int64_t size,
+                       const std::shared_ptr<CudaContext>& context, bool own_data,
+                       bool is_ipc)
+    : Buffer(data, size), context_(context), own_data_(own_data), is_ipc_(is_ipc) {
+  is_mutable_ = true;
+  mutable_data_ = data;
+}
+
+CudaBuffer::~CudaBuffer() { DCHECK(Close().ok()); }
+
+Status CudaBuffer::Close() {
   if (own_data_) {
-    DCHECK(context_->Free(mutable_data_, size_).ok());
+    if (is_ipc_) {
+      CU_RETURN_NOT_OK(cuIpcCloseMemHandle(reinterpret_cast<CUdeviceptr>(mutable_data_)));
+    } else {
+      return context_->Free(mutable_data_, size_);
+    }
   }
+  return Status::OK();
 }
 
 CudaBuffer::CudaBuffer(const std::shared_ptr<CudaBuffer>& parent, const int64_t offset,
                        const int64_t size)
-    : Buffer(parent, offset, size), context_(parent->context()) {}
+    : Buffer(parent, offset, size),
+      context_(parent->context()),
+      own_data_(false),
+      is_ipc_(false) {}
 
 Status CudaBuffer::CopyToHost(const int64_t position, const int64_t nbytes,
                               uint8_t* out) const {
@@ -53,12 +110,15 @@ Status CudaBuffer::CopyFromHost(const int64_t position, const uint8_t* data,
   return context_->CopyHostToDevice(mutable_data_ + position, data, nbytes);
 }
 
-Status AllocateCudaBuffer(const int64_t size, const std::shared_ptr<CudaContext>& context,
-                          std::shared_ptr<CudaBuffer>* out) {
-  DCHECK(context);
-  uint8_t* data = nullptr;
-  RETURN_NOT_OK(context->Allocate(size, &data));
-  *out = std::make_shared<CudaBuffer>(data, size, context);
+Status CudaBuffer::ExportForIpc(std::unique_ptr<CudaIpcMemHandle>* handle) {
+  if (is_ipc_) {
+    return Status::Invalid("Buffer has already been exported for IPC");
+  }
+  CUipcMemHandle cu_handle;
+  CU_RETURN_NOT_OK(
+      cuIpcGetMemHandle(&cu_handle, reinterpret_cast<CUdeviceptr>(mutable_data_)));
+  is_ipc_ = true;
+  *handle = std::unique_ptr<CudaIpcMemHandle>(new CudaIpcMemHandle(&cu_handle));
   return Status::OK();
 }
 
diff --git a/cpp/src/arrow/gpu/cuda_memory.h b/cpp/src/arrow/gpu/cuda_memory.h

Original file line number	Diff line number	Diff line change
`@@ -45,7 +45,8 @@ ADD_ARROW_LIB(arrow_gpu`
`45`	`45`	`)`
`46`	`46`
`47`	`47`	`install(FILES`
`48`		`- cuda_common.h`
	`48`	`+ cuda_api.h`
	`49`	`+ cuda_context.h`
`49`	`50`	`cuda_memory.h`
`50`	`51`	`DESTINATION "${CMAKE_INSTALL_INCLUDEDIR}/arrow/gpu")`
`51`	`52`