Tpetra parallel linear algebra  Version of the Day
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Friends Macros Pages
Tpetra_applyDirichletBoundaryCondition.hpp
Go to the documentation of this file.
1 // @HEADER
2 // *****************************************************************************
3 // Tpetra: Templated Linear Algebra Services Package
4 //
5 // Copyright 2008 NTESS and the Tpetra contributors.
6 // SPDX-License-Identifier: BSD-3-Clause
7 // *****************************************************************************
8 // @HEADER
9 
10 #ifndef TPETRA_APPLYDIRICHLETBOUNDARYCONDITION_HPP
11 #define TPETRA_APPLYDIRICHLETBOUNDARYCONDITION_HPP
12 
16 
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"
22 
23 namespace Tpetra {
24 
34 template<class CrsMatrixType>
35 void
37 (const typename CrsMatrixType::execution_space& execSpace,
38  CrsMatrixType& A,
39  const Kokkos::View<
40  typename CrsMatrixType::local_ordinal_type*,
41  typename CrsMatrixType::device_type> & lclRowInds);
42 
52 template<class CrsMatrixType>
53 void
55 (CrsMatrixType& A,
56  const Kokkos::View<
57  typename CrsMatrixType::local_ordinal_type*,
58  typename CrsMatrixType::device_type>& lclRowInds);
59 
69 template<class CrsMatrixType>
70 void
72 (CrsMatrixType& A,
73  const Kokkos::View<
74  typename CrsMatrixType::local_ordinal_type*,
75  Kokkos::HostSpace> & lclRowInds);
76 
77 
87 template<class CrsMatrixType>
88 void
90 (const typename CrsMatrixType::execution_space& execSpace,
91  CrsMatrixType& A,
92  const Kokkos::View<
93  typename CrsMatrixType::local_ordinal_type*,
94  typename CrsMatrixType::device_type> & lclRowInds);
95 
96 
97 
98 
107 
109 template<class CrsMatrixType>
110 void
112 (CrsMatrixType& A,
113  const Kokkos::View<
114  typename CrsMatrixType::local_ordinal_type*,
115  Kokkos::HostSpace> & lclRowInds);
116 
126 template<class CrsMatrixType>
127 void
129 (CrsMatrixType& A,
130  const Kokkos::View<
131  typename CrsMatrixType::local_ordinal_type*,
132  typename CrsMatrixType::device_type> & lclRowInds);
133 
134 
135 namespace Details {
136 
137 template<class SC, class LO, class GO, class NT>
138 struct ApplyDirichletBoundaryConditionToLocalMatrixRows {
139  using crs_matrix_type = ::Tpetra::CrsMatrix<SC, LO, GO, NT>;
140  using execution_space = typename crs_matrix_type::execution_space;
141  using local_row_indices_type =
142  Kokkos::View<const LO*, Kokkos::AnonymousSpace>;
143 
144  static void
145  run (const execution_space& execSpace,
146  crs_matrix_type& A,
147  const local_row_indices_type& lclRowInds,
148  const bool runOnHost)
149  {
150  // Notes for future refactoring: This routine seems to have one more layer
151  // of options than it probably needs. For instance, if you passed a Kokkos::Serial
152  // execution_space instance as the first argument you probably wound't need the runOnHost
153  // option and then the code below could be collapsed out removing one of the parallel_for's
154 
155  using IST = typename crs_matrix_type::impl_scalar_type;
156  using KAT = Kokkos::ArithTraits<IST>;
157 
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 ();
167 
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 "
173  "at least once.");
174 
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;
180 
181  const bool wasFillComplete = A.isFillComplete ();
182  if (wasFillComplete) {
183  A.resumeFill ();
184  }
185 
186  const LO numInputRows = lclRowInds.extent (0);
187  if (! runOnHost) {
188  using range_type = Kokkos::RangePolicy<execution_space, LO>;
189  Kokkos::parallel_for
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) {
196  const bool diagEnt =
197  lclColMap.getGlobalElement (colind(j)) == row_gid;
198  values(j) = diagEnt ? KAT::one () : KAT::zero ();
199  }
200  });
201  }
202  else {
203  using range_type =
204  Kokkos::RangePolicy<Kokkos::DefaultHostExecutionSpace, LO>;
205  Kokkos::parallel_for
206  ("Tpetra::CrsMatrix apply Dirichlet: Host",
207  range_type (0, numInputRows),
208  [&] (const LO i) {
209  LO row = lclRowInds(i);
210  const GO row_gid = lclRowMap.getGlobalElement(row);
211  for (auto j = rowptr(row); j < rowptr(row+1); ++j) {
212  const bool diagEnt =
213  lclColMap.getGlobalElement (colind(j)) == row_gid;
214  values(j) = diagEnt ? KAT::one () : KAT::zero ();
215  }
216  });
217  }
218  if (wasFillComplete) {
219  A.fillComplete (A.getDomainMap (), A.getRangeMap ());
220  }
221  }
222 };
223 
224 
225 
226 template<class SC, class LO, class GO, class NT>
227 struct ApplyDirichletBoundaryConditionToLocalMatrixColumns {
228  using crs_matrix_type = ::Tpetra::CrsMatrix<SC, LO, GO, NT>;
229  using execution_space = typename crs_matrix_type::execution_space;
230  using local_col_flag_type =
231  Kokkos::View<bool*, Kokkos::AnonymousSpace>;
232 
233 
234  static void
235  run (const execution_space& execSpace,
236  crs_matrix_type& A,
237  const local_col_flag_type& lclColFlags,
238  const bool runOnHost)
239  {
240  // Notes for future refactoring: This routine seems to have one more layer
241  // of options than it probably needs. For instance, if you passed a Kokkos::Serial
242  // execution_space instance as the first argument you probably wound't need the runOnHost
243  // option and then the code below could be collapsed out removing one of the parallel_for's
244 
245  using IST = typename crs_matrix_type::impl_scalar_type;
246  using KAT = Kokkos::ArithTraits<IST>;
247 
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 ();
257 
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 "
263  "at least once.");
264 
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;
270 
271  const bool wasFillComplete = A.isFillComplete ();
272  if (wasFillComplete) {
273  A.resumeFill ();
274  }
275 
276  const LO numRows = (LO) A.getRowMap()->getLocalNumElements();
277  if (! runOnHost) {
278  using range_type = Kokkos::RangePolicy<execution_space, LO>;
279  Kokkos::parallel_for
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();
286  }
287  });
288  }
289  else {
290  using range_type =
291  Kokkos::RangePolicy<Kokkos::DefaultHostExecutionSpace, LO>;
292  Kokkos::parallel_for
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();
299  }
300  });
301  }
302  if (wasFillComplete) {
303  A.fillComplete (A.getDomainMap (), A.getRangeMap ());
304  }
305  }
306 };
307 
308 
309 
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) {
312  using crs_matrix_type = ::Tpetra::CrsMatrix<SC, LO, GO, NT>;
313  using execution_space = typename crs_matrix_type::execution_space;
314  using memory_space = typename crs_matrix_type::device_type::memory_space;
315 
316  // Need a colMap
317  TEUCHOS_TEST_FOR_EXCEPTION(A.getColMap().get () == nullptr, std::invalid_argument,"The matrix must have a column Map.");
318 
319  // NOTE: We assume that the RowMap and DomainMap of the matrix match.
320  // This could get relaxed at a later date, if we need that functionality
321  TEUCHOS_TEST_FOR_EXCEPTION(!A.getRowMap()->isSameAs(*A.getDomainMap()),std::invalid_argument, "localRowsToColumns: Row/Domain maps do not match");
322 
323  // Assume that the dirichletColFlags array is the correct size
324  TEUCHOS_TEST_FOR_EXCEPTION(A.getColMap()->getLocalNumElements() != dirichletColFlags.size(), std::invalid_argument,"localRowsToColumns: dirichletColFlags must be the correct size");
325 
326  LO numDirichletRows = (LO) dirichletRowIds.size();
327  LO LO_INVALID = Teuchos::OrdinalTraits<LO>::invalid();
328 
329  if(A.getCrsGraph()->getImporter().is_null()) {
330  // Serial case: If A doesn't have an importer, just set the flags from the dirichletRowIds
331  using range_type = Kokkos::RangePolicy<execution_space, LO>;
332  auto lclRowMap = A.getRowMap()->getLocalMap();
333  auto lclColMap = A.getColMap()->getLocalMap();
334 
335  Kokkos::deep_copy(execSpace,dirichletColFlags,false);
336  using range_type = Kokkos::RangePolicy<execution_space, LO>;
337  Kokkos::parallel_for
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;
345  });
346  }
347  else {
348  // Parallel case: Use A's importer to halo-exchange Dirichlet information
349  auto Importer = A.getCrsGraph()->getImporter();
350  auto lclRowMap = A.getRowMap()->getLocalMap();
351  auto lclColMap = A.getColMap()->getLocalMap();
352  ::Tpetra::Vector<LO,LO,GO,NT> domainDirichlet(A.getDomainMap());
353  ::Tpetra::Vector<LO,LO,GO,NT> colDirichlet(A.getColMap());
354  const LO one = Teuchos::OrdinalTraits<LO>::one();
355  using range_type = Kokkos::RangePolicy<execution_space, LO>;
356  {
357  auto domain_data = domainDirichlet.template getLocalView<memory_space>(Access::ReadWrite);
358  Kokkos::parallel_for
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;
366  });
367  }
368  colDirichlet.doImport(domainDirichlet,*Importer,::Tpetra::INSERT);
369  LO numCols = (LO) A.getColMap()->getLocalNumElements();
370  {
371  auto col_data = colDirichlet.template getLocalView<memory_space>(Access::ReadOnly);
372  Kokkos::parallel_for
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;
377  });
378  }
379  }
380 }
381 
382 } // namespace Details
383 
384 template<class CrsMatrixType>
385 void
387 (const typename CrsMatrixType::execution_space& execSpace,
388  CrsMatrixType& A,
389  const Kokkos::View<
390  typename CrsMatrixType::local_ordinal_type*,
391  typename CrsMatrixType::device_type> & lclRowInds)
392 {
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;
397 
398  using local_row_indices_type =
399  Kokkos::View<const LO*, Kokkos::AnonymousSpace>;
400  const local_row_indices_type lclRowInds_a (lclRowInds);
401 
402  using Details::ApplyDirichletBoundaryConditionToLocalMatrixRows;
403  using impl_type =
404  ApplyDirichletBoundaryConditionToLocalMatrixRows<SC, LO, GO, NT>;
405  const bool runOnHost = false;
406  impl_type::run (execSpace, A, lclRowInds_a, runOnHost);
407 }
408 
409 template<class CrsMatrixType>
410 void
412 (CrsMatrixType& A,
413  const Kokkos::View<
414  typename CrsMatrixType::local_ordinal_type*,
415  typename CrsMatrixType::device_type> & lclRowInds)
416 {
417  using execution_space = typename CrsMatrixType::execution_space;
418  applyDirichletBoundaryConditionToLocalMatrixRows (execution_space (), A, lclRowInds);
419 }
420 
421 template<class CrsMatrixType>
422 void
424 (CrsMatrixType& A,
425  const Kokkos::View<
426  typename CrsMatrixType::local_ordinal_type*,
427  Kokkos::HostSpace> & lclRowInds)
428 {
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;
433  using crs_matrix_type = ::Tpetra::CrsMatrix<SC, LO, GO, NT>;
434  using execution_space = typename crs_matrix_type::execution_space;
435  using memory_space = typename crs_matrix_type::device_type::memory_space;
436 
437  using Details::ApplyDirichletBoundaryConditionToLocalMatrixRows;
438  using impl_type =
439  ApplyDirichletBoundaryConditionToLocalMatrixRows<SC, LO, GO, NT>;
440 
441  // Only run on host if we can access the data
442  const bool runOnHost = Kokkos::SpaceAccessibility<Kokkos::Serial,memory_space>::accessible;
443  if(runOnHost) {
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);
447  }
448  else {
449  auto lclRowInds_a = Kokkos::create_mirror_view_and_copy(execution_space(),lclRowInds);
450  impl_type::run (execution_space (), A, lclRowInds_a, false);
451  }
452 }
453 
454 
455 template<class CrsMatrixType>
456 void
458 (CrsMatrixType& A,
459  const Kokkos::View<
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;
466  using crs_matrix_type = ::Tpetra::CrsMatrix<SC, LO, GO, NT>;
467  using execution_space = typename crs_matrix_type::execution_space;
468  using memory_space = typename crs_matrix_type::device_type::memory_space;
469 
470  TEUCHOS_TEST_FOR_EXCEPTION(A.getColMap().get () == nullptr, std::invalid_argument,"The matrix must have a column Map.");
471 
472  // Copy the Host array to device
473  auto lclRowInds_d = Kokkos::create_mirror_view_and_copy(execution_space(),lclRowInds);
474 
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);
478 
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);
481 }
482 
483 
484 template<class CrsMatrixType>
485 void
487 (CrsMatrixType& A,
488  const Kokkos::View<
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;
495  using crs_matrix_type = ::Tpetra::CrsMatrix<SC, LO, GO, NT>;
496  using execution_space = typename crs_matrix_type::execution_space;
497  using memory_space = typename crs_matrix_type::device_type::memory_space;
498 
499  TEUCHOS_TEST_FOR_EXCEPTION(A.getColMap().get () == nullptr, std::invalid_argument,"The matrix must have a column Map.");
500 
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);
504 
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);
507 
508 }
509 
510 
511 template<class CrsMatrixType>
512 void
514 (const typename CrsMatrixType::execution_space& execSpace,
515  CrsMatrixType& A,
516  const Kokkos::View<
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;
523  using crs_matrix_type = ::Tpetra::CrsMatrix<SC, LO, GO, NT>;
524  // using execution_space = typename crs_matrix_type::execution_space;
525  using memory_space = typename crs_matrix_type::device_type::memory_space;
526 
527  TEUCHOS_TEST_FOR_EXCEPTION(A.getColMap().get () == nullptr, std::invalid_argument,"The matrix must have a column Map.");
528 
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);
532 
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);
535 
536 }
537 
538 
539 } // namespace Tpetra
540 
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&#39;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...