Test CudaBuffer::CopyFromHost

Change-Id: I205b2861797c57756a38f4c695937d354d497b53
NonVolatileComputing · Aug 22, 2017 · 1cf1196 · 1cf1196
1 parent a2708f2
commit 1cf1196
Show file tree

Hide file tree

Showing 3 changed files with 30 additions and 17 deletions.
diff --git a/cpp/src/arrow/gpu/cuda-test.cc b/cpp/src/arrow/gpu/cuda-test.cc
@@ -34,16 +34,34 @@ constexpr int kGpuNumber = 0;
 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_OK(AllocateCudaBuffer(kGpuNumber, kSize, &buffer));
   ASSERT_EQ(kSize, buffer->size());
 }
 
-TEST_F(TestCudaBuffer, CopyFromHost) {}
+void AssertCudaBufferEquals(const CudaBuffer& buffer, const uint8_t* host_data,
+                            const int64_t nbytes) {
+  std::shared_ptr<MutableBuffer> result;
+  ASSERT_OK(AllocateBuffer(default_memory_pool(), nbytes, &result));
+  ASSERT_OK(buffer.CopyToHost(result->mutable_data()));
+  ASSERT_EQ(0, std::memcmp(result->data(), host_data, nbytes));
+}
+
+TEST_F(TestCudaBuffer, CopyFromHost) {
+  const int64_t kSize = 1000;
+  std::shared_ptr<CudaBuffer> device_buffer;
+  ASSERT_OK(AllocateCudaBuffer(kGpuNumber, kSize, &device_buffer));
+
+  std::shared_ptr<PoolBuffer> host_buffer;
+  ASSERT_OK(test::MakeRandomBytePoolBuffer(kSize, default_memory_pool(), &host_buffer));
+
+  ASSERT_OK(device_buffer->CopyFromHost(0, host_buffer->data(), 500));
+  ASSERT_OK(device_buffer->CopyFromHost(500, host_buffer->data() + 500, kSize - 500));
+
+  AssertCudaBufferEquals(*device_buffer, host_buffer->data(), kSize);
+}
 
 class TestCudaBufferWriter : public ::testing::Test {
  public:
@@ -83,12 +101,7 @@ class TestCudaBufferWriter : public ::testing::Test {
 
     ASSERT_OK(writer_->Flush());
 
-    std::shared_ptr<MutableBuffer> result;
-    ASSERT_OK(AllocateBuffer(default_memory_pool(), total_bytes, &result));
-
-    ASSERT_OK(device_buffer_->CopyToHost(result->mutable_data()));
-
-    ASSERT_EQ(0, std::memcmp(result->data(), buffer->data(), total_bytes));
+    AssertCudaBufferEquals(*device_buffer_, buffer->data(), total_bytes);
   }
 
  protected:
@@ -151,10 +164,7 @@ TEST_F(TestCudaBufferWriter, EdgeCases) {
   ASSERT_EQ(0, writer_->num_bytes_buffered());
 
   // Check that everything was written
-  std::shared_ptr<MutableBuffer> result;
-  ASSERT_OK(AllocateBuffer(default_memory_pool(), 1000, &result));
-  ASSERT_OK(device_buffer_->CopyToHost(result->mutable_data()));
-  ASSERT_EQ(0, std::memcmp(result->data(), host_data, 1000));
+  AssertCudaBufferEquals(*device_buffer_, host_data, 1000);
 }
 
 }  // namespace gpu

diff --git a/cpp/src/arrow/gpu/cuda_memory.cc b/cpp/src/arrow/gpu/cuda_memory.cc
@@ -37,14 +37,14 @@ CudaBuffer::~CudaBuffer() {
   }
 }
 
-Status CudaBuffer::CopyToHost(uint8_t* out) {
+Status CudaBuffer::CopyToHost(uint8_t* out) const {
   CUDA_RETURN_NOT_OK(cudaMemcpy(out, data_, size_, cudaMemcpyDeviceToHost));
   return Status::OK();
 }
 
 Status CudaBuffer::CopyFromHost(const int64_t position, const uint8_t* data,
                                 int64_t nbytes) {
-  DCHECK_LT(nbytes, size_ - position) << "Copy would overflow buffer";
+  DCHECK_LE(nbytes, size_ - position) << "Copy would overflow buffer";
   CUDA_RETURN_NOT_OK(
       cudaMemcpy(mutable_data_ + position, data, nbytes, cudaMemcpyHostToDevice));
   return Status::OK();

diff --git a/cpp/src/arrow/gpu/cuda_memory.h b/cpp/src/arrow/gpu/cuda_memory.h
@@ -43,7 +43,7 @@ class ARROW_EXPORT CudaBuffer : public MutableBuffer {
   /// \brief Copy memory from GPU device to CPU host
   /// \param[out] out a pre-allocated output buffer
   /// \return Status
-  Status CopyToHost(uint8_t* out);
+  Status CopyToHost(uint8_t* out) const;
 
   /// \brief Copy memory to device at position
   /// \param[in] position start position to copy bytes
@@ -108,7 +108,10 @@ class ARROW_EXPORT CudaBufferWriter : public io::FixedSizeBufferWriter {
   /// By default writes are unbuffered
   Status SetBufferSize(const int64_t buffer_size);
 
+  /// \brief Returns size of host (CPU) buffer, 0 for unbuffered
   int64_t buffer_size() const { return buffer_size_; }
+
+  /// \brief Returns number of bytes buffered on host
   int64_t num_bytes_buffered() const { return buffer_position_; }
 
  private: