compute balance on cpus, but show warning only on gpus

Author: glesur

src/timeIntegrator.cpp

@@ -197,34 +197,33 @@ double TimeIntegrator::ComputeBalance() {
   // Check MPI imbalance
     double imbalance = 0;
     #ifdef WITH_MPI
-      // only do this on GPUs
-      #if defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP) || defined(KOKKOS_ENABLE_SYCL)
-        const double allowedImbalance = 20.0;
-        std::vector<double> computeLogPerCore(idfx::psize);
-        MPI_Gather(&computeLastLog, 1, MPI_DOUBLE, computeLogPerCore.data(), 1, MPI_DOUBLE, 0,
-                    MPI_COMM_WORLD);
-        computeLastLog = 0; // reset timer for all cores
-        if(idfx::prank==0) {
-          // Compute the average, the min and the max
-          double computeMin = computeLogPerCore[0];
-          double computeMax = computeLogPerCore[0];
-          double computeMean = 0;
-
-          for(int i = 0 ; i < idfx::psize ; i++) {
-            computeMean += computeLogPerCore[i];
-            if(computeLogPerCore[i]>computeMax) {
-              computeMax = computeLogPerCore[i];
-            }
-            if(computeLogPerCore[i]<computeMin) {
-              computeMin = computeLogPerCore[i];
-            }
+      const double allowedImbalance = 20.0;
+      std::vector<double> computeLogPerCore(idfx::psize);
+      MPI_Gather(&computeLastLog, 1, MPI_DOUBLE, computeLogPerCore.data(), 1, MPI_DOUBLE, 0,
+                  MPI_COMM_WORLD);
+      computeLastLog = 0; // reset timer for all cores
+      if(idfx::prank==0) {
+        // Compute the average, the min and the max
+        double computeMin = computeLogPerCore[0];
+        double computeMax = computeLogPerCore[0];
+        double computeMean = 0;
+
+        for(int i = 0 ; i < idfx::psize ; i++) {
+          computeMean += computeLogPerCore[i];
+          if(computeLogPerCore[i]>computeMax) {
+            computeMax = computeLogPerCore[i];
           }
-          computeMean /= idfx::psize;
-          imbalance = (computeMax-computeMin)/computeMean*100;
-
+          if(computeLogPerCore[i]<computeMin) {
+            computeMin = computeLogPerCore[i];
+          }
+        }
+        computeMean /= idfx::psize;
+        imbalance = (computeMax-computeMin)/computeMean*100;
+        // only show warnings if the imbalance is higher than 20%on GPUs
+        #if defined(KOKKOS_ENABLE_CUDA) || defined(KOKKOS_ENABLE_HIP) || defined(KOKKOS_ENABLE_SYCL)
           if(imbalance>allowedImbalance ) {
             idfx::cout << "-------------------------------------------------------------"
-                       << std::endl;
+                        << std::endl;
             idfx::cout << "Warning: MPI imbalance found in this run " << std::endl;
             idfx::cout << std::fixed;
             for(int i = 0 ; i < idfx::psize ; i++) {
@@ -238,12 +237,12 @@ double TimeIntegrator::ComputeBalance() {
               }
             }
             idfx::cout << "You should probably check these nodes are running properly."
-                       << std::endl;
+                        << std::endl;
             idfx::cout << "-------------------------------------------------------------"
-                       << std::endl;
+                        << std::endl;
           }
-        }
-      #endif
+        #endif
+      }
     #endif
     return(imbalance);
 }