10 #ifndef TPETRA_APPLYDIRICHLETBOUNDARYCONDITION_HPP
11 #define TPETRA_APPLYDIRICHLETBOUNDARYCONDITION_HPP
17 #include "Tpetra_CrsMatrix.hpp"
18 #include "Tpetra_Vector.hpp"
19 #include "Tpetra_Map.hpp"
20 #include "KokkosSparse_CrsMatrix.hpp"
21 #include "Kokkos_ArithTraits.hpp"
34 template<
class CrsMatrixType>
37 (
const typename CrsMatrixType::execution_space& execSpace,
40 typename CrsMatrixType::local_ordinal_type*,
41 typename CrsMatrixType::device_type> & lclRowInds);
52 template<
class CrsMatrixType>
57 typename CrsMatrixType::local_ordinal_type*,
58 typename CrsMatrixType::device_type>& lclRowInds);
69 template<
class CrsMatrixType>
74 typename CrsMatrixType::local_ordinal_type*,
75 Kokkos::HostSpace> & lclRowInds);
87 template<
class CrsMatrixType>
90 (
const typename CrsMatrixType::execution_space& execSpace,
93 typename CrsMatrixType::local_ordinal_type*,
94 typename CrsMatrixType::device_type> & lclRowInds);
109 template<
class CrsMatrixType>
114 typename CrsMatrixType::local_ordinal_type*,
115 Kokkos::HostSpace> & lclRowInds);
126 template<
class CrsMatrixType>
131 typename CrsMatrixType::local_ordinal_type*,
132 typename CrsMatrixType::device_type> & lclRowInds);
137 template<
class SC,
class LO,
class GO,
class NT>
138 struct ApplyDirichletBoundaryConditionToLocalMatrixRows {
141 using local_row_indices_type =
142 Kokkos::View<const LO*, Kokkos::AnonymousSpace>;
145 run (
const execution_space& execSpace,
147 const local_row_indices_type& lclRowInds,
148 const bool runOnHost)
156 using KAT = Kokkos::ArithTraits<IST>;
158 const auto rowMap = A.getRowMap ();
159 TEUCHOS_TEST_FOR_EXCEPTION
160 (rowMap.get () ==
nullptr, std::invalid_argument,
161 "The matrix must have a row Map.");
162 const auto colMap = A.getColMap ();
163 TEUCHOS_TEST_FOR_EXCEPTION
164 (colMap.get () ==
nullptr, std::invalid_argument,
165 "The matrix must have a column Map.");
166 auto A_lcl = A.getLocalMatrixDevice ();
168 const LO lclNumRows =
static_cast<LO
> (rowMap->getLocalNumElements ());
169 TEUCHOS_TEST_FOR_EXCEPTION
170 (lclNumRows != 0 && static_cast<LO>(A_lcl.graph.numRows ()) != lclNumRows,
171 std::invalid_argument,
"The matrix must have been either created "
172 "with a KokkosSparse::CrsMatrix, or must have been fill-completed "
175 auto lclRowMap = A.getRowMap ()->getLocalMap ();
176 auto lclColMap = A.getColMap ()->getLocalMap ();
177 auto rowptr = A_lcl.graph.row_map;
178 auto colind = A_lcl.graph.entries;
179 auto values = A_lcl.values;
181 const bool wasFillComplete = A.isFillComplete ();
182 if (wasFillComplete) {
186 const LO numInputRows = lclRowInds.extent (0);
188 using range_type = Kokkos::RangePolicy<execution_space, LO>;
190 (
"Tpetra::CrsMatrix apply Dirichlet: Device",
191 range_type (execSpace, 0, numInputRows),
192 KOKKOS_LAMBDA (
const LO i) {
193 LO row = lclRowInds(i);
194 const GO row_gid = lclRowMap.getGlobalElement(row);
195 for (
auto j = rowptr(row); j < rowptr(row+1); ++j) {
197 lclColMap.getGlobalElement (colind(j)) == row_gid;
198 values(j) = diagEnt ? KAT::one () : KAT::zero ();
204 Kokkos::RangePolicy<Kokkos::DefaultHostExecutionSpace, LO>;
206 (
"Tpetra::CrsMatrix apply Dirichlet: Host",
207 range_type (0, numInputRows),
209 LO row = lclRowInds(i);
210 const GO row_gid = lclRowMap.getGlobalElement(row);
211 for (
auto j = rowptr(row); j < rowptr(row+1); ++j) {
213 lclColMap.getGlobalElement (colind(j)) == row_gid;
214 values(j) = diagEnt ? KAT::one () : KAT::zero ();
218 if (wasFillComplete) {
219 A.fillComplete (A.getDomainMap (), A.getRangeMap ());
226 template<
class SC,
class LO,
class GO,
class NT>
227 struct ApplyDirichletBoundaryConditionToLocalMatrixColumns {
230 using local_col_flag_type =
231 Kokkos::View<bool*, Kokkos::AnonymousSpace>;
235 run (
const execution_space& execSpace,
237 const local_col_flag_type& lclColFlags,
238 const bool runOnHost)
246 using KAT = Kokkos::ArithTraits<IST>;
248 const auto rowMap = A.getRowMap ();
249 TEUCHOS_TEST_FOR_EXCEPTION
250 (rowMap.get () ==
nullptr, std::invalid_argument,
251 "The matrix must have a row Map.");
252 const auto colMap = A.getColMap ();
253 TEUCHOS_TEST_FOR_EXCEPTION
254 (colMap.get () ==
nullptr, std::invalid_argument,
255 "The matrix must have a column Map.");
256 auto A_lcl = A.getLocalMatrixDevice ();
258 const LO lclNumRows =
static_cast<LO
> (rowMap->getLocalNumElements ());
259 TEUCHOS_TEST_FOR_EXCEPTION
260 (lclNumRows != 0 && static_cast<LO>(A_lcl.graph.numRows ()) != lclNumRows,
261 std::invalid_argument,
"The matrix must have been either created "
262 "with a KokkosSparse::CrsMatrix, or must have been fill-completed "
265 auto lclRowMap = A.getRowMap()->getLocalMap ();
266 auto lclColMap = A.getColMap()->getLocalMap ();
267 auto rowptr = A_lcl.graph.row_map;
268 auto colind = A_lcl.graph.entries;
269 auto values = A_lcl.values;
271 const bool wasFillComplete = A.isFillComplete ();
272 if (wasFillComplete) {
276 const LO numRows = (LO) A.getRowMap()->getLocalNumElements();
278 using range_type = Kokkos::RangePolicy<execution_space, LO>;
280 (
"Tpetra::CrsMatrix apply Dirichlet cols: Device",
281 range_type (execSpace, 0, numRows),
282 KOKKOS_LAMBDA (
const LO i) {
283 for (
auto j = rowptr(i); j < rowptr(i+1); ++j) {
284 if(lclColFlags[colind[j]])
285 values(j) = KAT::zero();
291 Kokkos::RangePolicy<Kokkos::DefaultHostExecutionSpace, LO>;
293 (
"Tpetra::CrsMatrix apply Dirichlet cols: Host",
294 range_type (0, numRows),
295 KOKKOS_LAMBDA (
const LO i) {
296 for (
auto j = rowptr(i); j < rowptr(i+1); ++j) {
297 if(lclColFlags[colind[j]])
298 values(j) = KAT::zero();
302 if (wasFillComplete) {
303 A.fillComplete (A.getDomainMap (), A.getRangeMap ());
310 template<
class SC,
class LO,
class GO,
class NT>
311 void localRowsToColumns(
const typename ::Tpetra::CrsMatrix<SC, LO, GO, NT>::execution_space& execSpace, const ::Tpetra::CrsMatrix<SC, LO, GO, NT>& A,
const Kokkos::View<const LO*,Kokkos::AnonymousSpace> & dirichletRowIds, Kokkos::View<bool*,Kokkos::AnonymousSpace> & dirichletColFlags) {
313 using execution_space =
typename crs_matrix_type::execution_space;
314 using memory_space =
typename crs_matrix_type::device_type::memory_space;
317 TEUCHOS_TEST_FOR_EXCEPTION(A.getColMap().get () ==
nullptr, std::invalid_argument,
"The matrix must have a column Map.");
321 TEUCHOS_TEST_FOR_EXCEPTION(!A.getRowMap()->isSameAs(*A.getDomainMap()),std::invalid_argument,
"localRowsToColumns: Row/Domain maps do not match");
324 TEUCHOS_TEST_FOR_EXCEPTION(A.getColMap()->getLocalNumElements() != dirichletColFlags.size(), std::invalid_argument,
"localRowsToColumns: dirichletColFlags must be the correct size");
326 LO numDirichletRows = (LO) dirichletRowIds.size();
327 LO LO_INVALID = Teuchos::OrdinalTraits<LO>::invalid();
329 if(A.getCrsGraph()->getImporter().is_null()) {
331 using range_type = Kokkos::RangePolicy<execution_space, LO>;
332 auto lclRowMap = A.getRowMap()->getLocalMap();
333 auto lclColMap = A.getColMap()->getLocalMap();
336 using range_type = Kokkos::RangePolicy<execution_space, LO>;
338 (
"Tpetra::CrsMatrix flag Dirichlet cols",
339 range_type (execSpace, 0, numDirichletRows),
340 KOKKOS_LAMBDA (
const LO i) {
341 GO row_gid = lclRowMap.getGlobalElement(dirichletRowIds[i]);
342 LO col_lid = lclColMap.getLocalElement(row_gid);
343 if(col_lid != LO_INVALID)
344 dirichletColFlags[col_lid] =
true;
349 auto Importer = A.getCrsGraph()->getImporter();
350 auto lclRowMap = A.getRowMap()->getLocalMap();
351 auto lclColMap = A.getColMap()->getLocalMap();
354 const LO one = Teuchos::OrdinalTraits<LO>::one();
355 using range_type = Kokkos::RangePolicy<execution_space, LO>;
357 auto domain_data = domainDirichlet.template getLocalView<memory_space>(Access::ReadWrite);
359 (
"Tpetra::CrsMatrix flag Dirichlet domains",
360 range_type (execSpace, 0, numDirichletRows),
361 KOKKOS_LAMBDA (
const LO i) {
362 GO row_gid = lclRowMap.getGlobalElement(dirichletRowIds[i]);
363 LO col_lid = lclColMap.getLocalElement(row_gid);
364 if(col_lid != LO_INVALID)
365 domain_data(col_lid,0) = one;
368 colDirichlet.doImport(domainDirichlet,*Importer,::
Tpetra::INSERT);
369 LO numCols = (LO) A.getColMap()->getLocalNumElements();
371 auto col_data = colDirichlet.template getLocalView<memory_space>(Access::ReadOnly);
373 (
"Tpetra::CrsMatrix flag Dirichlet cols",
374 range_type (execSpace, 0, numCols),
375 KOKKOS_LAMBDA (
const LO i) {
376 dirichletColFlags[i] = (col_data(i,0) == one) ?
true :
false;
384 template<
class CrsMatrixType>
387 (
const typename CrsMatrixType::execution_space& execSpace,
390 typename CrsMatrixType::local_ordinal_type*,
391 typename CrsMatrixType::device_type> & lclRowInds)
393 using SC =
typename CrsMatrixType::scalar_type;
394 using LO =
typename CrsMatrixType::local_ordinal_type;
395 using GO =
typename CrsMatrixType::global_ordinal_type;
396 using NT =
typename CrsMatrixType::node_type;
398 using local_row_indices_type =
399 Kokkos::View<const LO*, Kokkos::AnonymousSpace>;
400 const local_row_indices_type lclRowInds_a (lclRowInds);
402 using Details::ApplyDirichletBoundaryConditionToLocalMatrixRows;
404 ApplyDirichletBoundaryConditionToLocalMatrixRows<SC, LO, GO, NT>;
405 const bool runOnHost =
false;
406 impl_type::run (execSpace, A, lclRowInds_a, runOnHost);
409 template<
class CrsMatrixType>
414 typename CrsMatrixType::local_ordinal_type*,
415 typename CrsMatrixType::device_type> & lclRowInds)
417 using execution_space =
typename CrsMatrixType::execution_space;
421 template<
class CrsMatrixType>
426 typename CrsMatrixType::local_ordinal_type*,
427 Kokkos::HostSpace> & lclRowInds)
429 using SC =
typename CrsMatrixType::scalar_type;
430 using LO =
typename CrsMatrixType::local_ordinal_type;
431 using GO =
typename CrsMatrixType::global_ordinal_type;
432 using NT =
typename CrsMatrixType::node_type;
434 using execution_space =
typename crs_matrix_type::execution_space;
435 using memory_space =
typename crs_matrix_type::device_type::memory_space;
437 using Details::ApplyDirichletBoundaryConditionToLocalMatrixRows;
439 ApplyDirichletBoundaryConditionToLocalMatrixRows<SC, LO, GO, NT>;
442 const bool runOnHost = Kokkos::SpaceAccessibility<Kokkos::Serial,memory_space>::accessible;
444 using local_row_indices_type = Kokkos::View<const LO*, Kokkos::AnonymousSpace>;
445 const local_row_indices_type lclRowInds_a (lclRowInds);
446 impl_type::run (execution_space (), A, lclRowInds_a,
true);
449 auto lclRowInds_a = Kokkos::create_mirror_view_and_copy(execution_space(),lclRowInds);
450 impl_type::run (execution_space (), A, lclRowInds_a,
false);
455 template<
class CrsMatrixType>
460 typename CrsMatrixType::local_ordinal_type*,
461 Kokkos::HostSpace> & lclRowInds) {
462 using SC =
typename CrsMatrixType::scalar_type;
463 using LO =
typename CrsMatrixType::local_ordinal_type;
464 using GO =
typename CrsMatrixType::global_ordinal_type;
465 using NT =
typename CrsMatrixType::node_type;
467 using execution_space =
typename crs_matrix_type::execution_space;
468 using memory_space =
typename crs_matrix_type::device_type::memory_space;
470 TEUCHOS_TEST_FOR_EXCEPTION(A.getColMap().get () ==
nullptr, std::invalid_argument,
"The matrix must have a column Map.");
473 auto lclRowInds_d = Kokkos::create_mirror_view_and_copy(execution_space(),lclRowInds);
475 Kokkos::View<bool*,memory_space> dirichletColFlags(
"dirichletColFlags",A.getColMap()->getLocalNumElements());
476 Kokkos::View<bool*, Kokkos::AnonymousSpace> dirichletColFlags_a(dirichletColFlags);
477 Details::localRowsToColumns<SC,LO,GO,NT>(execution_space(),A,lclRowInds_d,dirichletColFlags_a);
479 Details::ApplyDirichletBoundaryConditionToLocalMatrixColumns<SC, LO, GO, NT>::run(execution_space(),A,dirichletColFlags,
false);
480 Details::ApplyDirichletBoundaryConditionToLocalMatrixRows<SC, LO, GO, NT>::run(execution_space(),A,lclRowInds_d,
false);
484 template<
class CrsMatrixType>
489 typename CrsMatrixType::local_ordinal_type*,
490 typename CrsMatrixType::device_type> & lclRowInds_d) {
491 using SC =
typename CrsMatrixType::scalar_type;
492 using LO =
typename CrsMatrixType::local_ordinal_type;
493 using GO =
typename CrsMatrixType::global_ordinal_type;
494 using NT =
typename CrsMatrixType::node_type;
496 using execution_space =
typename crs_matrix_type::execution_space;
497 using memory_space =
typename crs_matrix_type::device_type::memory_space;
499 TEUCHOS_TEST_FOR_EXCEPTION(A.getColMap().get () ==
nullptr, std::invalid_argument,
"The matrix must have a column Map.");
501 Kokkos::View<bool*,memory_space> dirichletColFlags(
"dirichletColFlags",A.getColMap()->getLocalNumElements());
502 Kokkos::View<bool*, Kokkos::AnonymousSpace> dirichletColFlags_a(dirichletColFlags);
503 Details::localRowsToColumns<SC,LO,GO,NT>(execution_space(),A,lclRowInds_d,dirichletColFlags_a);
505 Details::ApplyDirichletBoundaryConditionToLocalMatrixColumns<SC, LO, GO, NT>::run(execution_space(),A,dirichletColFlags,
false);
506 Details::ApplyDirichletBoundaryConditionToLocalMatrixRows<SC, LO, GO, NT>::run(execution_space(),A,lclRowInds_d,
false);
511 template<
class CrsMatrixType>
514 (
const typename CrsMatrixType::execution_space& execSpace,
517 typename CrsMatrixType::local_ordinal_type*,
518 typename CrsMatrixType::device_type> & lclRowInds_d) {
519 using SC =
typename CrsMatrixType::scalar_type;
520 using LO =
typename CrsMatrixType::local_ordinal_type;
521 using GO =
typename CrsMatrixType::global_ordinal_type;
522 using NT =
typename CrsMatrixType::node_type;
525 using memory_space =
typename crs_matrix_type::device_type::memory_space;
527 TEUCHOS_TEST_FOR_EXCEPTION(A.getColMap().get () ==
nullptr, std::invalid_argument,
"The matrix must have a column Map.");
529 Kokkos::View<bool*,memory_space> dirichletColFlags(
"dirichletColFlags",A.getColMap()->getLocalNumElements());
530 Kokkos::View<bool*, Kokkos::AnonymousSpace> dirichletColFlags_a(dirichletColFlags);
531 Details::localRowsToColumns<SC,LO,GO,NT>(execSpace,A,lclRowInds_d,dirichletColFlags_a);
533 Details::ApplyDirichletBoundaryConditionToLocalMatrixColumns<SC, LO, GO, NT>::run(execSpace,A,dirichletColFlags,
false);
534 Details::ApplyDirichletBoundaryConditionToLocalMatrixRows<SC, LO, GO, NT>::run(execSpace,A,lclRowInds_d,
false);
541 #endif // TPETRA_APPLYDIRICHLETBOUNDARYCONDITION_HPP
Sparse matrix that presents a row-oriented interface that lets users read or modify entries...
typename device_type::execution_space execution_space
The Kokkos execution space.
void applyDirichletBoundaryConditionToLocalMatrixRowsAndColumns(const typename CrsMatrixType::execution_space &execSpace, CrsMatrixType &A, const Kokkos::View< typename CrsMatrixType::local_ordinal_type *, typename CrsMatrixType::device_type > &lclRowInds)
For all k in [0, lclRowInds.extent(0)), set local row and column lclRowInds[k] of A to have 1 on the ...
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.
typename row_matrix_type::impl_scalar_type impl_scalar_type
The type used internally in place of Scalar.
A distributed dense vector.
void applyDirichletBoundaryConditionToLocalMatrixRows(const typename CrsMatrixType::execution_space &execSpace, CrsMatrixType &A, const Kokkos::View< typename CrsMatrixType::local_ordinal_type *, typename CrsMatrixType::device_type > &lclRowInds)
For all k in [0, lclRowInds.extent(0)), set local row lclRowInds[k] of A to have 1 on the diagonal an...