Skip to content

Commit f3c724e

Browse files
wesmwesm
committed
Get things compiling / linking using driver API
Change-Id: I24d9d9510c8164dea36d83028b3c4bdbddbe2d85
1 parent 5d686fe commit f3c724e

6 files changed

Lines changed: 68 additions & 51 deletions

File tree

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

Lines changed: 7 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -79,28 +79,16 @@ set(ARROW_GPU_SRCS
7979

8080
set(ARROW_GPU_SHARED_LINK_LIBS
8181
arrow_shared
82+
${CUDA_LIBRARIES}
83+
${CUDA_CUDA_LIBRARY}
8284
)
8385

84-
add_library(arrow_gpu_objlib OBJECT
85-
${ARROW_GPU_SRCS}
86+
ADD_ARROW_LIB(arrow_gpu
87+
SOURCES ${ARROW_GPU_SRCS}
88+
SHARED_LINK_FLAGS ""
89+
SHARED_LINK_LIBS ${ARROW_GPU_SHARED_LINK_LIBS}
90+
STATIC_LINK_LIBS ""
8691
)
87-
set_property(TARGET arrow_gpu_objlib PROPERTY POSITION_INDEPENDENT_CODE 1)
88-
89-
if (ARROW_BUILD_SHARED)
90-
cuda_add_library(arrow_gpu_shared SHARED $<TARGET_OBJECTS:arrow_gpu_objlib>)
91-
install(TARGETS arrow_gpu_shared
92-
RUNTIME DESTINATION bin
93-
LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
94-
ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR})
95-
endif()
96-
97-
if (ARROW_BUILD_STATIC)
98-
add_library(arrow_gpu_static STATIC $<TARGET_OBJECTS:arrow_gpu_objlib>)
99-
install(TARGETS arrow_gpu_static
100-
RUNTIME DESTINATION bin
101-
LIBRARY DESTINATION ${CMAKE_INSTALL_LIBDIR}
102-
ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR})
103-
endif()
10492

