nan_check.hpp

Go to the documentation of this file.

#ifndef NAN_CHECK_HPP
#define NAN_CHECK_HPP
 
#include "space_settings.hpp"
 
template <typename T>
inline void nan_check_by_dim(const T& view, std::string name) {
    constexpr int Rank = T::rank;
    using value_type = typename T::value_type;
 
    // Allocate one bool array per dimension
    size_t max_ndim = 0;
    for (int d = 0; d < Rank; d++) max_ndim = std::max(max_ndim, view.extent(d));
    View<bool**,CLayout,typename T::device_type> dim_nan_flags(NoInit("dim_nan_flags"), Rank, max_ndim);
    Kokkos::deep_copy(dim_nan_flags, false);
 
    // Compute flattened range for simplicity
    const size_t N = view.size();
 
    if constexpr (Rank == 1) {
        Kokkos::parallel_for("nan_check_by_dim", Kokkos::RangePolicy<typename T::execution_space>(0, N), KOKKOS_LAMBDA(size_t idx_flat) {
            int i = idx_flat;
            if (!isfinite(view(i))) {
                dim_nan_flags(0,i) = true;
            }
        });
    } else if constexpr (Rank == 2) {
        Kokkos::parallel_for("nan_check_by_dim", Kokkos::RangePolicy<typename T::execution_space>(0, N), KOKKOS_LAMBDA(size_t idx_flat) {
            int e1 = view.extent(1);
            int i = idx_flat / e1;
            int j = idx_flat % e1;
            if (!isfinite(view(i, j))) {
                dim_nan_flags(0,i) = true;
                dim_nan_flags(1,j) = true;
            }
        });
    } else if constexpr (Rank == 3) {
        Kokkos::parallel_for("nan_check_by_dim", Kokkos::RangePolicy<typename T::execution_space>(0, N), KOKKOS_LAMBDA(size_t idx_flat) {
            int e1 = view.extent(1), e2 = view.extent(2);
            int i = idx_flat / (e1 * e2);
            int j = (idx_flat / e2) % e1;
            int k = idx_flat % e2;
            if (!isfinite(view(i, j, k))) {
                dim_nan_flags(0,i) = true;
                dim_nan_flags(1,j) = true;
                dim_nan_flags(2,k) = true;
            }
        });
    } else if constexpr (Rank == 4) {
        Kokkos::parallel_for("nan_check_by_dim", Kokkos::RangePolicy<typename T::execution_space>(0, N), KOKKOS_LAMBDA(size_t idx_flat) {
            int e1 = view.extent(1), e2 = view.extent(2), e3 = view.extent(3);
            int i = idx_flat / (e1 * e2 * e3);
            int j = (idx_flat / (e2 * e3)) % e1;
            int k = (idx_flat / e3) % e2;
            int l = idx_flat % e3;
            if (!isfinite(view(i, j, k, l))) {
                dim_nan_flags(0,i) = true;
                dim_nan_flags(1,j) = true;
                dim_nan_flags(2,k) = true;
                dim_nan_flags(3,l) = true;
            }
        });
    }
    Kokkos::fence();
 
    // Copy flags back to host and print results
    auto flags_host = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), dim_nan_flags);
    for (int d = 0; d < Rank; d++) {
        std::cout << "Dimension " << d << ": Found NaNs in: {";
        int range_start = -1;
        bool inside_range = false;
        for (size_t i = 0; i < view.extent(d); i++) {
            if(flags_host(d,i)) {
                if(!inside_range) range_start = i;
                inside_range = true;
            }else{
                if(inside_range){
                    // Just left range
                    inside_range = false;
 
                    if(i>range_start+1){
                        // Print range if there are consecutive values
                        std::cout << range_start << "-" << i-1 << ", ";
                    }else{
                        // Otherwise just print the value
                        std::cout << range_start << ", ";
                    }
                }
            }
        }
        if(inside_range){
            std::cout << range_start << "-" << view.extent(d)-1;
        }
        std::cout << "}\n";
    }
}
 
 
 
// Sums all elements of a view of any shape
template<typename T>
inline int nan_check(const T& view, std::string name){
    // Cast to 1D with an unmanaged view for a generic element-wise operation
    View<typename T::value_type*,CLayout, typename T::device_type, Kokkos::MemoryTraits<Kokkos::Unmanaged>> view_1D(view.data(), view.size());
    
    int n_nan = 0;
    Kokkos::parallel_reduce("nan_check", Kokkos::RangePolicy<typename T::execution_space>( 0,view.size() ), KOKKOS_LAMBDA( const int i, int& n_nan_l){
        n_nan_l += (!isfinite(view_1D(i)) ? 1 : 0);
    }, n_nan);
    Kokkos::fence();
 
    // No NaNs found
    if(n_nan==0) return 0;
 
    printf("Warning: %d NaNs found in %s!!\n", n_nan, name.c_str());
 
    // NaNs found
    // This should never happen. Do some analysis without worrying about performance.
    nan_check_by_dim(view, name);
 
    return 1;
}
 
#endif