Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

finish c_api.cu

Browse Source
This commit is contained in:
amdsc21 2023-03-10 05:12:51 +01:00
parent a76ccff390
commit bb6adda8a3
1 changed files with 8 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
src/c_api/c_api.cu
View File

@ -50,13 +50,21 @@ 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.
#if defined(XGBOOST_USE_CUDA)
cudaGetDevice(&device_id_);
#elif defined(XGBOOST_USE_HIP)
hipGetDevice(&device_id_);
#endif
}
void XGBoostAPIGuard::RestoreGPUAttribute() {
// Not calling `safe_cuda` to avoid unnecessary exception handling overhead.
// If errors, do nothing, assuming running on CPU only machine.
#if defined(XGBOOST_USE_CUDA)
cudaSetDevice(device_id_);
#elif defined(XGBOOST_USE_HIP)
hipSetDevice(device_id_);
#endif
}
} // namespace xgboost