MueLu: Guard against potential NaN/infs in strength-of-connection algorithms

Signed-off-by: malphil <[email protected]>

Author: malphil

packages/muelu/src/Graph/MatrixTransformation/MueLu_CutDrop.hpp

@@ -679,6 +679,8 @@ class CutDropFunctor {
     size_t keepStart = 0;
     size_t dropStart = nnz;
     // find index where dropping starts
+
+    bool nonFiniteValueEncountered = false;
     for (size_t i = 1; i < nnz; ++i) {
       auto const& x = row_permutation(i - 1);
       auto const& y = row_permutation(i);
@@ -693,6 +695,17 @@ class CutDropFunctor {
           dropStart = i;
         }
       }
+
+      nonFiniteValueEncountered |= !is_finite_type_safe(x_aij);
+      nonFiniteValueEncountered |= !is_finite_type_safe(y_aij);
+    }
+
+    if (nonFiniteValueEncountered) {
+      const char* message =
+          "Error encountered in MueLu::CutDrop::CutDropFunctor::operator():\n"
+          "Non-finite values encountered in strength-of-connection measure.\n"
+          "A potential fix is to enable rebalancing and/or perform an initial rebalance.\n";
+      Kokkos::abort(message);
     }
 
     // drop everything to the right of where values stop passing threshold

packages/muelu/src/Graph/MatrixTransformation/MueLu_DistanceLaplacianDropping.hpp

@@ -697,6 +697,7 @@ class DropFunctor {
     }
 #endif
 
+    bool nonFiniteValueEncountered = false;
     for (local_ordinal_type k = 0; k < row.length; ++k) {
       auto clid = row.colidx(k);
 
@@ -711,11 +712,14 @@ class DropFunctor {
         auto aiiajj = ATS::magnitude(diag(rlid)) * ATS::magnitude(diag(clid));  // |a_ii|*|a_jj|
         auto aij2   = ATS::magnitude(val) * ATS::magnitude(val);                // |a_ij|^2
 
+        nonFiniteValueEncountered |= !is_finite_type_safe(aiiajj) || !is_finite_type_safe(aij2);
+
         results(offset + k) = Kokkos::max((aij2 <= eps * eps * aiiajj) ? DROP : KEEP,
                                           results(offset + k));
       } else if constexpr (measure == Misc::SignedRugeStuebenMeasure) {
-        auto neg_aij        = -ATS::real(val);
-        auto max_neg_aik    = eps * ATS::real(diag(rlid));
+        auto neg_aij     = -ATS::real(val);
+        auto max_neg_aik = eps * ATS::real(diag(rlid));
+        nonFiniteValueEncountered |= !is_finite_type_safe(neg_aij) || !is_finite_type_safe(max_neg_aik);
         results(offset + k) = Kokkos::max((neg_aij < max_neg_aik) ? DROP : KEEP,
                                           results(offset + k));
       } else if constexpr (measure == Misc::SignedSmoothedAggregationMeasure) {
@@ -725,10 +729,19 @@ class DropFunctor {
         // + |a_ij|^2, if a_ij < 0, - |a_ij|^2 if a_ij >=0
         if (!is_nonpositive)
           aij2 = -aij2;
+        nonFiniteValueEncountered |= !is_finite_type_safe(aiiajj) || !is_finite_type_safe(aij2);
         results(offset + k) = Kokkos::max((aij2 <= eps * eps * aiiajj) ? DROP : KEEP,
                                           results(offset + k));
       }
     }
+
+    if (nonFiniteValueEncountered) {
+      const char* message =
+          "Error encountered in MueLu::DistanceLaplacian::DropFunctor::operator():\n"
+          "Non-finite values encountered in strength-of-connection measure.\n"
+          "A potential fix is to enable rebalancing and/or perform an initial rebalance.\n";
+      Kokkos::abort(message);
+    }
   }
 };
 
@@ -832,6 +845,7 @@ class VectorDropFunctor {
     }
 #endif
 
+    bool nonFiniteValueEncountered = false;
     for (local_ordinal_type k = 0; k < row.length; ++k) {
       auto clid  = row.colidx(k);
       auto bclid = ghosted_point_to_block(clid);
@@ -847,11 +861,14 @@ class VectorDropFunctor {
         auto aiiajj = ATS::magnitude(diag(brlid)) * ATS::magnitude(diag(bclid));  // |a_ii|*|a_jj|
         auto aij2   = ATS::magnitude(val) * ATS::magnitude(val);                  // |a_ij|^2
 
+        nonFiniteValueEncountered |= !is_finite_type_safe(aiiajj) || !is_finite_type_safe(aij2);
+
         results(offset + k) = Kokkos::max((aij2 <= eps * eps * aiiajj) ? DROP : KEEP,
                                           results(offset + k));
       } else if constexpr (measure == Misc::SignedRugeStuebenMeasure) {
-        auto neg_aij        = -ATS::real(val);
-        auto max_neg_aik    = eps * ATS::real(diag(brlid));
+        auto neg_aij     = -ATS::real(val);
+        auto max_neg_aik = eps * ATS::real(diag(brlid));
+        nonFiniteValueEncountered |= !is_finite_type_safe(neg_aij) || !is_finite_type_safe(max_neg_aik);
         results(offset + k) = Kokkos::max((neg_aij < max_neg_aik) ? DROP : KEEP,
                                           results(offset + k));
       } else if constexpr (measure == Misc::SignedSmoothedAggregationMeasure) {
@@ -861,10 +878,19 @@ class VectorDropFunctor {
         // + |a_ij|^2, if a_ij < 0, - |a_ij|^2 if a_ij >=0
         if (!is_nonpositive)
           aij2 = -aij2;
+        nonFiniteValueEncountered |= !is_finite_type_safe(aiiajj) || !is_finite_type_safe(aij2);
         results(offset + k) = Kokkos::max((aij2 <= eps * eps * aiiajj) ? DROP : KEEP,
                                           results(offset + k));
       }
     }
+
+    if (nonFiniteValueEncountered) {
+      const char* message =
+          "Error encountered in MueLu::DistanceLaplacian::VectorDropFunctor::operator():\n"
+          "Non-finite values encountered in strength-of-connection measure.\n"
+          "A potential fix is to enable rebalancing and/or perform an initial rebalance.\n";
+      Kokkos::abort(message);
+    }
   }
 };
 

packages/muelu/src/Graph/MatrixTransformation/MueLu_DroppingCommon.hpp

@@ -36,6 +36,15 @@ enum DecisionType : char {
   BOUNDARY  = 3   // entry is a boundary
 };
 
+template <typename T>
+KOKKOS_INLINE_FUNCTION bool is_finite_type_safe(T value) {
+  if constexpr (std::is_floating_point_v<T>) {
+    return Kokkos::isfinite(value);
+  } else {
+    return true;
+  }
+}
+
 namespace Misc {
 
 template <class local_ordinal_type>