[REVIEW] Fix race condition in device_scalar::set_value

CHANGELOG.md

@@ -27,6 +27,7 @@
 - PR #545 Fix build to support using `clang` as the host compiler
 - PR #534 Fix `pool_memory_resource` failure when init and max pool sizes are equal
 - PR #546 Remove CUDA driver linking and correct NVTX macro.
+- PR #569 Correct `device_scalar::set_value` to pass host value by reference to avoid copying from invalid value
 
 
 # RMM 0.15.0 (26 Aug 2020)

include/rmm/device_scalar.hpp

@@ -114,7 +114,25 @@ class device_scalar {
    * (e.g. using `cudaStreamWaitEvent()` or `cudaStreamSynchronize()`) before and after calling
    * this function, otherwise there may be a race condition.
    *
-   * Does not synchronize `stream`.
+   * This function does not synchronize stream `stream` before returning. Therefore, the object
+   * referenced by `host_value` should not be destroyed or modified until `stream` has been
+   * synchronized. Otherwise, behavior is undefined.
+   *
+   * @note: This function incurs a host to device memcpy and should be used sparingly.
+
+   * Example:
+   * \code{cpp}
+   * rmm::device_scalar<int32_t> s;
+   *
+   * int v{42};
+   *
+   * // Copies 42 to element 0 on `stream`. Does _not_ synchronize
+   * vec.set_value(v, stream);
+   * ...
+   * cudaStreamSynchronize(stream);
+   * // Synchronization is required before `v` can be modified
+   * v = 13;
+   * \endcode
    *
    * @throws `rmm::cuda_error` if copying `host_value` to device memory fails.
    * @throws `rmm::cuda_error` if synchronizing `stream` fails.
@@ -123,7 +141,7 @@ class device_scalar {
    * @param stream CUDA stream on which to perform the copy
    */
   template <typename Dummy = void>
-  auto set_value(T host_value, cudaStream_t stream = 0)
+  auto set_value(T const &host_value, cudaStream_t stream = 0)
     -> std::enable_if_t<std::is_fundamental<T>::value, Dummy>
   {
     if (host_value == T{0}) {
@@ -141,7 +159,25 @@ class device_scalar {
    * (e.g. using `cudaStreamWaitEvent()` or `cudaStreamSynchronize()`) before and after calling
    * this function, otherwise there may be a race condition.
    *
-   * Does not synchronize `stream`.
+   * This function does not synchronize stream `stream` before returning. Therefore, the object
+   * referenced by `host_value` should not be destroyed or modified until `stream` has been
+   * synchronized. Otherwise, behavior is undefined.
+   *
+   * @note: This function incurs a host to device memcpy and should be used sparingly.
+
+   * Example:
+   * \code{cpp}
+   * rmm::device_scalar<int32_t> s;
+   *
+   * int v{42};
+   *
+   * // Copies 42 to element 0 on `stream`. Does _not_ synchronize
+   * vec.set_value(v, stream);
+   * ...
+   * cudaStreamSynchronize(stream);
+   * // Synchronization is required before `v` can be modified
+   * v = 13;
+   * \endcode
    *
    * @throws `rmm::cuda_error` if copying `host_value` to device memory fails
    * @throws `rmm::cuda_error` if synchronizing `stream` fails
@@ -150,7 +186,7 @@ class device_scalar {
    * @param stream CUDA stream on which to perform the copy
    */
   template <typename Dummy = void>
-  auto set_value(T host_value, cudaStream_t stream = 0)
+  auto set_value(T const &host_value, cudaStream_t stream = 0)
     -> std::enable_if_t<not std::is_fundamental<T>::value, Dummy>
   {
     _memcpy(buffer.data(), &host_value, stream);