10593
install(FILES
10694
cuda_common.h

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

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,10 +38,10 @@ static void CudaBufferWriterBenchmark(benchmark::State& state, const int64_t tot
3838
CudaDeviceManager* manager;
3939
ABORT_NOT_OK(CudaDeviceManager::GetInstance(&manager));
4040
std::shared_ptr<CudaContext> context;
41-
ABORT_NOT_OK(manager->CreateContext(kGpuNumber, &context));
41+
ABORT_NOT_OK(manager->GetContext(kGpuNumber, &context));
4242

4343
std::shared_ptr<CudaBuffer> device_buffer;
44-
ABORT_NOT_OK(context->Allocate(total_bytes, &device_buffer));
44+
ABORT_NOT_OK(AllocateCudaBuffer(total_bytes, context, &device_buffer));
4545
CudaBufferWriter writer(device_buffer);
4646

4747
if (buffer_size > 0) {

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

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -35,7 +35,7 @@ class TestCudaBufferBase : public ::testing::Test {
3535
public:
3636
void SetUp() {
3737
ASSERT_OK(CudaDeviceManager::GetInstance(&manager_));
38-
ASSERT_OK(manager_->CreateContext(kGpuNumber, &context_));
38+
ASSERT_OK(manager_->GetContext(kGpuNumber, &context_));
3939
}
4040

4141
protected:
@@ -53,7 +53,7 @@ class TestCudaBuffer : public TestCudaBufferBase {
5353
TEST_F(TestCudaBuffer, Allocate) {
5454
const int64_t kSize = 100;
5555
std::shared_ptr<CudaBuffer> buffer;
56-
ASSERT_OK(context_->Allocate(kSize, &buffer));
56+
ASSERT_OK(AllocateCudaBuffer(kSize, context_, &buffer));
5757
ASSERT_EQ(kSize, buffer->size());
5858
}
5959

@@ -68,7 +68,7 @@ void AssertCudaBufferEquals(const CudaBuffer& buffer, const uint8_t* host_data,
6868
TEST_F(TestCudaBuffer, CopyFromHost) {
6969
const int64_t kSize = 1000;
7070
std::shared_ptr<CudaBuffer> device_buffer;
71-
ASSERT_OK(context_->Allocate(kSize, &device_buffer));
71+
ASSERT_OK(AllocateCudaBuffer(kSize, context_, &device_buffer));
7272

7373
std::shared_ptr<PoolBuffer> host_buffer;
7474
ASSERT_OK(test::MakeRandomBytePoolBuffer(kSize, default_memory_pool(), &host_buffer));
@@ -86,7 +86,7 @@ class TestCudaBufferWriter : public TestCudaBufferBase {
8686
}
8787

8888
void Allocate(const int64_t size) {
89-
ASSERT_OK(context_->Allocate(size, &device_buffer_));
89+
ASSERT_OK(AllocateCudaBuffer(size, context_, &device_buffer_));
9090
writer_.reset(new CudaBufferWriter(device_buffer_));
9191
}
9292

@@ -198,7 +198,7 @@ TEST_F(TestCudaBufferReader, Basics) {
198198
std::shared_ptr<CudaBuffer> device_buffer;
199199

200200
const int64_t size = 1000;
201-
ASSERT_OK(context_->Allocate(size, &device_buffer));
201+
ASSERT_OK(AllocateCudaBuffer(size, context_, &device_buffer));
202202

203203
std::shared_ptr<PoolBuffer> buffer;
204204
ASSERT_OK(test::MakeRandomBytePoolBuffer(1000, default_memory_pool(), &buffer));

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

Lines changed: 44 additions & 19 deletions
Original file line numberDiff line numberDiff line change
@@ -21,9 +21,14 @@
2121
#include <cstdint>
2222
#include <memory>
2323
#include <string>
24+
#include <unordered_map>
25+
#include <vector>
2426

2527
#include <cuda.h>
2628

29+
#include "arrow/gpu/cuda_common.h"
30+
#include "arrow/gpu/cuda_memory.h"
31+
2732
namespace arrow {
2833
namespace gpu {
2934

@@ -39,7 +44,7 @@ class CudaContext::CudaContextImpl {
3944

4045
Status Init(const CudaDevice& device) {
4146
device_ = device;
42-
CU_RETURN_NOT_OK(cuCtxCreate(context_, 0, device_.handle));
47+
CU_RETURN_NOT_OK(cuCtxCreate(&context_, 0, device_.handle));
4348
is_open_ = true;
4449
return Status::OK();
4550
}
@@ -58,22 +63,23 @@ class CudaContext::CudaContextImpl {
5863
CU_RETURN_NOT_OK(cuCtxSetCurrent(context_));
5964

6065
CUdeviceptr data;
61-
CU_RETURN_NOT_OK(cuMemAlloc(&data, nbytes));
66+
CU_RETURN_NOT_OK(cuMemAlloc(&data, static_cast<size_t>(nbytes)));
6267
*out = reinterpret_cast<uint8_t*>(data);
6368
return Status::OK();
6469
}
6570

6671
Status CopyHostToDevice(uint8_t* dst, const uint8_t* src, int64_t nbytes) {
6772
CU_RETURN_NOT_OK(cuCtxSetCurrent(context_));
68-
CU_RETURN_NOT_OK(cuMemcpyDtoH(reinterpret_cast<CUdeviceptr>(dst),
69-
src, nbytes));
70-
return Statsu::OK();
73+
CU_RETURN_NOT_OK(cuMemcpyHtoD(reinterpret_cast<CUdeviceptr>(dst),
74+
reinterpret_cast<const void*>(src),
75+
static_cast<size_t>(nbytes)));
76+
return Status::OK();
7177
}
7278

7379
Status CopyDeviceToHost(uint8_t* dst, const uint8_t* src, int64_t nbytes) {
7480
CU_RETURN_NOT_OK(cuCtxSetCurrent(context_));
75-
CU_RETURN_NOT_OK(cuMemcpyHtoD(src, reinterpret_cast<const CUdeviceptr>(src),
76-
nbytes));
81+
CU_RETURN_NOT_OK(cuMemcpyDtoH(dst, reinterpret_cast<const CUdeviceptr>(src),
82+
static_cast<size_t>(nbytes)));
7783
return Status::OK();
7884
}
7985

@@ -113,7 +119,8 @@ class CudaDeviceManager::CudaDeviceManagerImpl {
113119

114120
Status AllocateHost(int64_t nbytes, uint8_t** out) {
115121
CU_RETURN_NOT_OK(cuMemHostAlloc(reinterpret_cast<void**>(out),
116-
nbytes, CU_MEMHOSTALLOC_PORTABLE));
122+
static_cast<size_t>(nbytes),
123+
CU_MEMHOSTALLOC_PORTABLE));
117124
host_bytes_allocated_ += nbytes;
118125
return Status::OK();
119126
}
@@ -134,9 +141,16 @@ class CudaDeviceManager::CudaDeviceManagerImpl {
134141
return Status::OK();
135142
}
136143

137-
Status CreateContext(int device_number, std::shared_ptr<CudaContext>* out) {
138-
*out = std::shared_ptr<CudaContext>(new CudaContext());
139-
return (*out)->impl_->Init(devices_[i]);
144+
Status GetContext(int device_number, std::shared_ptr<CudaContext>* out) {
145+
auto it = contexts_.find(device_number);
146+
if (it == contexts_.end()) {
147+
auto ctx = std::shared_ptr<CudaContext>(new CudaContext());
148+
RETURN_NOT_OK(ctx->impl_->Init(devices_[device_number]));
149+
contexts_[device_number] = *out = ctx;
150+
} else {
151+
*out = it->second;
152+
}
153+
return Status::OK();
140154
}
141155

142156
int num_devices() const { return num_devices_; }
@@ -145,13 +159,18 @@ class CudaDeviceManager::CudaDeviceManagerImpl {
145159
int num_devices_;
146160
std::vector<CudaDevice> devices_;
147161

162+
// device_number -> CudaContext
163+
std::unordered_map<int, std::shared_ptr<CudaContext>> contexts_;
164+
148165
int host_bytes_allocated_;
149166
};
150167

151168
CudaDeviceManager::CudaDeviceManager() {
152169
impl_.reset(new CudaDeviceManagerImpl());
153170
}
154171

172+
std::unique_ptr<CudaDeviceManager> CudaDeviceManager::instance_ = nullptr;
173+
155174
Status CudaDeviceManager::GetInstance(CudaDeviceManager** manager) {
156175
if (!instance_) {
157176
instance_.reset(new CudaDeviceManager());
@@ -161,21 +180,21 @@ Status CudaDeviceManager::GetInstance(CudaDeviceManager** manager) {
161180
return Status::OK();
162181
}
163182

164-
Status CudaDeviceManager::Create(int device_number,
165-
std::shared_ptr<CudaContext>* out) {
166-
return impl_->Create(device_number, out);
183+
Status CudaDeviceManager::GetContext(int device_number,
184+
std::shared_ptr<CudaContext>* out) {
185+
return impl_->GetContext(device_number, out);
167186
}
168187

169188
Status CudaDeviceManager::AllocateHost(int64_t nbytes,
170189
std::shared_ptr<CudaHostBuffer>* out) {
171-
uint8_t* data;
190+
uint8_t* data = nullptr;
172191
RETURN_NOT_OK(impl_->AllocateHost(nbytes, &data));
173-
*out = std::shared_ptr<CudaHostBuffer>(data, nbytes);
192+
*out = std::make_shared<CudaHostBuffer>(data, nbytes);
174193
return Status::OK();
175194
}
176195

177196
Status CudaDeviceManager::FreeHost(uint8_t* data, int64_t nbytes) {
178-
return impl_->FreeHost(data, nbytes));
197+
return impl_->FreeHost(data, nbytes);
179198
}
180199

181200
int CudaDeviceManager::num_devices() const {
@@ -185,8 +204,14 @@ int CudaDeviceManager::num_devices() const {
185204
// ----------------------------------------------------------------------
186205
// CudaContext public API
187206

188-
Status CudaContext::Allocate(int64_t nbytes, std::shared_ptr<CudaBuffer>* out) {
189-
return impl_->AllocateHost(nbytes, out);
207+
CudaContext::CudaContext() {
208+
impl_.reset(new CudaContextImpl());
209+
}
210+
211+
CudaContext::~CudaContext() {}
212+
213+
Status CudaContext::Allocate(int64_t nbytes, uint8_t** out) {
214+
return impl_->Allocate(nbytes, out);
190215
}
191216

192217
Status CudaContext::CopyHostToDevice(uint8_t* dst, const uint8_t* src, int64_t nbytes) {

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -37,16 +37,17 @@ class ARROW_EXPORT CudaDeviceManager {
3737
public:
3838
static Status GetInstance(CudaDeviceManager** manager);
3939

40-
/// \brief Create a CUDA driver context for a particular device
41-
Status CreateContext(int gpu_number, std::shared_ptr<CudaContext>* ctx);
40+
/// \brief Get the CUDA driver context for a particular device
41+
Status GetContext(int gpu_number, std::shared_ptr<CudaContext>* ctx);
4242

4343
Status AllocateHost(int64_t nbytes, std::shared_ptr<CudaHostBuffer>* buffer);
4444
Status FreeHost(uint8_t* data, int64_t nbytes);
4545

4646
int num_devices() const;
4747

4848
private:
49-
std::unique_ptr<CudaDeviceManager> instance_;
49+
CudaDeviceManager();
50+
static std::unique_ptr<CudaDeviceManager> instance_;
5051

5152
class CudaDeviceManagerImpl;
5253
std::unique_ptr<CudaDeviceManagerImpl> impl_;
@@ -67,7 +68,7 @@ class ARROW_EXPORT CudaContext {
6768
Status CopyHostToDevice(uint8_t* dst, const uint8_t* src, int64_t nbytes);
6869
Status CopyDeviceToHost(uint8_t* dst, const uint8_t* src, int64_t nbytes);
6970

70-
Status Allocate(int64_t nbytes, std::shared_ptr<CudaBuffer>* buffer);
71+
Status Allocate(int64_t nbytes, uint8_t** out);
7172
Status Free(uint8_t* device_ptr, int64_t nbytes);
7273

7374
int64_t bytes_allocated() const;

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

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -55,9 +55,12 @@ Status CudaBuffer::CopyFromHost(const int64_t position, const uint8_t* data,
5555
}
5656

5757
Status AllocateCudaBuffer(const int64_t size,
58-
std::shared_ptr<CudaContext>& context,
58+
const std::shared_ptr<CudaContext>& context,
5959
std::shared_ptr<CudaBuffer>* out) {
60-
return context->Allocate(size, out);
60+
uint8_t* data = nullptr;
61+
RETURN_NOT_OK(context->Allocate(size, &data));
62+
*out = std::make_shared<CudaBuffer>(data, size, context);
63+
return Status::OK();
6164
}
6265

6366
CudaHostBuffer::~CudaHostBuffer() {

0 commit comments

Comments
 (0)