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

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## New Features
 
  - PR #106 Added multi-GPU initialization
+ - PR #167 Added value setter to `device_scalar`
 
 ## Improvements
 

include/rmm/device_scalar.hpp

@@ -39,49 +39,49 @@ class device_scalar {
    * @brief Construct a new `device_scalar`
    *
    * @param initial_value The initial value of the object in device memory
-   * @param stream_ Optional, stream on which to perform allocation and copy
-   * @param mr_ Optional, resource with which to allocate
+   * @param stream Optional, stream on which to perform allocation and copy
+   * @param mr Optional, resource with which to allocate
    *---------------------------------------------------------------------------**/
   explicit 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."};
-    }
+      T const &initial_value, cudaStream_t stream = 0,
+      rmm::mr::device_memory_resource *mr = rmm::mr::get_default_resource())
+      : buff{sizeof(T), stream, mr} {
+
+    _memcpy(buff.data(), &initial_value, stream);
   }
 
   /**---------------------------------------------------------------------------*
-   * @brief Copies the value from device to host and returns the value.
+   * @brief Copies the value from device to host synchronously and returns the
+   * value.
    *
-   * @return T The value of the scalar after synchronizing its stream
+   * @return T The value of the scalar
+   * @param stream CUDA stream on which to perform the copy
    *---------------------------------------------------------------------------**/
-  T value() const {
+  T value(cudaStream_t stream = 0) 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."};
-    }
+    _memcpy(&host_value, buff.data(), stream);
     return host_value;
   }
 
+  /**---------------------------------------------------------------------------*
+   * @brief Copies the value from host to device synchronously.
+   *
+   * @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 = 0) {
+    _memcpy(buff.data(), &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 +92,14 @@ class device_scalar {
 
  private:
   rmm::device_buffer buff{sizeof(T)};
+
+  inline void _memcpy(void *dst, const void *src, cudaStream_t stream) const {
+    auto status = cudaMemcpyAsync(dst, src, sizeof(T), cudaMemcpyDefault,
+                                  stream);
+    if (cudaSuccess != status) {
+      throw std::runtime_error{"Device memcpy 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());
 }