xgboost/include/xgboost/context.h
Jiaming Yuan c6a8754c62
Define CUDA Context. (#8604)
We will transition to non-default and non-blocking CUDA stream.
2022-12-20 15:15:07 +08:00

88 lines
2.6 KiB
C++

/*!
* Copyright 2014-2022 by Contributors
* \file context.h
*/
#ifndef XGBOOST_CONTEXT_H_
#define XGBOOST_CONTEXT_H_
#include <xgboost/logging.h>
#include <xgboost/parameter.h>
#include <memory> // std::shared_ptr
#include <string>
namespace xgboost {
struct CUDAContext;
struct Context : public XGBoostParameter<Context> {
public:
// Constant representing the device ID of CPU.
static std::int32_t constexpr kCpuId = -1;
static std::int64_t constexpr kDefaultSeed = 0;
public:
Context();
// stored random seed
std::int64_t seed{kDefaultSeed};
// whether seed the PRNG each iteration
bool seed_per_iteration{false};
// number of threads to use if OpenMP is enabled
// if equals 0, use system default
std::int32_t nthread{0};
// primary device, -1 means no gpu.
std::int32_t gpu_id{kCpuId};
// fail when gpu_id is invalid
bool fail_on_invalid_gpu_id{false};
bool validate_parameters{false};
/*!
* \brief Configure the parameter `gpu_id'.
*
* \param require_gpu Whether GPU is explicitly required from user.
*/
void ConfigureGpuId(bool require_gpu);
/*!
* Return automatically chosen threads.
*/
std::int32_t Threads() const;
bool IsCPU() const { return gpu_id == kCpuId; }
bool IsCUDA() const { return !IsCPU(); }
CUDAContext const* CUDACtx() const;
// declare parameters
DMLC_DECLARE_PARAMETER(Context) {
DMLC_DECLARE_FIELD(seed)
.set_default(kDefaultSeed)
.describe("Random number seed during training.");
DMLC_DECLARE_ALIAS(seed, random_state);
DMLC_DECLARE_FIELD(seed_per_iteration)
.set_default(false)
.describe("Seed PRNG determnisticly via iterator number.");
DMLC_DECLARE_FIELD(nthread).set_default(0).describe("Number of threads to use.");
DMLC_DECLARE_ALIAS(nthread, n_jobs);
DMLC_DECLARE_FIELD(gpu_id).set_default(-1).set_lower_bound(-1).describe(
"The primary GPU device ordinal.");
DMLC_DECLARE_FIELD(fail_on_invalid_gpu_id)
.set_default(false)
.describe("Fail with error when gpu_id is invalid.");
DMLC_DECLARE_FIELD(validate_parameters)
.set_default(false)
.describe("Enable checking whether parameters are used or not.");
}
private:
// mutable for lazy initialization for cuda context to avoid initializing CUDA at load.
// shared_ptr is used instead of unique_ptr as with unique_ptr it's difficult to define p_impl
// while trying to hide CUDA code from host compiler.
mutable std::shared_ptr<CUDAContext> cuctx_;
// cached value for CFS CPU limit. (used in containerized env)
std::int32_t cfs_cpu_count_; // NOLINT
};
} // namespace xgboost
#endif // XGBOOST_CONTEXT_H_