Tpetra_Details_EquilibrationInfo.hpp

Go to the documentation of this file.

// @HEADER
// *****************************************************************************
//          Tpetra: Templated Linear Algebra Services Package
//
// Copyright 2008 NTESS and the Tpetra contributors.
// SPDX-License-Identifier: BSD-3-Clause
// *****************************************************************************
// @HEADER

#ifndef TPETRA_DETAILS_EQUILIBRATIONINFO_HPP
#define TPETRA_DETAILS_EQUILIBRATIONINFO_HPP


#include "TpetraCore_config.h"
#if KOKKOS_VERSION >= 40799
#include "KokkosKernels_ArithTraits.hpp"
#else
#include "Kokkos_ArithTraits.hpp"
#endif
#include "Kokkos_Core.hpp"

namespace Tpetra {
namespace Details {

template <class ScalarType, class DeviceType>
struct EquilibrationInfo {
#if KOKKOS_VERSION >= 40799
  using val_type = typename KokkosKernels::ArithTraits<ScalarType>::val_type;
#else
  using val_type = typename Kokkos::ArithTraits<ScalarType>::val_type;
#endif
#if KOKKOS_VERSION >= 40799
  using mag_type = typename KokkosKernels::ArithTraits<val_type>::mag_type;
#else
  using mag_type = typename Kokkos::ArithTraits<val_type>::mag_type;
#endif
  using device_type      = typename DeviceType::device_type;
  using host_device_type = typename Kokkos::View<mag_type*, device_type>::host_mirror_type::device_type;
  using HostMirror       = EquilibrationInfo<val_type, host_device_type>;

  EquilibrationInfo()
    : foundInf(false)
    , foundNan(false)
    , foundZeroDiag(false)
    , foundZeroRowNorm(false) {}

  EquilibrationInfo(const std::size_t lclNumRows,
                    const std::size_t lclNumCols,
                    const bool assumeSymmetric_)
    : rowNorms(Kokkos::View<mag_type*, device_type>("rowNorms", lclNumRows))
    , rowDiagonalEntries(Kokkos::View<val_type*, device_type>("rowDiagonalEntries", lclNumRows))
    , colNorms(Kokkos::View<mag_type*, device_type>("colNorms", lclNumCols))
    , colDiagonalEntries(Kokkos::View<val_type*, device_type>("colDiagonalEntries",
                                                              assumeSymmetric_ ? std::size_t(0) : lclNumCols))
    , rowScaledColNorms(Kokkos::View<mag_type*, device_type>("rowScaledColNorms",
                                                             assumeSymmetric_ ? std::size_t(0) : lclNumCols))
    , assumeSymmetric(assumeSymmetric_)
    , foundInf(false)
    , foundNan(false)
    , foundZeroDiag(false)
    , foundZeroRowNorm(false) {}

  EquilibrationInfo(const Kokkos::View<mag_type*, device_type>& rowNorms_,
                    const Kokkos::View<val_type*, device_type>& rowDiagonalEntries_,
                    const Kokkos::View<mag_type*, device_type>& colNorms_,
                    const Kokkos::View<val_type*, device_type>& colDiagonalEntries_,
                    const Kokkos::View<mag_type*, device_type>& rowScaledColNorms_,
                    const bool assumeSymmetric_,
                    const bool foundInf_,
                    const bool foundNan_,
                    const bool foundZeroDiag_,
                    const bool foundZeroRowNorm_)
    : rowNorms(rowNorms_)
    , rowDiagonalEntries(rowDiagonalEntries_)
    , colNorms(colNorms_)
    , colDiagonalEntries(colDiagonalEntries_)
    , rowScaledColNorms(rowScaledColNorms_)
    , assumeSymmetric(assumeSymmetric_)
    , foundInf(foundInf_)
    , foundNan(foundNan_)
    , foundZeroDiag(foundZeroDiag_)
    , foundZeroRowNorm(foundZeroRowNorm_) {}

  template <class SrcDeviceType>
  void
  assign(const EquilibrationInfo<ScalarType, SrcDeviceType>& src) {
    using execution_space = typename device_type::execution_space;
    // DEEP_COPY REVIEW - DEVICE-TO-DEVICE
    Kokkos::deep_copy(execution_space(), rowNorms, src.rowNorms);
    // DEEP_COPY REVIEW - DEVICE-TO-DEVICE
    Kokkos::deep_copy(execution_space(), rowDiagonalEntries, src.rowDiagonalEntries);
    // DEEP_COPY REVIEW - DEVICE-TO-DEVICE
    Kokkos::deep_copy(execution_space(), colNorms, src.colNorms);
    if (src.colDiagonalEntries.extent(0) == 0) {
      colDiagonalEntries =
          Kokkos::View<val_type*, device_type>("colDiagonalEntries", 0);
    } else {
      // DEEP_COPY REVIEW - DEVICE-TO-DEVICE
      Kokkos::deep_copy(execution_space(), colDiagonalEntries, src.colDiagonalEntries);
    }
    if (src.rowScaledColNorms.extent(0) == 0) {
      rowScaledColNorms =
          Kokkos::View<mag_type*, device_type>("rowScaledColNorms", 0);
    } else {
      // DEEP_COPY REVIEW - DEVICE-TO-DEVICE
      Kokkos::deep_copy(execution_space(), rowScaledColNorms, src.rowScaledColNorms);
    }

    assumeSymmetric  = src.assumeSymmetric;
    foundInf         = src.foundInf;
    foundNan         = src.foundNan;
    foundZeroDiag    = src.foundZeroDiag;
    foundZeroRowNorm = src.foundZeroRowNorm;
  }

  typename EquilibrationInfo<val_type, device_type>::HostMirror
  createMirrorView() {
    auto rowNorms_h           = Kokkos::create_mirror_view(rowNorms);
    auto rowDiagonalEntries_h = Kokkos::create_mirror_view(rowDiagonalEntries);
    auto colNorms_h           = Kokkos::create_mirror_view(colNorms);
    auto colDiagonalEntries_h = Kokkos::create_mirror_view(colDiagonalEntries);
    auto rowScaledColNorms_h  = Kokkos::create_mirror_view(rowScaledColNorms);

    return HostMirror{rowNorms_h, rowDiagonalEntries_h, colNorms_h,
                      colDiagonalEntries_h, rowScaledColNorms_h, assumeSymmetric,
                      foundInf, foundNan, foundZeroDiag, foundZeroRowNorm};
  }

  // We call a row a "diagonally dominant row" if the absolute value
  // of the diagonal entry is >= the sum of the absolute values of the
  // off-diagonal entries.  The row norm is the sum of those two
  // things, so this means diagAbsVal >= rowNorm - diagAbsVal.  Ditto
  // for a column.

  Kokkos::View<mag_type*, device_type> rowNorms;

  Kokkos::View<val_type*, device_type> rowDiagonalEntries;

  Kokkos::View<mag_type*, device_type> colNorms;

  Kokkos::View<val_type*, device_type> colDiagonalEntries;

  Kokkos::View<mag_type*, device_type> rowScaledColNorms;

  bool assumeSymmetric;

  bool foundInf;

  bool foundNan;

  bool foundZeroDiag;

  bool foundZeroRowNorm;
};

}  // namespace Details
}  // namespace Tpetra

#endif  // TPETRA_DETAILS_EQUILIBRATIONINFO_HPP