GPU Mapping

For perlmutter and frontier, if there are multiple device available, we will
try to map GPUs to the closest core.

For an FFT test on perlmutter using 256 nodes, the correct mapping reduced
the run time from 0.172 to 0.127. Note that you can achieve the similar
effect with `srun ... bash -c "export
CUDA_VISIBLE_DEVICES=\$((3-SLURM_LOCALID)); ..."` by manually limiting the
number of visible devices. But in this commit, we are trying to do this
automatically for the user. Also note that MPI appears to crash with
gpu-bind=closest on perlmutter. So we need to use gpu-bind=none.

For frontier, you could use gpu-bind=closest. But if your use gpu-bind=none,
this commit will try to do the correct mapping for you.

In this commit, we also removed the old machin stuff and added new code for
machine detection.

Src/Base/AMReX.cpp

@@ -506,6 +506,8 @@ amrex::Initialize (int& argc, char**& argv, bool build_parm_parse,
     }
 #endif
 
+    Machine::Initialize();
+
 #ifdef AMREX_USE_GPU
     // Initialize after ParmParse so that we can read inputs.
     Gpu::Device::Initialize();
@@ -670,8 +672,6 @@ amrex::Initialize (int& argc, char**& argv, bool build_parm_parse,
     BL_PROFILE_INITPARAMS();
 #endif // ifndef BL_AMRPROF
 
-    machine::Initialize();
-
 #ifdef AMREX_USE_HYPRE
     if (init_hypre) {
         HYPRE_Init();

Src/Base/AMReX_GpuDevice.cpp

@@ -1,9 +1,10 @@
 
 #include <AMReX_GpuDevice.H>
+#include <AMReX_GpuLaunch.H>
+#include <AMReX_Machine.H>
 #include <AMReX_ParallelDescriptor.H>
 #include <AMReX_ParmParse.H>
 #include <AMReX_Print.H>
-#include <AMReX_GpuLaunch.H>
 
 #ifdef AMREX_USE_HYPRE
 #  include <_hypre_utilities.h>
@@ -208,7 +209,11 @@ Device::Initialize ()
     }
     else {
         if (amrex::Verbose() && ParallelDescriptor::IOProcessor()) {
-            amrex::Warning("Multiple GPUs are visible to each MPI rank. This is usually not an issue. But this may lead to incorrect or suboptimal rank-to-GPU mapping.");
+            if ((Machine::name() != "nersc.perlmutter") &&
+                (Machine::name() != "olcf.frontier"))
+            {
+                amrex::Warning("Multiple GPUs are visible to each MPI rank. This is usually not an issue. But this may lead to incorrect or suboptimal rank-to-GPU mapping.");
+            }
         }
         if (ParallelDescriptor::NProcsPerNode() == gpu_device_count) {
             device_id = ParallelDescriptor::MyRankInNode();
@@ -219,6 +224,20 @@ Device::Initialize ()
         }
     }
 
+    if (gpu_device_count > 1){
+        if (Machine::name() == "nersc.perlmutter") {
+            // The CPU/GPU mapping on perlmutter has the reverse order.
+            device_id = gpu_device_count - device_id - 1;
+        } else if (Machine::name() == "olcf.frontier") {
+            // The CPU/GPU mapping on fronter is documented at
+            // https://docs.olcf.ornl.gov/systems/frontier_user_guide.html
+            if (gpu_device_count == 8) {
+                constexpr std::array<int,8> gpu_order = {4,5,2,3,6,7,0,1};
+                device_id = gpu_order[device_id];
+            }
+        }
+    }
+
     AMREX_HIP_OR_CUDA(AMREX_HIP_SAFE_CALL (hipSetDevice(device_id));,
                       AMREX_CUDA_SAFE_CALL(cudaSetDevice(device_id)); );
 

Src/Base/AMReX_Machine.H

@@ -2,20 +2,15 @@
 #define AMREX_MACHINE_H
 #include <AMReX_Config.H>
 
-#include <AMReX_Vector.H>
+#include <string>
 
-namespace amrex::machine {
+namespace amrex::Machine {
 
 void Initialize (); //!< called in amrex::Initialize()
 
-#ifdef AMREX_USE_MPI
 void Finalize ();
-/**
-* find the best topologically close neighborhood of ranks
-* returns a vector of global or local rank IDs based on flag_local_ranks
-*/
-Vector<int> find_best_nbh (int rank_n, bool flag_local_ranks = false);
-#endif
+
+std::string const& name ();
 
 }