context.h

前往此文件文档。

 
#ifndef XGBOOST_CONTEXT_H_
#define XGBOOST_CONTEXT_H_
 
#include <xgboost/base.h> // 对于 bst_d_ordinal_t
#include <xgboost/logging.h> // 对于 CHECK_GE
#include <xgboost/parameter.h> // 对于 XGBoostParameter
 
#include <cstdint> // 对于 int16_t, int32_t, int64_t
#include <memory> // 对于 shared_ptr
#include <string> // 对于 string, to_string
#include <type_traits> // 对于 invoke_result_t, is_same_v, underlying_type_t
 
namespace xgboost {
 
struct CUDAContext;
 
// 符号名称
struct DeviceSym {
 static auto constexpr CPU() { return "cpu"; }
 static auto constexpr CUDA() { return "cuda"; }
 static auto constexpr SyclDefault() { return "sycl"; }
 static auto constexpr SyclCPU() { return "sycl:cpu"; }
 static auto constexpr SyclGPU() { return "sycl:gpu"; }
};
 
struct DeviceOrd {
 // 代表 CPU 设备 ID 的常量。
 static bst_d_ordinal_t constexpr CPUOrdinal() { return -1; }
 static bst_d_ordinal_t constexpr InvalidOrdinal() { return -2; }
 
 enum Type : std::int16_t { kCPU = 0, kCUDA = 1,
 kSyclDefault = 2, kSyclCPU = 3, kSyclGPU = 4} device{kCPU};
 // CUDA 或 Sycl 设备序号。
 bst_d_ordinal_t ordinal{CPUOrdinal()};
 
 [[nodiscard]] bool IsCUDA() const { return device == kCUDA; }
 [[nodiscard]] bool IsCPU() const { return device == kCPU; }
 [[nodiscard]] bool IsSyclDefault() const { return device == kSyclDefault; }
 [[nodiscard]] bool IsSyclCPU() const { return device == kSyclCPU; }
 [[nodiscard]] bool IsSyclGPU() const { return device == kSyclGPU; }
 [[nodiscard]] bool IsSycl() const { return (IsSyclDefault() ||
 IsSyclCPU() ||
 IsSyclGPU()); }
 
 constexpr DeviceOrd() = default;
 constexpr DeviceOrd(Type type, bst_d_ordinal_t ord) : device{type}, ordinal{ord} {}
 
 constexpr DeviceOrd(DeviceOrd const& that) = default;
 constexpr DeviceOrd& operator=(DeviceOrd const& that) = default;
 constexpr DeviceOrd(DeviceOrd&& that) = default;
 constexpr DeviceOrd& operator=(DeviceOrd&& that) = default;
 
 [[nodiscard]] constexpr static auto CPU() { return DeviceOrd{kCPU, CPUOrdinal()}; }
 [[nodiscard]] static constexpr auto CUDA(bst_d_ordinal_t ordinal) {
 return DeviceOrd{kCUDA, ordinal};
  }
 [[nodiscard]] constexpr static auto SyclDefault(bst_d_ordinal_t ordinal = -1) {
 return DeviceOrd{kSyclDefault, ordinal};
  }
 [[nodiscard]] constexpr static auto SyclCPU(bst_d_ordinal_t ordinal = -1) {
 return DeviceOrd{kSyclCPU, ordinal};
  }
 
 [[nodiscard]] constexpr static auto SyclGPU(bst_d_ordinal_t ordinal = -1) {
 return DeviceOrd{kSyclGPU, ordinal};
  }
 
 [[nodiscard]] bool operator==(DeviceOrd const& that) const {
 return device == that.device && ordinal == that.ordinal;
  }
 [[nodiscard]] bool operator!=(DeviceOrd const& that) const { return !(*this == that); }
 [[nodiscard]] std::string Name() const {
 switch (device) {
 case DeviceOrd::kCPU:
 return DeviceSym::CPU();
 case DeviceOrd::kCUDA:
 return DeviceSym::CUDA() + (':' + std::to_string(ordinal));
 case DeviceOrd::kSyclDefault:
 return DeviceSym::SyclDefault() + (':' + std::to_string(ordinal));
 case DeviceOrd::kSyclCPU:
 return DeviceSym::SyclCPU() + (':' + std::to_string(ordinal));
 case DeviceOrd::kSyclGPU:
 return DeviceSym::SyclGPU() + (':' + std::to_string(ordinal));
 default: {
 LOG(FATAL) << "未知设备。";
 return "";
      }
    }
  }
};
 
static_assert(sizeof(DeviceOrd) == sizeof(std::int32_t));
 
std::ostream& operator<<(std::ostream& os, DeviceOrd ord);
 
struct Context : public XGBoostParameter<Context> {
 private
 // 用户接口参数，用于设备序号
 std::string device{DeviceSym::CPU()}; // NOLINT
 // 用户设置的设备序号
 DeviceOrd device_{DeviceOrd::CPU()};
 
 public
 static std::int64_t constexpr kDefaultSeed = 0;
 
 public
 Context();
 
 void Init(Args const& kwargs);
 
 template <typename Container>
 Args UpdateAllowUnknown(Container const& kwargs) {
 auto args = XGBoostParameter<Context>::UpdateAllowUnknown(kwargs);
 this->SetDeviceOrdinal(kwargs);
 return args;
  }
 
