device_scalar.hpp

/*
 * Copyright (c) 2019-2021, NVIDIA CORPORATION.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
 
#pragma once
 
#include <rmm/cuda_stream_view.hpp>
#include <rmm/device_uvector.hpp>
#include <rmm/mr/device/device_memory_resource.hpp>
#include <rmm/mr/device/per_device_resource.hpp>
 
#include <type_traits>
 
namespace rmm {
 
template <typename T>
class device_scalar {
 public:
  static_assert(std::is_trivially_copyable<T>::value, "Scalar type must be trivially copyable");
 
  using value_type      = typename device_uvector<T>::value_type;
  using reference       = typename device_uvector<T>::reference;
  using const_reference = typename device_uvector<T>::const_reference;
  using pointer         = typename device_uvector<T>::pointer;
  using const_pointer   = typename device_uvector<T>::const_pointer;
 
  RMM_EXEC_CHECK_DISABLE
  ~device_scalar() = default;
 
  RMM_EXEC_CHECK_DISABLE
  device_scalar(device_scalar&&) noexcept = default;
 
  device_scalar& operator=(device_scalar&&) noexcept = default;
 
  device_scalar(device_scalar const&) = delete;
 
  device_scalar& operator=(device_scalar const&) = delete;
 
  device_scalar() = delete;
 
  explicit device_scalar(
    cuda_stream_view stream,
    rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource())
    : _storage{1, stream, mr}
  {
  }
 
  explicit device_scalar(
    value_type const& initial_value,
    cuda_stream_view stream,
    rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource())
    : _storage{1, stream, mr}
  {
    set_value_async(initial_value, stream);
  }
 
  device_scalar(device_scalar const& other,
                cuda_stream_view stream,
                rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource())
    : _storage{other._storage, stream, mr}
  {
  }
 
  [[nodiscard]] value_type value(cuda_stream_view stream) const
  {
    return _storage.front_element(stream);
  }
 
  void set_value_async(value_type const& value, cuda_stream_view stream)
  {
    _storage.set_element_async(0, value, stream);
  }
 
  // Disallow passing literals to set_value to avoid race conditions where the memory holding the
  // literal can be freed before the async memcpy / memset executes.
  void set_value_async(value_type&&, cuda_stream_view) = delete;
 
  void set_value_to_zero_async(cuda_stream_view stream)
  {
    _storage.set_element_to_zero_async(value_type{0}, stream);
  }
 
  [[nodiscard]] pointer data() noexcept { return static_cast<pointer>(_storage.data()); }
 
  [[nodiscard]] const_pointer data() const noexcept
  {
    return static_cast<const_pointer>(_storage.data());
  }
 
  [[nodiscard]] cuda_stream_view stream() const noexcept { return _storage.stream(); }
 
  void set_stream(cuda_stream_view stream) noexcept { _storage.set_stream(stream); }
 
 private:
  rmm::device_uvector<T> _storage;
};
}  // namespace rmm