Set device in more places

Change-Id: Ic152224d842b7c19079487d3ba1ca531e51e0687
NonVolatileComputing · Aug 24, 2017 · b646f96 · b646f96
1 parent 750b77d
commit b646f96
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 3 deletions.
diff --git a/cpp/src/arrow/gpu/cuda_common.h b/cpp/src/arrow/gpu/cuda_common.h
@@ -45,6 +45,17 @@ namespace gpu {
     }                                                                  \
   } while (0)
 
+#define CUDADRV_RETURN_NOT_OK(STMT)                                           \
+  do {                                                                        \
+    CUresult ret = (STMT);                                                    \
+    if (ret != CUDA_SUCCESS) {                                                \
+      std::stringstream ss;                                                   \
+      ss << "Cuda Driver API call in " << __FILE__ << " at line " << __LINE__ \
+         << " failed: " << #STMT;                                             \
+      return Status::IOError(ss.str());                                       \
+    }                                                                         \
+  } while (0)
+
 }  // namespace gpu
 }  // namespace arrow
 

diff --git a/cpp/src/arrow/gpu/cuda_memory.cc b/cpp/src/arrow/gpu/cuda_memory.cc
@@ -71,12 +71,13 @@ CudaHostBuffer::~CudaHostBuffer() { CUDA_DCHECK(cudaFreeHost(mutable_data_)); }
 // CudaBufferReader
 
 CudaBufferReader::CudaBufferReader(const std::shared_ptr<CudaBuffer>& buffer)
-    : io::BufferReader(buffer), cuda_buffer_(buffer) {}
+    : io::BufferReader(buffer), cuda_buffer_(buffer), gpu_number_(buffer->gpu_number()) {}
 
 CudaBufferReader::~CudaBufferReader() {}
 
 Status CudaBufferReader::Read(int64_t nbytes, int64_t* bytes_read, uint8_t* buffer) {
   nbytes = std::min(nbytes, size_ - position_);
+  CUDA_RETURN_NOT_OK(cudaSetDevice(gpu_number_));
   CUDA_RETURN_NOT_OK(
       cudaMemcpy(buffer, data_ + position_, nbytes, cudaMemcpyDeviceToHost));
   *bytes_read = nbytes;
@@ -95,7 +96,10 @@ Status CudaBufferReader::Read(int64_t nbytes, std::shared_ptr<Buffer>* out) {
 // CudaBufferWriter
 
 CudaBufferWriter::CudaBufferWriter(const std::shared_ptr<CudaBuffer>& buffer)
-    : io::FixedSizeBufferWriter(buffer), buffer_size_(0), buffer_position_(0) {}
+    : io::FixedSizeBufferWriter(buffer),
+      gpu_number_(buffer->gpu_number()),
+      buffer_size_(0),
+      buffer_position_(0) {}
 
 CudaBufferWriter::~CudaBufferWriter() {}
 
@@ -104,6 +108,7 @@ Status CudaBufferWriter::Close() { return Flush(); }
 Status CudaBufferWriter::Flush() {
   if (buffer_size_ > 0 && buffer_position_ > 0) {
     // Only need to flush when the write has been buffered
+    CUDA_RETURN_NOT_OK(cudaSetDevice(gpu_number_));
     CUDA_RETURN_NOT_OK(cudaMemcpy(mutable_data_ + position_ - buffer_position_,
                                   host_buffer_data_, buffer_position_,
                                   cudaMemcpyHostToDevice));
@@ -132,6 +137,7 @@ Status CudaBufferWriter::Write(const uint8_t* data, int64_t nbytes) {
     if (nbytes + buffer_position_ >= buffer_size_) {
       // Reach end of buffer, write everything
       RETURN_NOT_OK(Flush());
+      CUDA_RETURN_NOT_OK(cudaSetDevice(gpu_number_));
       CUDA_RETURN_NOT_OK(
           cudaMemcpy(mutable_data_ + position_, data, nbytes, cudaMemcpyHostToDevice));
     } else {
@@ -141,6 +147,7 @@ Status CudaBufferWriter::Write(const uint8_t* data, int64_t nbytes) {
     }
   } else {
     // Unbuffered write
+    CUDA_RETURN_NOT_OK(cudaSetDevice(gpu_number_));
     CUDA_RETURN_NOT_OK(
         cudaMemcpy(mutable_data_ + position_, data, nbytes, cudaMemcpyHostToDevice));
   }

diff --git a/cpp/src/arrow/gpu/cuda_memory.h b/cpp/src/arrow/gpu/cuda_memory.h
@@ -97,8 +97,8 @@ class ARROW_EXPORT CudaBufferReader : public io::BufferReader {
   Status Read(int64_t nbytes, std::shared_ptr<Buffer>* out) override;
 
  private:
-  // In case we need to access anything GPU-specific, like device number
   std::shared_ptr<CudaBuffer> cuda_buffer_;
+  int gpu_number_;
 };
 
 /// \class CudaBufferWriter
@@ -132,6 +132,8 @@ class ARROW_EXPORT CudaBufferWriter : public io::FixedSizeBufferWriter {
   int64_t num_bytes_buffered() const { return buffer_position_; }
 
  private:
+  int gpu_number_;
+
   // Pinned host buffer for buffering writes on CPU before calling cudaMalloc
   int64_t buffer_size_;
   int64_t buffer_position_;