GPU Mapping (#4326)

For perlmutter and frontier, if there are multiple devices 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 machine 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>
@@ -207,9 +208,6 @@ Device::Initialize ()
         device_id = 0;
     }
     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 (ParallelDescriptor::NProcsPerNode() == gpu_device_count) {
             device_id = ParallelDescriptor::MyRankInNode();
         } else if (ParallelDescriptor::NProcsPerProcessor() == gpu_device_count) {
@@ -219,6 +217,30 @@ 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;
+            if (amrex::Verbose()) {
+                amrex::Print() << "Multiple GPUs are visible to each MPI rank. Fixing GPU assignment for Perlmuuter according to heuristics.\n";
+            }
+        } 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];
+                if (amrex::Verbose()) {
+                    amrex::Print() << "Multiple GPUs are visible to each MPI rank. Fixing GPU assignment for Frontier according to heuristics.\n";
+                }
+            }
+        } 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.");
+            }
+        }
+    }
+
     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 ();
 
 }