Skip to content

Commit e8f04a8

Browse files
wesmwesm
committed
Set up libarrow_gpu, add simple unit test that allocates memory on device
Change-Id: I40e6b5d77f895d7803dfad64dd9fc2d9f882c0f3
1 parent 6ad976e commit e8f04a8

10 files changed

Lines changed: 388 additions & 1 deletion

File tree

‎cpp/CMakeLists.txt‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -98,6 +98,10 @@ if("${CMAKE_SOURCE_DIR}" STREQUAL "${CMAKE_CURRENT_SOURCE_DIR}")
9898
"Build the Arrow IPC extensions"
9999
ON)
100100

101+
option(ARROW_GPU
102+
"Build the Arrow GPU extensions (requires CUDA installation)"
103+
OFF)
104+
101105
option(ARROW_JEMALLOC
102106
"Build the Arrow jemalloc-based allocator"
103107
OFF)
@@ -713,6 +717,10 @@ if (ARROW_IPC)
713717
add_dependencies(arrow_dependencies metadata_fbs)
714718
endif()
715719
720+
if (ARROW_GPU)
721+
add_subdirectory(src/arrow/gpu)
722+
endif()
723+
716724
set(ARROW_SRCS
717725
src/arrow/array.cc
718726
src/arrow/buffer.cc

‎cpp/src/arrow/builder.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -124,6 +124,8 @@ class ARROW_EXPORT ArrayBuilder {
124124
std::shared_ptr<DataType> type() const { return type_; }
125125

126126
protected:
127+
ArrayBuilder() {}
128+
127129
std::shared_ptr<DataType> type_;
128130
MemoryPool* pool_;
129131

‎cpp/src/arrow/gpu/CMakeLists.txt‎

Lines changed: 111 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,111 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
function(ADD_ARROW_CUDA_TEST REL_TEST_NAME)
19+
set(options)
20+
set(single_value_args)
21+
set(multi_value_args STATIC_LINK_LIBS)
22+
cmake_parse_arguments(ARG "${options}" "${one_value_args}" "${multi_value_args}" ${ARGN})
23+
if(ARG_UNPARSED_ARGUMENTS)
24+
message(SEND_ERROR "Error: unrecognized arguments: ${ARG_UNPARSED_ARGUMENTS}")
25+
endif()
26+
27+
if(NO_TESTS OR NOT ARROW_BUILD_STATIC)
28+
return()
29+
endif()
30+
get_filename_component(TEST_NAME ${REL_TEST_NAME} NAME_WE)
31+
32+
if(EXISTS ${CMAKE_CURRENT_SOURCE_DIR}/${REL_TEST_NAME}.cc)
33+
# This test has a corresponding .cc file, set it up as an executable.
34+
set(TEST_PATH "${EXECUTABLE_OUTPUT_PATH}/${TEST_NAME}")
35+
cuda_add_executable(${TEST_NAME} "${REL_TEST_NAME}.cc")
36+
37+
if (ARG_STATIC_LINK_LIBS)
38+
# Customize link libraries
39+
target_link_libraries(${TEST_NAME} ${ARG_STATIC_LINK_LIBS})
40+
else()
41+
target_link_libraries(${TEST_NAME} ${ARROW_TEST_LINK_LIBS})
42+
endif()
43+
add_dependencies(unittest ${TEST_NAME})
44+
else()
45+
# No executable, just invoke the test (probably a script) directly.
46+
set(TEST_PATH ${CMAKE_CURRENT_SOURCE_DIR}/${REL_TEST_NAME})
47+
endif()
48+
49+
if (ARROW_TEST_MEMCHECK)
50+
SET_PROPERTY(TARGET ${TEST_NAME}
51+
APPEND_STRING PROPERTY
52+
COMPILE_FLAGS " -DARROW_VALGRIND")
53+
add_test(${TEST_NAME}
54+
bash -c "cd ${EXECUTABLE_OUTPUT_PATH}; valgrind --tool=memcheck --leak-check=full --leak-check-heuristics=stdstring --error-exitcode=1 ${TEST_PATH}")
55+
elseif(MSVC)
56+
add_test(${TEST_NAME} ${TEST_PATH})
57+
else()
58+
add_test(${TEST_NAME}
59+
${BUILD_SUPPORT_DIR}/run-test.sh ${CMAKE_BINARY_DIR} test ${TEST_PATH})
60+
endif()
61+
set_tests_properties(${TEST_NAME} PROPERTIES LABELS "unittest")
62+
endfunction()
63+
64+
#######################################
65+
# arrow_gpu
66+
#######################################
67+
68+
if (DEFINED ENV{CUDA_HOME})
69+
set(CUDA_TOOLKIT_ROOT_DIR "$ENV{CUDA_HOME}")
70+
endif()
71+
72+
find_package(CUDA REQUIRED)
73+
include_directories(SYSTEM ${CUDA_INCLUDE_DIRS})
74+
75+
set(ARROW_GPU_SRCS
76+
cuda_memory.cc
77+
)
78+
79+
set(ARROW_GPU_SHARED_LINK_LIBS
80+
arrow_shared
81+
)
82+
83+
cuda_add_library(arrow_gpu SHARED
84+
${ARROW_GPU_SRCS}
85+
)
86+
87+
install(FILES
88+
cuda_common.h
89+
cuda_memory.h
90+
DESTINATION "${CMAKE_INSTALL_INCLUDEDIR}/arrow/gpu")
91+
92+
set(ARROW_GPU_TEST_LINK_LIBS
93+
arrow_shared
94+
arrow_gpu_shared
95+
)
96+
97+
# pkg-config support
98+
configure_file(arrow-gpu.pc.in
99+
"${CMAKE_CURRENT_BINARY_DIR}/arrow-gpu.pc"
100+
@ONLY)
101+
install(
102+
FILES "${CMAKE_CURRENT_BINARY_DIR}/arrow-gpu.pc"
103+
DESTINATION "${CMAKE_INSTALL_LIBDIR}/pkgconfig/")
104+
105+
if (ARROW_BUILD_TESTS)
106+
set(ARROW_GPU_TEST_LINK_LIBS
107+
${ARROW_TEST_LINK_LIBS}
108+
arrow_gpu)
109+
ADD_ARROW_CUDA_TEST(cuda-test
110+
STATIC_LINK_LIBS ${ARROW_GPU_TEST_LINK_LIBS})
111+
endif()

‎cpp/src/arrow/gpu/arrow-gpu.pc.in‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
# Licensed to the Apache Software Foundation (ASF) under one
2+
# or more contributor license agreements. See the NOTICE file
3+
# distributed with this work for additional information
4+
# regarding copyright ownership. The ASF licenses this file
5+
# to you under the Apache License, Version 2.0 (the
6+
# "License"); you may not use this file except in compliance
7+
# with the License. You may obtain a copy of the License at
8+
#
9+
# http://www.apache.org/licenses/LICENSE-2.0
10+
#
11+
# Unless required by applicable law or agreed to in writing,
12+
# software distributed under the License is distributed on an
13+
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
# KIND, either express or implied. See the License for the
15+
# specific language governing permissions and limitations
16+
# under the License.
17+
18+
libdir=@CMAKE_INSTALL_FULL_LIBDIR@
19+
includedir=@CMAKE_INSTALL_FULL_INCLUDEDIR@
20+
21+
Name: Apache Arrow GPU
22+
Description: GPU integration library for Apache Arrow
23+
Version: @ARROW_VERSION@
24+
Requires: arrow
25+
Libs: -L${libdir} -larrow_gpu
26+
Cflags: -I${includedir}

‎cpp/src/arrow/gpu/cuda-test.cc‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
#include <cstdint>
19+
#include <limits>
20+
#include <string>
21+
22+
#include "gtest/gtest.h"
23+
24+
#include "arrow/status.h"
25+
#include "arrow/test-util.h"
26+
27+
#include "arrow/gpu/cuda_memory.h"
28+
29+
namespace arrow {
30+
namespace gpu {
31+
32+
class TestCudaBuffer : public ::testing::Test {};
33+
34+
TEST_F(TestCudaBuffer, Allocate) {
35+
const int device = 0;
36+
37+
const int64_t kSize = 100;
38+
std::shared_ptr<CudaBuffer> buffer;
39+
40+
ASSERT_OK(AllocateCudaBuffer(device, kSize, &buffer));
41+
ASSERT_EQ(kSize, buffer->size());
42+
}
43+
44+
} // namespace gpu
45+
} // namespace arrow

‎cpp/src/arrow/gpu/cuda_common.h‎

Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
// Non-public header
19+
20+
#ifndef ARROW_GPU_CUDA_COMMON_H
21+
#define ARROW_GPU_CUDA_COMMON_H
22+
23+
#include <cuda_runtime_api.h>
24+
25+
namespace arrow {
26+
namespace gpu {
27+
28+
#define CUDA_DCHECK(STMT) \
29+
do { \
30+
int ret = (STMT); \
31+
DCHECK_EQ(0, ret); \
32+
(void)ret; \
33+
} while (0)
34+
35+
#define CUDA_RETURN_NOT_OK(STMT) \
36+
do { \
37+
cudaError_t ret = (STMT); \
38+
if (ret != cudaSuccess) { \
39+
return Status::IOError("Cuda API call failed: " #STMT); \
40+
} \
41+
} while (0)
42+
43+
} // namespace gpu
44+
} // namespace arrow
45+
46+
#endif // ARROW_GPU_CUDA_COMMON_H

‎cpp/src/arrow/gpu/cuda_memory.cc‎

Lines changed: 65 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,65 @@
1+
// Licensed to the Apache Software Foundation (ASF) under one
2+
// or more contributor license agreements. See the NOTICE file
3+
// distributed with this work for additional information
4+
// regarding copyright ownership. The ASF licenses this file
5+
// to you under the Apache License, Version 2.0 (the
6+
// "License"); you may not use this file except in compliance
7+
// with the License. You may obtain a copy of the License at
8+
//
9+
// http://www.apache.org/licenses/LICENSE-2.0
10+
//
11+
// Unless required by applicable law or agreed to in writing,
12+
// software distributed under the License is distributed on an
13+
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14+
// KIND, either express or implied. See the License for the
15+
// specific language governing permissions and limitations
16+
// under the License.
17+
18+
#include "arrow/gpu/cuda_memory.h"
19+
20+
#include <cstdint>
21+
#include <memory>
22+
23+
#include "arrow/buffer.h"
24+
#include "arrow/status.h"
25+
#include "arrow/util/logging.h"
26+
27+
#include "arrow/gpu/cuda_common.h"
28+
29+
namespace arrow {
30+
namespace gpu {
31+
32+
CudaBuffer::~CudaBuffer() {
33+
if (own_data_) {
34+
CUDA_DCHECK(cudaFree(mutable_data_));
35+
}
36+
}
37+
38+
Status CudaBuffer::CopyHost(uint8_t* out) {
39+
CUDA_RETURN_NOT_OK(cudaMemcpy(out, data_, size_, cudaMemcpyDeviceToHost));
40+
return Status::OK();
41+
}
42+
43+
Status AllocateCudaBuffer(int gpu_number, const int64_t size,
44+
std::shared_ptr<CudaBuffer>* out) {
45+
CUDA_RETURN_NOT_OK(cudaSetDevice(gpu_number));
46+
uint8_t* data = nullptr;
47+
CUDA_RETURN_NOT_OK(
48+
cudaMalloc(reinterpret_cast<void**>(&data), static_cast<size_t>(size)));
49+
*out = std::make_shared<CudaBuffer>(data, size, gpu_number, true);
50+
return Status::OK();
51+
}
52+
53+
CudaHostBuffer::~CudaHostBuffer() { CUDA_DCHECK(cudaFreeHost(mutable_data_)); }
54+
55+
Status AllocateCudaHostBuffer(const int gpu_number, const int64_t size,
56+
std::shared_ptr<CudaHostBuffer>* out) {
57+
uint8_t* data = nullptr;
58+
CUDA_RETURN_NOT_OK(
59+
cudaMallocHost(reinterpret_cast<void**>(&data), static_cast<size_t>(size)));
60+
*out = std::make_shared<CudaHostBuffer>(data, size);
61+
return Status::OK();
62+
}
63+
64+
} // namespace gpu
65+
} // namespace arrow

0 commit comments

Comments
 (0)