From e8f04a868a3c209a4bf4393e3aaa7bd5f5adf494 Mon Sep 17 00:00:00 2001
From: Wes McKinney <wes.mckinney@twosigma.com>
Date: Sun, 20 Aug 2017 17:06:13 -0400
Subject: [PATCH] Set up libarrow_gpu, add simple unit test that allocates
 memory on device

Change-Id: I40e6b5d77f895d7803dfad64dd9fc2d9f882c0f3
---
 cpp/CMakeLists.txt                  |   8 ++
 cpp/src/arrow/builder.h             |   2 +
 cpp/src/arrow/gpu/CMakeLists.txt    | 111 ++++++++++++++++++++++++++++
 cpp/src/arrow/gpu/arrow-gpu.pc.in   |  26 +++++++
 cpp/src/arrow/gpu/cuda-test.cc      |  45 +++++++++++
 cpp/src/arrow/gpu/cuda_common.h     |  46 ++++++++++++
 cpp/src/arrow/gpu/cuda_memory.cc    |  65 ++++++++++++++++
 cpp/src/arrow/gpu/cuda_memory.h     |  78 +++++++++++++++++++
 cpp/src/arrow/python/CMakeLists.txt |   6 ++
 cpp/src/arrow/type_traits.h         |   2 +-
 10 files changed, 388 insertions(+), 1 deletion(-)
 create mode 100644 cpp/src/arrow/gpu/CMakeLists.txt
 create mode 100644 cpp/src/arrow/gpu/arrow-gpu.pc.in
 create mode 100644 cpp/src/arrow/gpu/cuda-test.cc
 create mode 100644 cpp/src/arrow/gpu/cuda_common.h
 create mode 100644 cpp/src/arrow/gpu/cuda_memory.cc
 create mode 100644 cpp/src/arrow/gpu/cuda_memory.h

diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index 7d73d1ffff089..b55a9bb1d0256 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -98,6 +98,10 @@ if("${CMAKE_SOURCE_DIR}" STREQUAL "${CMAKE_CURRENT_SOURCE_DIR}")
     "Build the Arrow IPC extensions"
     ON)
 
+  option(ARROW_GPU
+    "Build the Arrow GPU extensions (requires CUDA installation)"
+    OFF)
+
   option(ARROW_JEMALLOC
     "Build the Arrow jemalloc-based allocator"
     OFF)
@@ -713,6 +717,10 @@ if (ARROW_IPC)
   add_dependencies(arrow_dependencies metadata_fbs)
 endif()
 
