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,
 // write implies read
 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]] 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_