1- /* !
2- * Copyright 2014-2022 by Contributors
1+ /* *
2+ * Copyright 2014-2023, XGBoost Contributors
33 * \file context.h
44 */
55#ifndef XGBOOST_CONTEXT_H_
66#define XGBOOST_CONTEXT_H_
77
8- #include < xgboost/logging.h>
9- #include < xgboost/parameter.h>
8+ #include < xgboost/base.h> // for bst_d_ordinal_t
9+ #include < xgboost/logging.h> // for CHECK_GE
10+ #include < xgboost/parameter.h> // for XGBoostParameter
1011
11- #include < memory> // std::shared_ptr
12- #include < string>
12+ #include < cstdint> // for int16_t, int32_t, int64_t
13+ #include < memory> // for shared_ptr
14+ #include < string> // for string, to_string
1315
1416namespace xgboost {
1517
1618struct CUDAContext ;
1719
20+ /* *
21+ * @brief A type for device ordinal. The type is packed into 32-bit for efficient use in
22+ * viewing types like `linalg::TensorView`.
23+ */
24+ struct DeviceOrd {
25+ enum Type : std::int16_t { kCPU = 0 , kCUDA = 1 } device{kCPU };
26+ // CUDA device ordinal.
27+ bst_d_ordinal_t ordinal{-1 };
28+
29+ [[nodiscard]] bool IsCUDA () const { return device == kCUDA ; }
30+ [[nodiscard]] bool IsCPU () const { return device == kCPU ; }
31+
32+ DeviceOrd () = default ;
33+ constexpr DeviceOrd (Type type, bst_d_ordinal_t ord) : device{type}, ordinal{ord} {}
34+
35+ DeviceOrd (DeviceOrd const & that) = default ;
36+ DeviceOrd& operator =(DeviceOrd const & that) = default ;
37+ DeviceOrd (DeviceOrd&& that) = default ;
38+ DeviceOrd& operator =(DeviceOrd&& that) = default ;
39+
40+ /* *
41+ * @brief Constructor for CPU.
42+ */
43+ [[nodiscard]] constexpr static auto CPU () { return DeviceOrd{kCPU , -1 }; }
44+ /* *
45+ * @brief Constructor for CUDA device.
46+ *
47+ * @param ordinal CUDA device ordinal.
48+ */
49+ [[nodiscard]] static auto CUDA (bst_d_ordinal_t ordinal) { return DeviceOrd{kCUDA , ordinal}; }
50+
51+ [[nodiscard]] bool operator ==(DeviceOrd const & that) const {
52+ return device == that.device && ordinal == that.ordinal ;
53+ }
54+ [[nodiscard]] bool operator !=(DeviceOrd const & that) const { return !(*this == that); }
55+ /* *
56+ * @brief Get a string representation of the device and the ordinal.
57+ */
58+ [[nodiscard]] std::string Name () const {
59+ switch (device) {
60+ case DeviceOrd::kCPU :
61+ return " CPU"
62+ case DeviceOrd::kCUDA :
63+ return " CUDA:" std::to_string (ordinal);
64+ default : {
65+ LOG (FATAL) << " Unknown device."
66+ return " "
67+ }
68+ }
69+ }
70+ };
71+
72+ static_assert (sizeof (DeviceOrd) == sizeof (std::int32_t ));
73+
74+ /* *
75+ * @brief Runtime context for XGBoost. Contains information like threads and device.
76+ */
1877struct Context : public XGBoostParameter <Context> {
1978 public:
2079 // Constant representing the device ID of CPU.
@@ -36,29 +95,59 @@ struct Context : public XGBoostParameter<Context> {
3695 // fail when gpu_id is invalid
3796 bool fail_on_invalid_gpu_id{false };
3897 bool validate_parameters{false };
39-
40- /* !
41- * \brief Configure the parameter `gpu_id'.
98+ /* *
99+ * @brief Configure the parameter `gpu_id'.
42100 *
43- * \param require_gpu Whether GPU is explicitly required from user.
101+ * @param require_gpu Whether GPU is explicitly required by the user through other
102+ * configurations.
44103 */
45104 void ConfigureGpuId (bool require_gpu);
46- /* !
47- * Return automatically chosen threads.
105+ /* *
106+ * @brief Returns the automatically chosen number of threads based on the `nthread`
107+ * parameter and the system settting.
48108 */
49- std::int32_t Threads () const ;
50-
51- bool IsCPU () const { return gpu_id == kCpuId ; }
52- bool IsCUDA () const { return !IsCPU (); }
53-
54- CUDAContext const * CUDACtx () const ;
55- // Make a CUDA context based on the current context.
56- Context MakeCUDA (std::int32_t device = 0 ) const {
109+ [[nodiscard]] std::int32_t Threads () const ;
110+ /* *
111+ * @brief Is XGBoost running on CPU?
112+ */
113+ [[nodiscard]] bool IsCPU () const { return gpu_id == kCpuId ; }
114+ /* *
115+ * @brief Is XGBoost running on a CUDA device?
116+ */
117+ [[nodiscard]] bool IsCUDA () const { return !IsCPU (); }
118+ /* *
119+ * @brief Get the current device and ordinal.
120+ */
121+ [[nodiscard]] DeviceOrd Device () const {
122+ return IsCPU () ? DeviceOrd::CPU () : DeviceOrd::CUDA (static_cast <bst_d_ordinal_t >(gpu_id));
123+ }
124+ /* *
125+ * @brief Get the CUDA device ordinal. -1 if XGBoost is running on CPU.
126+ */
127+ [[nodiscard]] bst_d_ordinal_t Ordinal () const { return this ->gpu_id ; }
128+ /* *
129+ * @brief Name of the current device.
130+ */
131+ [[nodiscard]] std::string DeviceName () const { return Device ().Name (); }
132+ /* *
133+ * @brief Get a CUDA device context for allocator and stream.
134+ */
135+ [[nodiscard]] CUDAContext const * CUDACtx () const ;
136+ /* *
137+ * @brief Make a CUDA context based on the current context.
138+ *
139+ * @param ordinal The CUDA device ordinal.
140+ */
141+ [[nodiscard]] Context MakeCUDA (std::int32_t ordinal = 0 ) const {
57142 Context ctx = *this ;
58- ctx.gpu_id = device;
143+ CHECK_GE (ordinal, 0 );
144+ ctx.gpu_id = ordinal;
59145 return ctx;
60146 }
61- Context MakeCPU () const {
147+ /* *
148+ * @brief Make a CPU context based on the current context.
149+ */
150+ [[nodiscard]] Context MakeCPU () const {
62151 Context ctx = *this ;
63152 ctx.gpu_id = kCpuId ;
64153 return ctx;
@@ -87,9 +176,9 @@ struct Context : public XGBoostParameter<Context> {
87176 }
88177
89178 private:
90- // mutable for lazy initialization for cuda context to avoid initializing CUDA at load.
91- // shared_ptr is used instead of unique_ptr as with unique_ptr it's difficult to define p_impl
92- // while trying to hide CUDA code from host compiler.
179+ // mutable for lazy cuda context initialization. This avoids initializing CUDA at load.
180+ // shared_ptr is used instead of unique_ptr as with unique_ptr it's difficult to define
181+ // p_impl while trying to hide CUDA code from the host compiler.
93182 mutable std::shared_ptr<CUDAContext> cuctx_;
94183 // cached value for CFS CPU limit. (used in containerized env)
95184 std::int32_t cfs_cpu_count_; // NOLINT
0 commit comments