host_device_vector.h

前往此文件文档。

 
#ifndef XGBOOST_HOST_DEVICE_VECTOR_H_
#define XGBOOST_HOST_DEVICE_VECTOR_H_
 
#include <xgboost/context.h> // 用于 DeviceOrd
#include <xgboost/span.h> // 用于 Span
 
#include <initializer_list>
#include <type_traits>
#include <vector>
 
namespace xgboost {
 
#ifdef __CUDACC__
// 设置一个函数来代替 cudaSetDevice() 进行调用；
// 仅为测试添加
void SetCudaSetDeviceHandler(void (*handler)(int));
#endif // __CUDACC__
 
template <typename T> struct HostDeviceVectorImpl;
 
enum GPUAccess {
 kNone, kRead,
 // 写操作意味着读操作
 kWrite
};
 
template <typename T>
class HostDeviceVector {
 static_assert(std::is_standard_layout_v<T>, "HostDeviceVector 仅支持 POD 类型");
 
 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_