From 91c19d9431e8424336ef65302412a4b231e07cf4 Mon Sep 17 00:00:00 2001 From: Chuck Atkins Date: Wed, 4 Sep 2024 15:09:25 -0400 Subject: [PATCH] C API: Cache CUDA availability so it's not re-evaluated every API call --- src/c_api/c_api.cu | 31 +++++++++++++++++++++++++------ 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/src/c_api/c_api.cu b/src/c_api/c_api.cu index 47868f466473..7b164efba4a5 100644 --- a/src/c_api/c_api.cu +++ b/src/c_api/c_api.cu @@ -19,7 +19,14 @@ #include #endif +#include + namespace xgboost { + +namespace { + std::atomic_bool MIGHT_HAVE_CUDA(true); +} + void XGBBuildInfoDevice(Json *p_info) { auto &info = *p_info; @@ -59,15 +66,27 @@ void XGBBuildInfoDevice(Json *p_info) { } void XGBoostAPIGuard::SetGPUAttribute() { - // Not calling `safe_cuda` to avoid unnecessary exception handling overhead. - // If errors, do nothing, assuming running on CPU only machine. - cudaGetDevice(&device_id_); + try { + if (MIGHT_HAVE_CUDA.load(std::memory_order_relaxed)) { + // Not calling `safe_cuda` to avoid unnecessary exception handling overhead. + // If errors, do nothing, assuming running on CPU only machine. + cudaGetDevice(&device_id_); + } + } catch (dmlc::Error const&) { + MIGHT_HAVE_CUDA.store(false, std::memory_order_relaxed); + } } void XGBoostAPIGuard::RestoreGPUAttribute() { - // Not calling `safe_cuda` to avoid unnecessary exception handling overhead. - // If errors, do nothing, assuming running on CPU only machine. - cudaSetDevice(device_id_); + try { + if (MIGHT_HAVE_CUDA.load(std::memory_order_relaxed)) { + // Not calling `safe_cuda` to avoid unnecessary exception handling overhead. + // If errors, do nothing, assuming running on CPU only machine. + cudaSetDevice(device_id_); + } + } catch (dmlc::Error const&) { + MIGHT_HAVE_CUDA.store(false, std::memory_order_relaxed); + } } void CopyGradientFromCUDAArrays(Context const *ctx, ArrayInterface<2, false> const &grad,