17 #include "alsutils/config.hpp" 18 #ifdef ALSVINN_HAVE_CUDA 19 #define ENV_LOCAL_RANK "OMPI_COMM_WORLD_LOCAL_RANK" 20 #include <cuda_runtime.h> 41 if (deviceCount > 1) {
50 const char* rankAsString = std::getenv(ENV_LOCAL_RANK);
54 int mpiRank = std::atoi(rankAsString);
62 const int device = mpiRank % deviceCount;
67 int currentDevice = -1;
71 if (currentDevice != device) {
73 currentDevice <<
")");
77 }
catch (std::runtime_error& e) {
79 "(Ignore this if you are not running with CUDA). Failed setting CUDA GPU: " <<
void setCudaDevice()
Definition: set_cuda_device.hpp:89
#define INFO
Definition: log.hpp:28
#define CUDA_SAFE_CALL_SILENT(x)
Does the same as CUDA_SAFE_CALL, but doesn't print an error message.
Definition: cuda_safe_call.hpp:31
Various utilities for mpi and cuda.
Definition: Factory.hpp:3
#define WARNING
Definition: log.hpp:30
#define ALSVINN_LOG(severity, message)
Definition: log.hpp:36