Tpetra_Details_normImpl.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_NORMIMPL_HPP
#define TPETRA_DETAILS_NORMIMPL_HPP


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

#ifndef DOXYGEN_SHOULD_SKIP_THIS
namespace Teuchos {
template <class T>
class ArrayView;  // forward declaration
template <class OrdinalType>
class Comm;  // forward declaration
}  // namespace Teuchos
#endif  // DOXYGEN_SHOULD_SKIP_THIS

// Declarations start here

namespace Tpetra {
namespace Details {

enum EWhichNorm {
  NORM_ONE,  //<! Use the one-norm
  NORM_TWO,  //<! Use the two-norm
  NORM_INF   //<! Use the infinity-norm
};

template <class ValueType,
          class ArrayLayout,
          class DeviceType,
          class MagnitudeType>
void normImpl(MagnitudeType norms[],
              const Kokkos::View<const ValueType**, ArrayLayout, DeviceType>& X,
              const EWhichNorm whichNorm,
              const Teuchos::ArrayView<const size_t>& whichVecs,
              const bool isConstantStride,
              const bool isDistributed,
              const Teuchos::Comm<int>* comm);

}  // namespace Details
}  // namespace Tpetra

// Definitions start here

namespace Tpetra {
namespace Details {
namespace Impl {

template <class RV, class XMV>
void lclNormImpl(const RV& normsOut,
                 const XMV& X,
                 const size_t numVecs,
                 const Teuchos::ArrayView<const size_t>& whichVecs,
                 const bool constantStride,
                 const EWhichNorm whichNorm) {
  using Kokkos::ALL;
  using Kokkos::subview;
  using mag_type = typename RV::non_const_value_type;

  static_assert(static_cast<int>(RV::rank) == 1,
                "Tpetra::MultiVector::lclNormImpl: "
                "The first argument normsOut must have rank 1.");
  static_assert(Kokkos::is_view<XMV>::value,
                "Tpetra::MultiVector::lclNormImpl: "
                "The second argument X is not a Kokkos::View.");
  static_assert(static_cast<int>(XMV::rank) == 2,
                "Tpetra::MultiVector::lclNormImpl: "
                "The second argument X must have rank 2.");

  const size_t lclNumRows = static_cast<size_t>(X.extent(0));
  TEUCHOS_TEST_FOR_EXCEPTION(lclNumRows != 0 && constantStride &&
                                 static_cast<size_t>(X.extent(1)) != numVecs,
                             std::logic_error, "Constant Stride X's dimensions are " << X.extent(0) << " x " << X.extent(1) << ", which differ from the local dimensions " << lclNumRows << " x " << numVecs << ".  Please report this bug to "
                                                                                                                                                                                                                "the Tpetra developers.");
  TEUCHOS_TEST_FOR_EXCEPTION(lclNumRows != 0 && !constantStride &&
                                 static_cast<size_t>(X.extent(1)) < numVecs,
                             std::logic_error, "Strided X's dimensions are " << X.extent(0) << " x " << X.extent(1) << ", which are incompatible with the local dimensions " << lclNumRows << " x " << numVecs << ".  Please report this bug to "
                                                                                                                                                                                                                  "the Tpetra developers.");

  if (lclNumRows == 0) {
#if KOKKOS_VERSION >= 40799
    const mag_type zeroMag = KokkosKernels::ArithTraits<mag_type>::zero();
#else
    const mag_type zeroMag = Kokkos::ArithTraits<mag_type>::zero();
#endif
    // DEEP_COPY REVIEW - VALUE-TO-DEVICE
    using execution_space = typename RV::execution_space;
    Kokkos::deep_copy(execution_space(), normsOut, zeroMag);
  } else {  // lclNumRows != 0
    if (constantStride) {
      if (whichNorm == NORM_INF) {
        KokkosBlas::nrminf(normsOut, X);
      } else if (whichNorm == NORM_ONE) {
        KokkosBlas::nrm1(normsOut, X);
      } else if (whichNorm == NORM_TWO) {
        KokkosBlas::nrm2_squared(normsOut, X);
      } else {
        TEUCHOS_TEST_FOR_EXCEPTION(true, std::logic_error, "Should never get here!");
      }
    } else {  // not constant stride
      // NOTE (mfh 15 Jul 2014, 11 Apr 2019) This does a kernel launch
      // for every column.  It might be better to have a kernel that
      // does the work all at once.  On the other hand, we don't
      // prioritize performance of MultiVector views of noncontiguous
      // columns.
      for (size_t k = 0; k < numVecs; ++k) {
        const size_t X_col = constantStride ? k : whichVecs[k];
        if (whichNorm == NORM_INF) {
          KokkosBlas::nrminf(subview(normsOut, k),
                             subview(X, ALL(), X_col));
        } else if (whichNorm == NORM_ONE) {
          KokkosBlas::nrm1(subview(normsOut, k),
                           subview(X, ALL(), X_col));
        } else if (whichNorm == NORM_TWO) {
          KokkosBlas::nrm2_squared(subview(normsOut, k),
                                   subview(X, ALL(), X_col));
        } else {
          TEUCHOS_TEST_FOR_EXCEPTION(true, std::logic_error, "Should never get here!");
        }
      }  // for each column
    }    // constantStride
  }      // lclNumRows != 0
}

// Kokkos::parallel_for functor for applying square root to each
// entry of a 1-D Kokkos::View.
template <class ViewType>
class SquareRootFunctor {
 public:
  typedef typename ViewType::execution_space execution_space;
  typedef typename ViewType::size_type size_type;

  SquareRootFunctor(const ViewType& theView)
    : theView_(theView) {}

