40 #ifndef TPETRA_APPLY_HELPERS_HPP
41 #define TPETRA_APPLY_HELPERS_HPP
42 #include <type_traits>
43 #include "Teuchos_ScalarTraits.hpp"
44 #include "Tpetra_CrsMatrix.hpp"
73 template <
class MatrixArray,
class MultiVectorArray>
75 const typename std::remove_pointer<typename MultiVectorArray::value_type>::type &X,
77 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(),
78 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(),
79 Teuchos::RCP<Teuchos::ParameterList> params = Teuchos::null) {
82 using size_type =
typename MatrixArray::size_type;
83 using matrix_type =
typename std::remove_pointer<typename MatrixArray::value_type>::type;
84 using map_type =
typename matrix_type::map_type;
85 using import_type =
typename matrix_type::import_type;
86 using export_type =
typename matrix_type::export_type;
87 using MV =
typename matrix_type::MV;
88 using scalar_type =
typename matrix_type::scalar_type;
91 using Teuchos::rcp_const_cast;
93 const scalar_type ONE = Teuchos::ScalarTraits<scalar_type>::one();
94 const scalar_type
ZERO = Teuchos::ScalarTraits<scalar_type>::zero();
96 size_type N = Matrices.size();
98 int numRanks = X.getMap()->getComm()->getSize();
101 for(size_type i=0; i<N-1; i++) {
102 TEUCHOS_TEST_FOR_EXCEPTION( &X == Y[i], std::runtime_error,
"Tpetra::batchedApply(): X cannot be aliased to any Y except the final one.");
106 RCP<const map_type> compare_colMap = Matrices[0]->getColMap();
107 RCP<const import_type> importer = Matrices[0]->getGraph()->getImporter();
110 bool can_batch, check_maps;
111 if(params.is_null() || !params->isParameter(
"can batch")) {
112 can_batch = (importer.is_null() || N==1) ?
false :
true;
116 can_batch = (!params->get<
bool>(
"can batch") || importer.is_null() || N==1) ?
false :
true;
121 for(size_type i=0; i<N && can_batch; i++) {
122 if(!Y[i]->isDistributed())
128 for(size_type i=1; i<N && check_maps && can_batch; i++) {
129 if(!Matrices[i]->getColMap()->isSameAs(*compare_colMap)) {
140 for(size_type i=0; i<N; i++) Y[i]->putScalar(ZERO);
141 }
else if (beta != ONE) {
142 for(size_type i=0; i<N; i++) Y[i]->scale(beta);
144 if(!params.is_null()) params->set(
"can batch",
true);
148 const bool Y_is_overwritten = (beta ==
ZERO);
151 RCP<const MV> X_colMap;
154 RCP<MV> X_colMapNonConst = Matrices[0]->getColumnMapMultiVector(X);
157 X_colMapNonConst->doImport(X, *importer,
INSERT);
158 X_colMap = rcp_const_cast<
const MV>(X_colMapNonConst);
161 for(size_type i=0; i<N; i++) {
162 RCP<const export_type> exporter = Matrices[i]->getGraph()->getExporter();
165 RCP<MV> Y_rowMap = Matrices[i]->getRowMapMultiVector(*Y[i]);
166 if (!exporter.is_null()) {
167 Matrices[i]->localApply(*X_colMap, *Y_rowMap, Teuchos::NO_TRANS, alpha, ZERO);
170 if (Y_is_overwritten) {
171 Y[i]->putScalar(ZERO);
176 Y[i]->doExport(*Y_rowMap, *exporter,
ADD_ASSIGN);
181 if (! Y[i]->isConstantStride() || X_colMap.getRawPtr() == Y[i]) {
182 Y_rowMap = Matrices[i]->getRowMapMultiVector(*Y[i],
true);
187 Matrices[i]->localApply(*X_colMap, *Y_rowMap, Teuchos::NO_TRANS, alpha, beta);
191 Matrices[i]->localApply(*X_colMap, *Y[i], Teuchos::NO_TRANS, alpha, beta);
195 if(!params.is_null()) params->set(
"can batch",
true);
199 for(size_type i=0; i<N; i++) {
200 Matrices[i]->apply(X,*Y[i],Teuchos::NO_TRANS, alpha, beta);
202 if(!params.is_null()) params->set(
"can batch",
false);
209 #endif // TPETRA_APPLY_HELPERS_HPP
Declaration of Tpetra::Details::Profiling, a scope guard for Kokkos Profiling.
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)
Does multiply matrix apply() calls with a single X vector.
void deep_copy(MultiVector< DS, DL, DG, DN > &dst, const MultiVector< SS, SL, SG, SN > &src)
Copy the contents of the MultiVector src into dst.
Insert new values that don't currently exist.
Replace old values with zero.
Accumulate new values into existing values (may not be supported in all classes)