MooseInit.C File Reference

Go to the source code of this file.

Functions
void	SigHandler (int signum)
void	RegisterSigHandler ()
int	getLocalRank ()
void	initializeCuda ()

Function Documentation

getLocalRank()

int getLocalRank

(

)

Definition at line 53 of file MooseInit.C.

{
  static constexpr std::array local_rank_vars{
      "OMPI_COMM_WORLD_LOCAL_RANK", // Open MPI
      "MV2_COMM_WORLD_LOCAL_RANK",  // MVAPICH2
      "MPI_LOCALRANKID",            // Intel MPI / some MPICH environments
      "PMI_LOCAL_RANK",             // PMI-based launchers
      "PALS_LOCAL_RANKID",          // HPE Cray PALS
      "SLURM_LOCALID",              // Slurm
      "FLUX_TASK_LOCAL_ID"          // Flux
  };
 
  for (const char * variable : local_rank_vars)
  {
    const char * value = std::getenv(variable);
    if (!value)
      continue;
 
    int rank;
    const char * end = value + std::char_traits<char>::length(value);
    const auto [ptr, error] = std::from_chars(value, end, rank);
 
    if (error != std::errc{} || ptr != end || rank < 0)
      throw std::runtime_error(std::string("Invalid local rank in ") + variable + ": " + value);
 
    return rank;
  }
 
  // Not launched through a recognized MPI/process-manager environment.
  return 0;
}

Referenced by initializeCuda().

initializeCuda()

void initializeCuda

(

)

NOTE: Some GPU-aware MPI implementations initialize accelerator-related resources during MPI_Init() and may associate those resources with the device that is current at that time. Therefore, select the GPU device before MPI initialization.

Open MPI explicitly recommends selecting the accelerator before MPI_Init(), using OMPI_COMM_WORLD_LOCAL_RANK for per-node device assignment. In particular, we have observed that failing to do so can prevent initialization of the smcuda BTL (the transfer component for fast shared-memory GPU communication), causing GPU communication to fall back to a much slower PCI path instead of using NVLink, with over 100x performance difference in some cases. MVAPICH2 also supports selecting the CUDA device before MPI_Init() using MV2_COMM_WORLD_LOCAL_RANK. MPICH does not appear to require early device selection, but initializing the GPU runtime before MPI_Init() is expected to be harmless.

The mapping used here (local_rank % num_gpus) is intended to establish the device used by Kokkos. It is not yet known whether this mapping is appropriate for other GPU packages supported by MOOSE.

TODO: Add HIP and SYCL device selection if needed.

Definition at line 87 of file MooseInit.C.

{
  int local_rank = getLocalRank();
 
  int num_gpus = 0;
  cudaGetDeviceCount(&num_gpus);
 
  if (num_gpus)
  {
    int device = local_rank % num_gpus;
    cudaSetDevice(device);
  }
}

Referenced by MooseDeviceInit::MooseDeviceInit().

RegisterSigHandler()

void RegisterSigHandler

(

)

Definition at line 47 of file MooseInit.C.

{
  signal(SIGUSR1, SigHandler);
}

Referenced by MooseInit::MooseInit().

SigHandler()

void SigHandler

(

int

signum

)

Definition at line 40 of file MooseInit.C.

{
  Moose::interrupt_signal_number = signum;
  return;
}

Referenced by RegisterSigHandler().