[REVIEW] impl and refactor device_scalar::set_value and value, respectively.

CHANGELOG.md

@@ -2,6 +2,8 @@
 
 ## New Features
 
+ - PR #167 Added value setter to `device_scalar`
+
 ## Improvements
 
 - PR #161 Use `std::atexit` to finalize RMM after Python interpreter shutdown

include/rmm/device_scalar.hpp

@@ -46,12 +46,8 @@ class device_scalar {
       T const &initial_value, cudaStream_t stream_ = 0,
       rmm::mr::device_memory_resource *mr_ = rmm::mr::get_default_resource())
       : buff{sizeof(T), stream_, mr_} {
-    auto status = cudaMemcpyAsync(buff.data(), &initial_value, sizeof(T),
-                                  cudaMemcpyDefault, buff.stream());
-
-    if (cudaSuccess != status) {
-      throw std::runtime_error{"Device memcpy failed."};
-    }
+
+    _set_value<false>(initial_value, buff.stream());
   }
 
   /**---------------------------------------------------------------------------*
@@ -60,28 +56,86 @@ class device_scalar {
    * @return T The value of the scalar after synchronizing its stream
    *---------------------------------------------------------------------------**/
   T value() const {
-    T host_value{};
-    auto status = cudaMemcpyAsync(&host_value, buff.data(), sizeof(T),
-                             cudaMemcpyDefault, buff.stream());
-    if (cudaSuccess != status) {
-      throw std::runtime_error{"Device memcpy failed."};
-    }
-    status = cudaStreamSynchronize(buff.stream());
-    if (cudaSuccess != status) {
-      throw std::runtime_error{"Stream sync failed."};
-    }
-    return host_value;
+    return _value<true>(buff.stream());
+  }
+
+  /**---------------------------------------------------------------------------*
+   * @brief Copies the value from device to host and returns the value.
+   *
+   * @return T The value of the scalar after synchronizing its stream
+   * @param stream CUDA stream on which to perform the copy
+   *---------------------------------------------------------------------------**/
+  T value(cudaStream_t stream) const {
+    return _value<true>(stream);
+  }
+
+  /**---------------------------------------------------------------------------*
+   * @brief Copies the value from device to host and returns the value.
+   *
+   * @return T The value of the scalar
+   *---------------------------------------------------------------------------**/
+  T value_async() const {
+    return _value<false>(buff.stream());
+  }
+
+  /**---------------------------------------------------------------------------*
+   * @brief Copies the value from device to host and returns the value.
+   *
+   * @return T The value of the scalar
+   * @param stream CUDA stream on which to perform the copy
+   *---------------------------------------------------------------------------**/
+  T value_async(cudaStream_t stream) const {
+    return _value<false>(stream);
+  }
+
+  /**---------------------------------------------------------------------------*
+   * @brief Copies the value from host to device and synchronizes.
+   *
+   * @param host_value The host value which will be copied to device
+   *---------------------------------------------------------------------------**/
+  void set_value(T host_value) {
+    _set_value<true>(host_value, buff.stream());
+  }
+
+
+  /**---------------------------------------------------------------------------*
+   * @brief Copies the value from host to device and synchronizes.
+   *
+   * @param host_value The host value which will be copied to device
+   * @param stream CUDA stream on which to perform the copy
+   *---------------------------------------------------------------------------**/
+  void set_value(T host_value, cudaStream_t stream) {
+    _set_value<true>(host_value, stream);
+  }
+
+  /**---------------------------------------------------------------------------*
+   * @brief Copies the value from host to device.
+   *
+   * @param host_value The host value which will be copied to device
+   *---------------------------------------------------------------------------**/
+  void set_value_async(T host_value) {
+    _set_value<false>(host_value, buff.stream());
+  }
+
+  /**---------------------------------------------------------------------------*
+   * @brief Copies the value from host to device.
+   *
+   * @param host_value The host value which will be copied to device
+   * @param stream CUDA stream on which to perform the copy
+   *---------------------------------------------------------------------------**/
+  void set_value_async(T host_value, cudaStream_t stream) {
+    _set_value<false>(host_value, stream);
   }
 
   /**---------------------------------------------------------------------------*
    * @brief Returns pointer to object in device memory.
    *---------------------------------------------------------------------------**/
-  T *get() noexcept { return static_cast<T *>(buff.data()); }
+  T *data() noexcept { return static_cast<T *>(buff.data()); }
 
   /**---------------------------------------------------------------------------*
    * @brief Returns pointer to object in device memory.
    *---------------------------------------------------------------------------**/
