#pragma once
#include <cstddef>
#include <cstdint>
#if defined(XGBOOST_USE_NVTX)
#include <nvtx3/nvtx3.hpp>
#endif
namespace xgboost::curt {
std::int32_t AllVisibleGPUs();
std::int32_t CurrentDevice(bool raise = true);
bool SupportsPageableMem();
bool SupportsAts();
void CheckComputeCapability();
void SetDevice(std::int32_t device);
[[nodiscard]] std::size_t TotalMemory();
void RtVersion(std::int32_t* major, std::int32_t* minor);
void DrVersion(std::int32_t* major, std::int32_t* minor);
struct NvtxDomain {
static constexpr char const *name{"libxgboost"}; };
#if defined(XGBOOST_USE_NVTX)
using NvtxScopedRange = ::nvtx3::scoped_range_in<NvtxDomain>;
using NvtxEventAttr = ::nvtx3::event_attributes;
using NvtxRgb = ::nvtx3::rgb;
#else
class NvtxScopedRange {
public:
template <typename... Args>
explicit NvtxScopedRange(Args &&...) {}
};
class NvtxEventAttr {
public:
template <typename... Args>
explicit NvtxEventAttr(Args &&...) {}
};
class NvtxRgb {
public:
template <typename... Args>
explicit NvtxRgb(Args &&...) {}
};
#endif }
#if defined(XGBOOST_USE_NVTX)
#define xgboost_NVTX_FN_RANGE() NVTX3_FUNC_RANGE_IN(::xgboost::curt::NvtxDomain)
#else
#define xgboost_NVTX_FN_RANGE()
#endif