Tpetra_Apply_Helpers.hpp

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

#ifndef TPETRA_APPLY_HELPERS_HPP
#define TPETRA_APPLY_HELPERS_HPP
#include <type_traits>
#include "Teuchos_ScalarTraits.hpp"
#include "Tpetra_CrsMatrix.hpp"
#include "Tpetra_Details_Profiling.hpp"

namespace Tpetra {


template <class MatrixArray, class MultiVectorArray>
void batchedApply(const MatrixArray &Matrices,
                  const typename std::remove_pointer<typename MultiVectorArray::value_type>::type &X,
                  MultiVectorArray &Y,
                  typename std::remove_pointer<typename MatrixArray::value_type>::type::scalar_type alpha = Teuchos::ScalarTraits<typename std::remove_pointer<typename MatrixArray::value_type>::type::scalar_type>::one(),
                  typename std::remove_pointer<typename MatrixArray::value_type>::type::scalar_type beta  = Teuchos::ScalarTraits<typename std::remove_pointer<typename MatrixArray::value_type>::type::scalar_type>::zero(),
                  Teuchos::RCP<Teuchos::ParameterList> params                                             = Teuchos::null) {
  Tpetra::Details::ProfilingRegion regionTotal("Tpetra::batchedApply");

  using size_type   = typename MatrixArray::size_type;
  using matrix_type = typename std::remove_pointer<typename MatrixArray::value_type>::type;
  using map_type    = typename matrix_type::map_type;
  using import_type = typename matrix_type::import_type;
  using export_type = typename matrix_type::export_type;
  using MV          = typename matrix_type::MV;
  using scalar_type = typename matrix_type::scalar_type;

  using Teuchos::RCP;
  using Teuchos::rcp_const_cast;

  const scalar_type ONE  = Teuchos::ScalarTraits<scalar_type>::one();
  const scalar_type ZERO = Teuchos::ScalarTraits<scalar_type>::zero();

  size_type N = Matrices.size();
  if (N == 0) return;
  int numRanks = X.getMap()->getComm()->getSize();

  // If X is aliased to any Y but the last one, throw
  for (size_type i = 0; i < N - 1; i++) {
    TEUCHOS_TEST_FOR_EXCEPTION(&X == Y[i], std::runtime_error, "Tpetra::batchedApply(): X cannot be aliased to any Y except the final one.");
  }

  /* Checks for whether the reduced communication path can be used */
  RCP<const map_type> compare_colMap = Matrices[0]->getColMap();
  RCP<const import_type> importer    = Matrices[0]->getGraph()->getImporter();

  bool can_batch, check_maps;
  if (params.is_null() || !params->isParameter("can batch")) {
    can_batch  = (importer.is_null() || N == 1) ? false : true;
    check_maps = true;
  } else {
    can_batch  = (!params->get<bool>("can batch") || importer.is_null() || N == 1) ? false : true;
    check_maps = false;
  }
  // We can't batch with replicated Y's
  if (numRanks > 1) {
    for (size_type i = 0; i < N && can_batch; i++) {
      if (!Y[i]->isDistributed())
        can_batch = false;
    }
  }

  // Do the domain/column maps all match?
  for (size_type i = 1; i < N && check_maps && can_batch; i++) {
    if (!Matrices[i]->getColMap()->isSameAs(*compare_colMap)) {
      can_batch = false;
    }
  }

  if (can_batch) {
    /* Batching path: Guarantees an existing importer and N>1 */

    // Special case for alpha = 0
    if (alpha == ZERO) {
      if (beta == ZERO) {
        for (size_type i = 0; i < N; i++) Y[i]->putScalar(ZERO);
      } else if (beta != ONE) {
        for (size_type i = 0; i < N; i++) Y[i]->scale(beta);
      }
      if (!params.is_null()) params->set("can batch", true);
      return;
    }

    const bool Y_is_overwritten = (beta == ZERO);

    // Start by importing X to Matrices[0]'s temporary
    RCP<const MV> X_colMap;
    {
      Tpetra::Details::ProfilingRegion regionImport("Tpetra::batchedApply: Import");
      RCP<MV> X_colMapNonConst = Matrices[0]->getColumnMapMultiVector(X);

      // Import from the domain Map MV to the column Map MV.
      X_colMapNonConst->doImport(X, *importer, INSERT);
      X_colMap = rcp_const_cast<const MV>(X_colMapNonConst);
    }

    for (size_type i = 0; i < N; i++) {
      RCP<const export_type> exporter = Matrices[i]->getGraph()->getExporter();

      // Temporary MV for doExport (if needed),
      RCP<MV> Y_rowMap = Matrices[i]->getRowMapMultiVector(*Y[i]);
      if (!exporter.is_null()) {
        Matrices[i]->localApply(*X_colMap, *Y_rowMap, Teuchos::NO_TRANS, alpha, ZERO);
        {
          Tpetra::Details::ProfilingRegion regionExport("Tpetra::batchedApply: Export");
          if (Y_is_overwritten) {
            Y[i]->putScalar(ZERO);
          } else {
            Y[i]->scale(beta);
          }
          Y[i]->doExport(*Y_rowMap, *exporter, ADD_ASSIGN);
        }
      } else {  // Don't do an Export: row Map and range Map are the same.
        // Check for aliasing
        if (!Y[i]->isConstantStride() || X_colMap.getRawPtr() == Y[i]) {
          Y_rowMap = Matrices[i]->getRowMapMultiVector(*Y[i], true);
          if (beta != ZERO) {
            Tpetra::deep_copy(*Y_rowMap, *Y[i]);
          }

          Matrices[i]->localApply(*X_colMap, *Y_rowMap, Teuchos::NO_TRANS, alpha, beta);
          Tpetra::deep_copy(*Y[i], *Y_rowMap);
        } else {
          Matrices[i]->localApply(*X_colMap, *Y[i], Teuchos::NO_TRANS, alpha, beta);
        }
      }
    }
    if (!params.is_null()) params->set("can batch", true);
  } else {
    /* Non-batching path */
    for (size_type i = 0; i < N; i++) {
      Matrices[i]->apply(X, *Y[i], Teuchos::NO_TRANS, alpha, beta);
    }
    if (!params.is_null()) params->set("can batch", false);
  }
}

}  // namespace Tpetra

#endif  // TPETRA_APPLY_HELPERS_HPP