ExaTrkX/src/printCudaMemInfo.hpp

0001 // This file is part of the ACTS project.
0002 //
0003 // Copyright (C) 2016 CERN for the benefit of the ACTS project
0004 //
0005 // This Source Code Form is subject to the terms of the Mozilla Public
0006 // License, v. 2.0. If a copy of the MPL was not distributed with this
0007 // file, You can obtain one at https://mozilla.org/MPL/2.0/.
0008
0009 #pragma once
0010
0011 #include <Acts/Utilities/Logger.hpp>
0012
0013 #ifndef ACTS_EXATRKX_CPUONLY
0014 #include <cuda_runtime_api.h>
0015 #endif
0016
0017 #include <cstdint>
0018
0019 #include <torch/torch.h>
0020
0021 namespace {
0022
0023 inline void printCudaMemInfo(const Acts::Logger& logger) {
0024 #ifndef ACTS_EXATRKX_CPUONLY
0025   if (torch::cuda::is_available() && logger.level() == Acts::Logging::VERBOSE) {
0026     constexpr float kb = 1024;
0027     constexpr float mb = kb * kb;
0028
0029     int device;
0030     std::size_t free, total;
0031     cudaMemGetInfo(&free, &total);
0032     cudaGetDevice(&device);
0033
0034     ACTS_VERBOSE("Current CUDA device: " << device);
0035     ACTS_VERBOSE("Memory (used / total) [in MB]: " << (total - free) / mb
0036                                                    << " / " << total / mb);
0037   } else {
0038     ACTS_VERBOSE("No memory info, CUDA disabled");
0039   }
0040 #else
0041   ACTS_VERBOSE("No memory info, CUDA disabled");
0042 #endif
0043 }
0044
0045 }  // namespace