  KOKKOS_INLINE_FUNCTION void
  operator()(const size_type& i) const {
    typedef typename ViewType::non_const_value_type value_type;
#if KOKKOS_VERSION >= 40799
    typedef KokkosKernels::ArithTraits<value_type> KAT;
#else
    typedef Kokkos::ArithTraits<value_type> KAT;
#endif
    theView_(i) = KAT::sqrt(theView_(i));
  }

 private:
  ViewType theView_;
};

template <class RV>
void gblNormImpl(const RV& normsOut,
                 const Teuchos::Comm<int>* const comm,
                 const bool distributed,
                 const EWhichNorm whichNorm) {
  using Teuchos::REDUCE_MAX;
  using Teuchos::REDUCE_SUM;
  using Teuchos::reduceAll;
  typedef typename RV::non_const_value_type mag_type;

  const size_t numVecs = normsOut.extent(0);

  // If the MultiVector is distributed over multiple processes, do
  // the distributed (interprocess) part of the norm.  We assume
  // that the MPI implementation can read from and write to device
  // memory.
  //
  // replaceMap() may have removed some processes.  Those processes
  // have a null Map.  They must not participate in any collective
  // operations.  We ask first whether the Map is null, because
  // isDistributed() defers that question to the Map.  We still
  // compute and return local norms for processes not participating
  // in collective operations; those probably don't make any sense,
  // but it doesn't hurt to do them, since it's illegal to call
  // norm*() on those processes anyway.
  if (distributed && comm != nullptr) {
    // The calling process only participates in the collective if
    // both the Map and its Comm on that process are nonnull.

    const int nv               = static_cast<int>(numVecs);
    const bool commIsInterComm = ::Tpetra::Details::isInterComm(*comm);

    if (commIsInterComm) {
      RV lclNorms(Kokkos::ViewAllocateWithoutInitializing("MV::normImpl lcl"), numVecs);
      // DEEP_COPY REVIEW - DEVICE-TO-DEVICE
      using execution_space = typename RV::execution_space;
      Kokkos::deep_copy(execution_space(), lclNorms, normsOut);
      const mag_type* const lclSum = lclNorms.data();
      mag_type* const gblSum       = normsOut.data();

      if (whichNorm == NORM_INF) {
        reduceAll<int, mag_type>(*comm, REDUCE_MAX, nv, lclSum, gblSum);
      } else {
        reduceAll<int, mag_type>(*comm, REDUCE_SUM, nv, lclSum, gblSum);
      }
    } else {
      mag_type* const gblSum = normsOut.data();
      if (whichNorm == NORM_INF) {
        reduceAll<int, mag_type>(*comm, REDUCE_MAX, nv, gblSum, gblSum);
      } else {
        reduceAll<int, mag_type>(*comm, REDUCE_SUM, nv, gblSum, gblSum);
      }
    }
  }

  if (whichNorm == NORM_TWO) {
    // Replace the norm-squared results with their square roots in
    // place, to get the final output.  If the device memory and
    // the host memory are the same, it probably doesn't pay to
    // launch a parallel kernel for that, since there isn't enough
    // parallelism for the typical MultiVector case.
    const bool inHostMemory =
        std::is_same<typename RV::memory_space,
                     typename RV::host_mirror_space::memory_space>::value;
    if (inHostMemory) {
      for (size_t j = 0; j < numVecs; ++j) {
#if KOKKOS_VERSION >= 40799
        normsOut(j) = KokkosKernels::ArithTraits<mag_type>::sqrt(normsOut(j));
#else
        normsOut(j) = Kokkos::ArithTraits<mag_type>::sqrt(normsOut(j));
#endif
      }
    } else {
      // There's not as much parallelism now, but that's OK.  The
      // point of doing parallel dispatch here is to keep the norm
      // results on the device, thus avoiding a copy to the host
      // and back again.
      SquareRootFunctor<RV> f(normsOut);
      typedef typename RV::execution_space execution_space;
      typedef Kokkos::RangePolicy<execution_space, size_t> range_type;
      Kokkos::parallel_for(range_type(0, numVecs), f);
    }
  }
}

}  // namespace Impl

template <class ValueType,
          class ArrayLayout,
          class DeviceType,
          class MagnitudeType>
void normImpl(MagnitudeType norms[],
              const Kokkos::View<const ValueType**, ArrayLayout, DeviceType>& X,
              const EWhichNorm whichNorm,
              const Teuchos::ArrayView<const size_t>& whichVecs,
              const bool isConstantStride,
              const bool isDistributed,
              const Teuchos::Comm<int>* comm) {
  using execution_space = typename DeviceType::execution_space;
  using RV              = Kokkos::View<MagnitudeType*, Kokkos::HostSpace>;
  // using XMV = Kokkos::View<const ValueType**, ArrayLayout, DeviceType>;
  // using pair_type = std::pair<size_t, size_t>;

  const size_t numVecs = isConstantStride ? static_cast<size_t>(X.extent(1)) : static_cast<size_t>(whichVecs.size());
  if (numVecs == 0) {
    return;  // nothing to do
  }
  RV normsOut(norms, numVecs);

  Impl::lclNormImpl(normsOut, X, numVecs, whichVecs,
                    isConstantStride, whichNorm);

  // lbv 03/15/23: the data from the local norm calculation
  // better really be available before communication happens
  // so fencing to make sure the local computations have
  // completed on device. We might want to make this an
  // execution space fence down the road?
  execution_space exec_space_instance = execution_space();
  exec_space_instance.fence();

  Impl::gblNormImpl(normsOut, comm, isDistributed, whichNorm);
}

}  // namespace Details
}  // namespace Tpetra

#endif  // TPETRA_DETAILS_NORMIMPL_HPP