-  T const *get() const noexcept { return static_cast<T const *>(buff.data()); }
+  T const *data() const noexcept { return static_cast<T const *>(buff.data()); }
 
   device_scalar() = default;
   ~device_scalar() = default;
@@ -92,6 +146,37 @@ class device_scalar {
 
  private:
   rmm::device_buffer buff{sizeof(T)};
+
+  template<bool synchronize>
+  inline T _value(cudaStream_t stream) const {
+    T host_value{};
+    _memcpy<synchronize>(&host_value, buff.data(), sizeof(T), stream);
+    return host_value;
+  }
+
+  template<bool synchronize>
+  inline void _set_value(T host_value, cudaStream_t stream) {
+    _memcpy<synchronize>(buff.data(), &host_value, sizeof(T), stream);
+  }
+
+  template<bool synchronize>
+  inline void _memcpy(void *dst, const void *src, size_t count, cudaStream_t stream) const {
+    auto status = cudaMemcpyAsync(dst, src, count, cudaMemcpyDefault, stream);
+
+    if (cudaSuccess != status) {
+      throw std::runtime_error{"Device memcpy failed."};
+    }
+
+    if (false == synchronize) {
+      return;
+    }
+
+    status = cudaStreamSynchronize(stream);
+
+    if (cudaSuccess != status) {
+      throw std::runtime_error{"Stream sync failed."};
+    }
+  }
 };
 
 }  // namespace rmm

tests/device_scalar_tests.cpp

@@ -53,37 +53,47 @@ TYPED_TEST_CASE(DeviceScalarTest, Types);
 
 TYPED_TEST(DeviceScalarTest, DefaultUninitialized) {
   rmm::device_scalar<TypeParam> scalar{};
-  EXPECT_NE(nullptr, scalar.get());
+  EXPECT_NE(nullptr, scalar.data());
 }
 
 TYPED_TEST(DeviceScalarTest, InitialValue) {
   rmm::device_scalar<TypeParam> scalar{this->value, this->stream, this->mr};
-  EXPECT_NE(nullptr, scalar.get());
+  EXPECT_NE(nullptr, scalar.data());
   EXPECT_EQ(this->value, scalar.value());
 }
 
 TYPED_TEST(DeviceScalarTest, CopyCtor) {
   rmm::device_scalar<TypeParam> scalar{this->value, this->stream, this->mr};
-  EXPECT_NE(nullptr, scalar.get());
+  EXPECT_NE(nullptr, scalar.data());
   EXPECT_EQ(this->value, scalar.value());
 
   rmm::device_scalar<TypeParam> copy{scalar};
-  EXPECT_NE(nullptr, copy.get());
-  EXPECT_NE(copy.get(), scalar.get());
+  EXPECT_NE(nullptr, copy.data());
+  EXPECT_NE(copy.data(), scalar.data());
   EXPECT_EQ(copy.value(), scalar.value());
 }
 
 TYPED_TEST(DeviceScalarTest, MoveCtor) {
   rmm::device_scalar<TypeParam> scalar{this->value, this->stream, this->mr};
-  EXPECT_NE(nullptr, scalar.get());
+  EXPECT_NE(nullptr, scalar.data());
   EXPECT_EQ(this->value, scalar.value());
 
-  auto original_pointer = scalar.get();
+  auto original_pointer = scalar.data();
   auto original_value = scalar.value();
 
   rmm::device_scalar<TypeParam> moved_to{std::move(scalar)};
-  EXPECT_NE(nullptr, moved_to.get());
-  EXPECT_EQ(moved_to.get(), original_pointer);
+  EXPECT_NE(nullptr, moved_to.data());
+  EXPECT_EQ(moved_to.data(), original_pointer);
   EXPECT_EQ(moved_to.value(), original_value);
-  EXPECT_EQ(nullptr, scalar.get());
+  EXPECT_EQ(nullptr, scalar.data());
+}
+
+TYPED_TEST(DeviceScalarTest, SetValue) {
+  rmm::device_scalar<TypeParam> scalar{this->value, this->stream, this->mr};
+  EXPECT_NE(nullptr, scalar.data());
+
+  auto expected = this->distribution(this->generator);
+
+  scalar.set_value(expected);
+  EXPECT_EQ(expected, scalar.value());
 }