host_device_vector.h

Go to the documentation of this file.

 
#ifndef XGBOOST_HOST_DEVICE_VECTOR_H_
#define XGBOOST_HOST_DEVICE_VECTOR_H_
 
#include <xgboost/context.h>  // for DeviceOrd
#include <xgboost/span.h>     // for Span
 
#include <initializer_list>
#include <type_traits>
#include <vector>
 
namespace xgboost {
 
#ifdef __CUDACC__
// Sets a function to call instead of cudaSetDevice();
// only added for testing
void SetCudaSetDeviceHandler(void (*handler)(int));
#endif  // __CUDACC__
 
template <typename T>
class HostDeviceVectorImpl;
 
enum GPUAccess {
  kNone,
  kRead,
  // write implies read
  kWrite
};
 
template <typename T>
class HostDeviceVector {
  static_assert(std::is_standard_layout_v<T>, "HostDeviceVector admits only POD types");
 
 public:
  explicit HostDeviceVector(size_t size = 0, T v = T(), DeviceOrd device = DeviceOrd::CPU());
  HostDeviceVector(std::initializer_list<T> init, DeviceOrd device = DeviceOrd::CPU());
  explicit HostDeviceVector(const std::vector<T>& init, DeviceOrd device = DeviceOrd::CPU());
  ~HostDeviceVector();
 
  HostDeviceVector(const HostDeviceVector<T>&) = delete;
  HostDeviceVector(HostDeviceVector<T>&&);
 
  HostDeviceVector<T>& operator=(const HostDeviceVector<T>&) = delete;
  HostDeviceVector<T>& operator=(HostDeviceVector<T>&&);
 
  [[nodiscard]] bool Empty() const { return Size() == 0; }
  [[nodiscard]] std::size_t Size() const;
  [[nodiscard]] std::size_t SizeBytes() const { return this->Size() * sizeof(T); }
  [[nodiscard]] DeviceOrd Device() const;
  common::Span<T> DeviceSpan();
  common::Span<const T> ConstDeviceSpan() const;
  common::Span<const T> DeviceSpan() const { return ConstDeviceSpan(); }
  T* DevicePointer();
  const T* ConstDevicePointer() const;
  const T* DevicePointer() const { return ConstDevicePointer(); }
 
  T* HostPointer() { return HostVector().data(); }
  common::Span<T> HostSpan() { return common::Span<T>{HostVector()}; }
  common::Span<T const> HostSpan() const { return common::Span<T const>{HostVector()}; }
  common::Span<T const> ConstHostSpan() const { return HostSpan(); }
  const T* ConstHostPointer() const { return ConstHostVector().data(); }
  const T* HostPointer() const { return ConstHostPointer(); }
 
  void Fill(T v);
  void Copy(const HostDeviceVector<T>& other);
  void Copy(const std::vector<T>& other);
  void Copy(std::initializer_list<T> other);
 
  void Extend(const HostDeviceVector<T>& other);
 
  std::vector<T>& HostVector();
  const std::vector<T>& ConstHostVector() const;
  const std::vector<T>& HostVector() const { return ConstHostVector(); }
 
  [[nodiscard]] bool HostCanRead() const;
  [[nodiscard]] bool HostCanWrite() const;
  [[nodiscard]] bool DeviceCanRead() const;
  [[nodiscard]] bool DeviceCanWrite() const;
  [[nodiscard]] GPUAccess DeviceAccess() const;
 
  void SetDevice(DeviceOrd device) const;
 
  void Resize(std::size_t new_size);
  void Resize(std::size_t new_size, T v);
 
  using value_type = T;  // NOLINT
 
 private:
  HostDeviceVectorImpl<T>* impl_;
};
 
}  // namespace xgboost
 
#endif  // XGBOOST_HOST_DEVICE_VECTOR_H_