+if (ARROW_GPU)
+  add_subdirectory(src/arrow/gpu)
+endif()
+
 set(ARROW_SRCS
   src/arrow/array.cc
   src/arrow/buffer.cc
diff --git a/cpp/src/arrow/builder.h b/cpp/src/arrow/builder.h
index 3b851f92c1726..687286d4eb88a 100644
--- a/cpp/src/arrow/builder.h
+++ b/cpp/src/arrow/builder.h
@@ -124,6 +124,8 @@ class ARROW_EXPORT ArrayBuilder {
   std::shared_ptr<DataType> type() const { return type_; }
 
  protected:
+  ArrayBuilder() {}
+
   std::shared_ptr<DataType> type_;
   MemoryPool* pool_;
 
diff --git a/cpp/src/arrow/gpu/CMakeLists.txt b/cpp/src/arrow/gpu/CMakeLists.txt
new file mode 100644
index 0000000000000..0a56e497fd8ab
--- /dev/null
+++ b/cpp/src/arrow/gpu/CMakeLists.txt
@@ -0,0 +1,111 @@
+# 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.
+
+function(ADD_ARROW_CUDA_TEST REL_TEST_NAME)
+  set(options)
+  set(single_value_args)
+  set(multi_value_args STATIC_LINK_LIBS)
+  cmake_parse_arguments(ARG "${options}" "${one_value_args}" "${multi_value_args}" ${ARGN})
+  if(ARG_UNPARSED_ARGUMENTS)
+    message(SEND_ERROR "Error: unrecognized arguments: ${ARG_UNPARSED_ARGUMENTS}")
+  endif()
+
+  if(NO_TESTS OR NOT ARROW_BUILD_STATIC)
+    return()
+  endif()
+  get_filename_component(TEST_NAME ${REL_TEST_NAME} NAME_WE)
+
+  if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${REL_TEST_NAME}.cc)
+    # This test has a corresponding .cc file, set it up as an executable.
+    set(TEST_PATH "${EXECUTABLE_OUTPUT_PATH}/${TEST_NAME}")
+    cuda_add_executable(${TEST_NAME} "${REL_TEST_NAME}.cc")
+
+    if (ARG_STATIC_LINK_LIBS)
+      # Customize link libraries
+      target_link_libraries(${TEST_NAME} ${ARG_STATIC_LINK_LIBS})
+    else()
+      target_link_libraries(${TEST_NAME} ${ARROW_TEST_LINK_LIBS})
+    endif()
+    add_dependencies(unittest ${TEST_NAME})
+  else()
+    # No executable, just invoke the test (probably a script) directly.
+    set(TEST_PATH ${CMAKE_CURRENT_SOURCE_DIR}/${REL_TEST_NAME})
+  endif()
+
+  if (ARROW_TEST_MEMCHECK)
+    SET_PROPERTY(TARGET ${TEST_NAME}
+      APPEND_STRING PROPERTY
+      COMPILE_FLAGS " -DARROW_VALGRIND")
+    add_test(${TEST_NAME}
+      bash -c "cd ${EXECUTABLE_OUTPUT_PATH}; valgrind --tool=memcheck --leak-check=full --leak-check-heuristics=stdstring --error-exitcode=1 ${TEST_PATH}")
+  elseif(MSVC)
+    add_test(${TEST_NAME} ${TEST_PATH})
+  else()
+    add_test(${TEST_NAME}
+      ${BUILD_SUPPORT_DIR}/run-test.sh ${CMAKE_BINARY_DIR} test ${TEST_PATH})
+  endif()
+  set_tests_properties(${TEST_NAME} PROPERTIES LABELS "unittest")
+endfunction()
+
+#######################################
+# arrow_gpu
+#######################################
+
+if (DEFINED ENV{CUDA_HOME})
+  set(CUDA_TOOLKIT_ROOT_DIR "$ENV{CUDA_HOME}")
+endif()
+
+find_package(CUDA REQUIRED)
+include_directories(SYSTEM ${CUDA_INCLUDE_DIRS})
+
+set(ARROW_GPU_SRCS
+  cuda_memory.cc
+)
+
+set(ARROW_GPU_SHARED_LINK_LIBS
+  arrow_shared
+)
+
+cuda_add_library(arrow_gpu SHARED
+  ${ARROW_GPU_SRCS}
+)
+
+install(FILES
+  cuda_common.h
+  cuda_memory.h
+  DESTINATION "${CMAKE_INSTALL_INCLUDEDIR}/arrow/gpu")
+
+set(ARROW_GPU_TEST_LINK_LIBS
+  arrow_shared
+  arrow_gpu_shared
+)
+
+# pkg-config support
+configure_file(arrow-gpu.pc.in
+  "${CMAKE_CURRENT_BINARY_DIR}/arrow-gpu.pc"
+  @ONLY)
+install(
+  FILES "${CMAKE_CURRENT_BINARY_DIR}/arrow-gpu.pc"
+  DESTINATION "${CMAKE_INSTALL_LIBDIR}/pkgconfig/")
+
+if (ARROW_BUILD_TESTS)
+  set(ARROW_GPU_TEST_LINK_LIBS
+    ${ARROW_TEST_LINK_LIBS}
+    arrow_gpu)
+  ADD_ARROW_CUDA_TEST(cuda-test
+    STATIC_LINK_LIBS ${ARROW_GPU_TEST_LINK_LIBS})
+endif()
diff --git a/cpp/src/arrow/gpu/arrow-gpu.pc.in b/cpp/src/arrow/gpu/arrow-gpu.pc.in
new file mode 100644
index 0000000000000..3889d03b204ca
--- /dev/null
+++ b/cpp/src/arrow/gpu/arrow-gpu.pc.in
@@ -0,0 +1,26 @@
+# 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.
+
+libdir=@CMAKE_INSTALL_FULL_LIBDIR@
+includedir=@CMAKE_INSTALL_FULL_INCLUDEDIR@
+
+Name: Apache Arrow GPU
+Description: GPU integration library for Apache Arrow
+Version: @ARROW_VERSION@
+Requires: arrow
+Libs: -L${libdir} -larrow_gpu
+Cflags: -I${includedir}
diff --git a/cpp/src/arrow/gpu/cuda-test.cc b/cpp/src/arrow/gpu/cuda-test.cc
new file mode 100644
index 0000000000000..9dc8e1aff946a
--- /dev/null
+++ b/cpp/src/arrow/gpu/cuda-test.cc
@@ -0,0 +1,45 @@
+// 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 <cstdint>
+#include <limits>
+#include <string>
+
+#include "gtest/gtest.h"
+
+#include "arrow/status.h"
+#include "arrow/test-util.h"
+
+#include "arrow/gpu/cuda_memory.h"
+
+namespace arrow {
+namespace gpu {
+
+class TestCudaBuffer : public ::testing::Test {};
+
+TEST_F(TestCudaBuffer, Allocate) {
+  const int device = 0;
+
+  const int64_t kSize = 100;
+  std::shared_ptr<CudaBuffer> buffer;
+
+  ASSERT_OK(AllocateCudaBuffer(device, kSize, &buffer));
+  ASSERT_EQ(kSize, buffer->size());
+}
+
+}  // namespace gpu
+}  // namespace arrow
diff --git a/cpp/src/arrow/gpu/cuda_common.h b/cpp/src/arrow/gpu/cuda_common.h
new file mode 100644
index 0000000000000..75f41c9a41eab
--- /dev/null
+++ b/cpp/src/arrow/gpu/cuda_common.h
@@ -0,0 +1,46 @@
+// 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.
+
+// Non-public header
+
+#ifndef ARROW_GPU_CUDA_COMMON_H
+#define ARROW_GPU_CUDA_COMMON_H
+
+#include <cuda_runtime_api.h>
+
+namespace arrow {
+namespace gpu {
+
+#define CUDA_DCHECK(STMT) \
+  do {                    \
+    int ret = (STMT);     \
+    DCHECK_EQ(0, ret);    \
+    (void)ret;            \
+  } while (0)
+
+#define CUDA_RETURN_NOT_OK(STMT)                              \
+  do {                                                        \
+    cudaError_t ret = (STMT);                                 \
+    if (ret != cudaSuccess) {                                 \
+      return Status::IOError("Cuda API call failed: " #STMT); \
+    }                                                         \
+  } while (0)
+
+}  // namespace gpu
+}  // namespace arrow
+
+#endif  // ARROW_GPU_CUDA_COMMON_H
diff --git a/cpp/src/arrow/gpu/cuda_memory.cc b/cpp/src/arrow/gpu/cuda_memory.cc
new file mode 100644
index 0000000000000..cf37bdae2f3ef
--- /dev/null
+++ b/cpp/src/arrow/gpu/cuda_memory.cc
@@ -0,0 +1,65 @@
+// 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 "arrow/gpu/cuda_memory.h"
+
+#include <cstdint>
+#include <memory>
+
+#include "arrow/buffer.h"
+#include "arrow/status.h"
+#include "arrow/util/logging.h"
+
+#include "arrow/gpu/cuda_common.h"
+
+namespace arrow {
+namespace gpu {
+
+CudaBuffer::~CudaBuffer() {
+  if (own_data_) {
+    CUDA_DCHECK(cudaFree(mutable_data_));
+  }
+}
+
+Status CudaBuffer::CopyHost(uint8_t* out) {
+  CUDA_RETURN_NOT_OK(cudaMemcpy(out, data_, size_, cudaMemcpyDeviceToHost));
+  return Status::OK();
+}
+
+Status AllocateCudaBuffer(int gpu_number, const int64_t size,
+                          std::shared_ptr<CudaBuffer>* out) {
+  CUDA_RETURN_NOT_OK(cudaSetDevice(gpu_number));
+  uint8_t* data = nullptr;
+  CUDA_RETURN_NOT_OK(
+      cudaMalloc(reinterpret_cast<void**>(&data), static_cast<size_t>(size)));
+  *out = std::make_shared<CudaBuffer>(data, size, gpu_number, true);
+  return Status::OK();
+}
+
+CudaHostBuffer::~CudaHostBuffer() { CUDA_DCHECK(cudaFreeHost(mutable_data_)); }
+
+Status AllocateCudaHostBuffer(const int gpu_number, const int64_t size,
+                              std::shared_ptr<CudaHostBuffer>* out) {
+  uint8_t* data = nullptr;
+  CUDA_RETURN_NOT_OK(
+      cudaMallocHost(reinterpret_cast<void**>(&data), static_cast<size_t>(size)));
+  *out = std::make_shared<CudaHostBuffer>(data, size);
+  return Status::OK();
+}
+
+}  // namespace gpu
+}  // namespace arrow
diff --git a/cpp/src/arrow/gpu/cuda_memory.h b/cpp/src/arrow/gpu/cuda_memory.h
new file mode 100644
index 0000000000000..885373a26c943
--- /dev/null
+++ b/cpp/src/arrow/gpu/cuda_memory.h
@@ -0,0 +1,78 @@
+// 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.
+
+#ifndef ARROW_GPU_CUDA_MEMORY_H
+#define ARROW_GPU_CUDA_MEMORY_H
+
+#include <cstdint>
+#include <memory>
+
+#include "arrow/buffer.h"
+#include "arrow/status.h"
+
+namespace arrow {
+namespace gpu {
+
+/// \brief An Arrow buffer located on a GPU device
+///
+/// Be careful using this in any Arrow code which may not be GPU-aware
+class ARROW_EXPORT CudaBuffer : public MutableBuffer {
+ public:
+  CudaBuffer(uint8_t* data, int64_t size, const int gpu_number, bool own_data = false)
+      : MutableBuffer(data, size), gpu_number_(gpu_number), own_data_(own_data) {}
+
+  ~CudaBuffer();
+
+  /// \brief Copy memory from GPU device to CPU host
+  /// \param[out] out a pre-allocated output buffer
+  /// \return Status
+  Status CopyHost(uint8_t* out);
+
+  int gpu_number() const { return gpu_number_; }
+
+ private:
+  const int gpu_number_;
+  bool own_data_;
+};
+
+/// \brief Device-accessible CPU memory created using cudaHostAlloc
+class ARROW_EXPORT CudaHostBuffer : public MutableBuffer {
+ public:
+  using MutableBuffer::MutableBuffer;
+  ~CudaHostBuffer();
+};
+
+/// \brief Allocate CUDA memory on a GPU device
+/// \param[in] gpu_number Device number to allocate
+/// \param[in] size number of bytes
+/// \param[out] out the allocated buffer
+/// \return Status
+ARROW_EXPORT
+Status AllocateCudaBuffer(const int gpu_number, const int64_t size,
+                          std::shared_ptr<CudaBuffer>* out);
+
+/// \brief Allocate CUDA-accessible memory on CPU host
+/// \param[in] size number of bytes
+/// \param[out] out the allocated buffer
+/// \return Status
+ARROW_EXPORT
+Status AllocateCudaHostBuffer(const int64_t size, std::shared_ptr<CudaHostBuffer>* out);
+
+}  // namespace gpu
+}  // namespace arrow
+
+#endif  // ARROW_GPU_CUDA_MEMORY_H
diff --git a/cpp/src/arrow/python/CMakeLists.txt b/cpp/src/arrow/python/CMakeLists.txt
index f2807b930a33c..84aad82e2a90e 100644
--- a/cpp/src/arrow/python/CMakeLists.txt
+++ b/cpp/src/arrow/python/CMakeLists.txt
@@ -8,6 +8,12 @@
 #
 #   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.
 
 #######################################
 # arrow_python
diff --git a/cpp/src/arrow/type_traits.h b/cpp/src/arrow/type_traits.h
index f05eb56718f5f..d424cc81ff92a 100644
--- a/cpp/src/arrow/type_traits.h
+++ b/cpp/src/arrow/type_traits.h
@@ -328,7 +328,7 @@ GET_ATTR(TypeClass, void);
   using TypeClass =                                                                 \
       typename std::conditional<std::is_base_of<DataType, T>::value, T,             \
                                 typename detail::GetAttr_TypeClass<T>::type>::type; \
-  using c_type = typename detail::GetAttr_c_type<TypeClass>::type;
+  using c_type = typename detail::GetAttr_c_type<TypeClass>::type
 
 template <typename T>
 struct IsUnsignedInt {