 // 如果启用 OpenMP，则使用的线程数。如果等于 0，则使用系统默认值。
 std::int32_t nthread{0}; // NOLINT
 // 存储的随机种子
 std::int64_t seed{kDefaultSeed};
 // 是否每次迭代都为 PRNG 播种
 bool seed_per_iteration{false};
 // 当 gpu_id 无效时失败
 bool fail_on_invalid_gpu_id{false};
 bool validate_parameters{false};
 
 [[nodiscard]] std::int32_t Threads() const;
 [[nodiscard]] bool IsCPU() const { return Device().IsCPU(); }
 [[nodiscard]] bool IsCUDA() const { return Device().IsCUDA(); }
 [[nodiscard]] bool IsSyclDefault() const { return Device().IsSyclDefault(); }
 [[nodiscard]] bool IsSyclCPU() const { return Device().IsSyclCPU(); }
 [[nodiscard]] bool IsSyclGPU() const { return Device().IsSyclGPU(); }
 [[nodiscard]] bool IsSycl() const { return IsSyclDefault()
 || IsSyclCPU()
 || IsSyclGPU(); }
 
 [[nodiscard]] DeviceOrd Device() const { return device_; }
 
 [[nodiscard]] DeviceOrd DeviceFP64() const;
 
 [[nodiscard]] bst_d_ordinal_t Ordinal() const { return Device().ordinal; }
 [[nodiscard]] std::string DeviceName() const { return Device().Name(); }
 [[nodiscard]] CUDAContext const* CUDACtx() const;
 
 [[nodiscard]] Context MakeCUDA(bst_d_ordinal_t ordinal = 0) const {
 Context ctx = *this;
 return ctx.SetDevice(DeviceOrd::CUDA(ordinal));
  }
 [[nodiscard]] Context MakeCPU() const {
 Context ctx = *this;
 return ctx.SetDevice(DeviceOrd::CPU());
  }
 
 template <typename CPUFn, typename CUDAFn>
 decltype(auto) DispatchDevice(CPUFn&& cpu_fn, CUDAFn&& cuda_fn) const {
 static_assert(std::is_same_v<std::invoke_result_t<CPUFn>, std::invoke_result_t<CUDAFn>>);
 switch (this->Device().device) {
 case DeviceOrd::kCPU:
 return cpu_fn();
 case DeviceOrd::kCUDA:
 return cuda_fn();
 default:
 // 不要使用设备名称，因为这很可能是一个内部错误，该名称
 // 将无效。
 if (this->Device().IsSycl()) {
 LOG(WARNING) << "所请求的功能尚无 SYCL 特定的实现。 "
 << "正在使用 CPU 实现";
 return cpu_fn();
 } else {
 LOG(FATAL) << "未知设备类型:"
 << static_cast<std::underlying_type_t<DeviceOrd::Type>>(this->Device().device);
 break;
        }
    }
 return std::invoke_result_t<CPUFn>();
  }
 
 template <typename CPUFn, typename CUDAFn, typename SYCLFn>
 decltype(auto) DispatchDevice(CPUFn&& cpu_fn, CUDAFn&& cuda_fn, SYCLFn&& sycl_fn) const {
 static_assert(std::is_same_v<std::invoke_result_t<CPUFn>, std::invoke_result_t<SYCLFn>>);
 if (this->Device().IsSycl()) {
 return sycl_fn();
 } else {
 return DispatchDevice(cpu_fn, cuda_fn);
    }
  }
 
 // 声明参数
 DMLC_DECLARE_PARAMETER(Context) {
 DMLC_DECLARE_FIELD(seed)
 .set_default(kDefaultSeed)
 .describe("训练期间的随机数种子。");
 DMLC_DECLARE_ALIAS(seed, random_state);
 DMLC_DECLARE_FIELD(seed_per_iteration)
 .set_default(false)
 .describe("通过迭代器编号确定性地为 PRNG 播种。");
 DMLC_DECLARE_FIELD(device).set_default(DeviceSym::CPU()).describe("设备序号。");
 DMLC_DECLARE_FIELD(nthread).set_default(0).describe("要使用的线程数。");
 DMLC_DECLARE_ALIAS(nthread, n_jobs);
 DMLC_DECLARE_FIELD(fail_on_invalid_gpu_id)
 .set_default(false)
 .describe("当 gpu_id 无效时，报错失败。");
 DMLC_DECLARE_FIELD(validate_parameters)
 .set_default(false)
 .describe("启用检查参数是否使用。");
  }
 
 private
 void SetDeviceOrdinal(Args const& kwargs);
 Context& SetDevice(DeviceOrd d) {
 this->device = (this->device_ = d).Name();
 return *this;
  }
 
 // 可变用于惰性 CUDA 上下文初始化。这避免了在加载时初始化 CUDA。
 // 使用 shared_ptr 而不是 unique_ptr，因为使用 unique_ptr 很难在尝试隐藏 CUDA 代码时定义 p_impl。
 // p_impl，同时尝试从主机编译器中隐藏 CUDA 代码。
 mutable std::shared_ptr<CUDAContext> cuctx_;
 // CFS CPU 限制的缓存值。（在容器化环境中使用）
 std::int32_t cfs_cpu_count_; // NOLINT
};
} // namespace xgboost
 
#endif // XGBOOST_CONTEXT_H_