device_uvector.hpp

/*
 * Copyright (c) 2020-2022, 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/detail/error.hpp>
#include <rmm/detail/exec_check_disable.hpp>
#include <rmm/device_buffer.hpp>
#include <rmm/mr/device/device_memory_resource.hpp>
#include <rmm/mr/device/per_device_resource.hpp>
 
#include <cstddef>
#include <vector>
 
namespace rmm {
 
template <typename T>
class device_uvector {
  static_assert(std::is_trivially_copyable<T>::value,
                "device_uvector only supports types that are trivially copyable.");
 
 public:
  using value_type      = T;
  using size_type       = std::size_t;
  using reference       = value_type&;
  using const_reference = value_type const&;
  using pointer         = value_type*;
  using const_pointer   = value_type const*;
  using iterator        = pointer;
  using const_iterator  = const_pointer;
 
  RMM_EXEC_CHECK_DISABLE
  ~device_uvector() = default;
 
  RMM_EXEC_CHECK_DISABLE
  device_uvector(device_uvector&&) noexcept = default;
 
  device_uvector& operator=(device_uvector&&) noexcept = default;
 
  device_uvector(device_uvector const&) = delete;
 
  device_uvector& operator=(device_uvector const&) = delete;
 
  device_uvector() = delete;
 
  explicit device_uvector(
    std::size_t size,
    cuda_stream_view stream,
    rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource())
    : _storage{elements_to_bytes(size), stream, mr}
  {
  }
 
  explicit device_uvector(
    device_uvector const& other,
    cuda_stream_view stream,
    rmm::mr::device_memory_resource* mr = rmm::mr::get_current_device_resource())
    : _storage{other._storage, stream, mr}
  {
  }
 
  [[nodiscard]] pointer element_ptr(std::size_t element_index) noexcept
  {
    assert(element_index < size());
    return data() + element_index;
  }
 
  [[nodiscard]] const_pointer element_ptr(std::size_t element_index) const noexcept
  {
    assert(element_index < size());
    return data() + element_index;
  }
 
  void set_element_async(std::size_t element_index,
                         value_type const& value,
                         cuda_stream_view stream)
  {
    RMM_EXPECTS(
      element_index < size(), rmm::out_of_range, "Attempt to access out of bounds element.");
 
    if constexpr (std::is_same<value_type, bool>::value) {
      RMM_CUDA_TRY(
        cudaMemsetAsync(element_ptr(element_index), value, sizeof(value), stream.value()));
      return;
    }
 
    if constexpr (std::is_fundamental<value_type>::value) {
      if (value == value_type{0}) {
        set_element_to_zero_async(element_index, stream);
        return;
      }
    }
 
    RMM_CUDA_TRY(cudaMemcpyAsync(
      element_ptr(element_index), &value, sizeof(value), cudaMemcpyDefault, stream.value()));
  }
 
  // We delete the r-value reference overload to prevent asynchronously copying from a literal or
  // implicit temporary value after it is deleted or goes out of scope.
  void set_element_async(std::size_t, value_type const&&, cuda_stream_view) = delete;
 
  void set_element_to_zero_async(std::size_t element_index, cuda_stream_view stream)
  {
    RMM_EXPECTS(
      element_index < size(), rmm::out_of_range, "Attempt to access out of bounds element.");
    RMM_CUDA_TRY(
      cudaMemsetAsync(element_ptr(element_index), 0, sizeof(value_type), stream.value()));
  }
 
  void set_element(std::size_t element_index, T const& value, cuda_stream_view stream)
  {
    set_element_async(element_index, value, stream);
    stream.synchronize_no_throw();
  }
 
  [[nodiscard]] value_type element(std::size_t element_index, cuda_stream_view stream) const
  {
    RMM_EXPECTS(
      element_index < size(), rmm::out_of_range, "Attempt to access out of bounds element.");
    value_type value;
    RMM_CUDA_TRY(cudaMemcpyAsync(
      &value, element_ptr(element_index), sizeof(value), cudaMemcpyDefault, stream.value()));
    stream.synchronize();
    return value;
  }
 
  [[nodiscard]] value_type front_element(cuda_stream_view stream) const
  {
    return element(0, stream);
  }
 
  [[nodiscard]] value_type back_element(cuda_stream_view stream) const
  {
    return element(size() - 1, stream);
  }
 
  void resize(std::size_t new_size, cuda_stream_view stream)
  {
    _storage.resize(elements_to_bytes(new_size), stream);
  }
 
  void shrink_to_fit(cuda_stream_view stream) { _storage.shrink_to_fit(stream); }
 
  device_buffer release() noexcept { return std::move(_storage); }
 
  [[nodiscard]] std::size_t capacity() const noexcept
  {
    return bytes_to_elements(_storage.capacity());
  }
 
  [[nodiscard]] pointer data() noexcept { return static_cast<pointer>(_storage.data()); }
 
  [[nodiscard]] const_pointer data() const noexcept
  {
    return static_cast<const_pointer>(_storage.data());
  }
 
  [[nodiscard]] iterator begin() noexcept { return data(); }
 
  [[nodiscard]] const_iterator cbegin() const noexcept { return data(); }
 
  [[nodiscard]] const_iterator begin() const noexcept { return cbegin(); }
 
  [[nodiscard]] iterator end() noexcept { return data() + size(); }
 
  [[nodiscard]] const_iterator cend() const noexcept { return data() + size(); }
 
  [[nodiscard]] const_iterator end() const noexcept { return cend(); }
 
  [[nodiscard]] std::size_t size() const noexcept { return bytes_to_elements(_storage.size()); }
 
  [[nodiscard]] std::int64_t ssize() const noexcept
  {
    assert(size() < static_cast<std::size_t>(std::numeric_limits<int64_t>::max()) &&
           "Size overflows signed integer");
    return static_cast<int64_t>(size());
  }
 
  [[nodiscard]] bool is_empty() const noexcept { return size() == 0; }
 
  [[nodiscard]] mr::device_memory_resource* memory_resource() const noexcept
  {
    return _storage.memory_resource();
  }
 
  [[nodiscard]] cuda_stream_view stream() const noexcept { return _storage.stream(); }
 
  void set_stream(cuda_stream_view stream) noexcept { _storage.set_stream(stream); }
 
 private:
  device_buffer _storage{};  
 
  [[nodiscard]] std::size_t constexpr elements_to_bytes(std::size_t num_elements) const noexcept
  {
    return num_elements * sizeof(value_type);
  }
 
  [[nodiscard]] std::size_t constexpr bytes_to_elements(std::size_t num_bytes) const noexcept
  {
    return num_bytes / sizeof(value_type);
  }
};
}  // namespace rmm