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 #if KOKKOS_VERSION >= 40799
22 #include "KokkosKernels_ArithTraits.hpp"
23 #else
24 #include "Kokkos_ArithTraits.hpp"
25 #endif
26 
27 namespace Tpetra {
28 
38 template <class CrsMatrixType>
39 void applyDirichletBoundaryConditionToLocalMatrixRows(const typename CrsMatrixType::execution_space& execSpace,
40  CrsMatrixType& A,
41  const Kokkos::View<
42  typename CrsMatrixType::local_ordinal_type*,
43  typename CrsMatrixType::device_type>& lclRowInds);
44 
54 template <class CrsMatrixType>
56  const Kokkos::View<
57  typename CrsMatrixType::local_ordinal_type*,
58  typename CrsMatrixType::device_type>& lclRowInds);
59 
69 template <class CrsMatrixType>
71  const Kokkos::View<
72  typename CrsMatrixType::local_ordinal_type*,
73  Kokkos::HostSpace>& lclRowInds);
74 
84 template <class CrsMatrixType>
85 void applyDirichletBoundaryConditionToLocalMatrixRowsAndColumns(const typename CrsMatrixType::execution_space& execSpace,
86  CrsMatrixType& A,
87  const Kokkos::View<
88  typename CrsMatrixType::local_ordinal_type*,
89  typename CrsMatrixType::device_type>& lclRowInds);
90 
99 
101 template <class CrsMatrixType>
103  const Kokkos::View<
104  typename CrsMatrixType::local_ordinal_type*,
105  Kokkos::HostSpace>& lclRowInds);
106 
116 template <class CrsMatrixType>
118  const Kokkos::View<
119  typename CrsMatrixType::local_ordinal_type*,
120  typename CrsMatrixType::device_type>& lclRowInds);
121 
122 namespace Details {
123 
124 template <class SC, class LO, class GO, class NT>
125 struct ApplyDirichletBoundaryConditionToLocalMatrixRows {
126  using crs_matrix_type = ::Tpetra::CrsMatrix<SC, LO, GO, NT>;
127  using execution_space = typename crs_matrix_type::execution_space;
128  using local_row_indices_type =
129  Kokkos::View<const LO*, Kokkos::AnonymousSpace>;
130 
131  static void
132  run(const execution_space& execSpace,
133  crs_matrix_type& A,
134  const local_row_indices_type& lclRowInds,
135  const bool runOnHost) {
136  // Notes for future refactoring: This routine seems to have one more layer
137  // of options than it probably needs. For instance, if you passed a Kokkos::Serial
138  // execution_space instance as the first argument you probably wound't need the runOnHost
139  // option and then the code below could be collapsed out removing one of the parallel_for's
140 
141  using IST = typename crs_matrix_type::impl_scalar_type;
142 #if KOKKOS_VERSION >= 40799
143  using KAT = KokkosKernels::ArithTraits<IST>;
144 #else
145  using KAT = Kokkos::ArithTraits<IST>;
146 #endif
147 
148  const auto rowMap = A.getRowMap();
149  TEUCHOS_TEST_FOR_EXCEPTION(rowMap.get() == nullptr, std::invalid_argument,
150  "The matrix must have a row Map.");
151  const auto colMap = A.getColMap();
152  TEUCHOS_TEST_FOR_EXCEPTION(colMap.get() == nullptr, std::invalid_argument,
153  "The matrix must have a column Map.");
154  auto A_lcl = A.getLocalMatrixDevice();
155 
156  const LO lclNumRows = static_cast<LO>(rowMap->getLocalNumElements());
157  TEUCHOS_TEST_FOR_EXCEPTION(lclNumRows != 0 && static_cast<LO>(A_lcl.graph.numRows()) != lclNumRows,
158  std::invalid_argument,
159  "The matrix must have been either created "
160  "with a KokkosSparse::CrsMatrix, or must have been fill-completed "
161  "at least once.");
162 
163  auto lclRowMap = A.getRowMap()->getLocalMap();
164  auto lclColMap = A.getColMap()->getLocalMap();
165  auto rowptr = A_lcl.graph.row_map;
166  auto colind = A_lcl.graph.entries;
167  auto values = A_lcl.values;
168 
169  const bool wasFillComplete = A.isFillComplete();
170  if (wasFillComplete) {
171  A.resumeFill();
172  }
173 
174  const LO numInputRows = lclRowInds.extent(0);
175  if (!runOnHost) {
176  using range_type = Kokkos::RangePolicy<execution_space, LO>;
177  Kokkos::parallel_for(
178  "Tpetra::CrsMatrix apply Dirichlet: Device",
179  range_type(execSpace, 0, numInputRows),
180  KOKKOS_LAMBDA(const LO i) {
181  LO row = lclRowInds(i);
182  const GO row_gid = lclRowMap.getGlobalElement(row);
183  for (auto j = rowptr(row); j < rowptr(row + 1); ++j) {
184  const bool diagEnt =
185  lclColMap.getGlobalElement(colind(j)) == row_gid;
186  values(j) = diagEnt ? KAT::one() : KAT::zero();
187  }
188  });
189  } else {
190  using range_type =
191  Kokkos::RangePolicy<Kokkos::DefaultHostExecutionSpace, LO>;
192  Kokkos::parallel_for("Tpetra::CrsMatrix apply Dirichlet: Host",
193  range_type(0, numInputRows),
194  [&](const LO i) {
195  LO row = lclRowInds(i);
196  const GO row_gid = lclRowMap.getGlobalElement(row);
197  for (auto j = rowptr(row); j < rowptr(row + 1); ++j) {
198  const bool diagEnt =
199  lclColMap.getGlobalElement(colind(j)) == row_gid;
200  values(j) = diagEnt ? KAT::one() : KAT::zero();
201  }
202  });
203  }
204  if (wasFillComplete) {
205  A.fillComplete(A.getDomainMap(), A.getRangeMap());
206  }
207  }
208 };
209 
210 template <class SC, class LO, class GO, class NT>
211 struct ApplyDirichletBoundaryConditionToLocalMatrixColumns {
212  using crs_matrix_type = ::Tpetra::CrsMatrix<SC, LO, GO, NT>;
213  using execution_space = typename crs_matrix_type::execution_space;
214  using local_col_flag_type =
215  Kokkos::View<bool*, Kokkos::AnonymousSpace>;
216 
217  static void
218  run(const execution_space& execSpace,
219  crs_matrix_type& A,
220  const local_col_flag_type& lclColFlags,
221  const bool runOnHost) {
222  // Notes for future refactoring: This routine seems to have one more layer
223  // of options than it probably needs. For instance, if you passed a Kokkos::Serial
224  // execution_space instance as the first argument you probably wound't need the runOnHost
225  // option and then the code below could be collapsed out removing one of the parallel_for's
226 
227  using IST = typename crs_matrix_type::impl_scalar_type;
228 #if KOKKOS_VERSION >= 40799
229  using KAT = KokkosKernels::ArithTraits<IST>;
230 #else
231  using KAT = Kokkos::ArithTraits<IST>;
232 #endif
233 
234  const auto rowMap = A.getRowMap();
235  TEUCHOS_TEST_FOR_EXCEPTION(rowMap.get() == nullptr, std::invalid_argument,
236  "The matrix must have a row Map.");
237  const auto colMap = A.getColMap();
238  TEUCHOS_TEST_FOR_EXCEPTION(colMap.get() == nullptr, std::invalid_argument,
239  "The matrix must have a column Map.");
240  auto A_lcl = A.getLocalMatrixDevice();
241 
242  const LO lclNumRows = static_cast<LO>(rowMap->getLocalNumElements());
243  TEUCHOS_TEST_FOR_EXCEPTION(lclNumRows != 0 && static_cast<LO>(A_lcl.graph.numRows()) != lclNumRows,
244  std::invalid_argument,
245  "The matrix must have been either created "
246  "with a KokkosSparse::CrsMatrix, or must have been fill-completed "
247  "at least once.");
248 
249  auto lclRowMap = A.getRowMap()->getLocalMap();
250  auto lclColMap = A.getColMap()->getLocalMap();
251  auto rowptr = A_lcl.graph.row_map;
252  auto colind = A_lcl.graph.entries;
253  auto values = A_lcl.values;
254 
255  const bool wasFillComplete = A.isFillComplete();
256  if (wasFillComplete) {
257  A.resumeFill();
258  }
259 
260  const LO numRows = (LO)A.getRowMap()->getLocalNumElements();
261  if (!runOnHost) {
262  using range_type = Kokkos::RangePolicy<execution_space, LO>;
263  Kokkos::parallel_for(
264  "Tpetra::CrsMatrix apply Dirichlet cols: Device",
265  range_type(execSpace, 0, numRows),
266  KOKKOS_LAMBDA(const LO i) {
267  for (auto j = rowptr(i); j < rowptr(i + 1); ++j) {
268  if (lclColFlags[colind[j]])
269  values(j) = KAT::zero();
270  }
271  });
272  } else {
273  using range_type =
274  Kokkos::RangePolicy<Kokkos::DefaultHostExecutionSpace, LO>;
275  Kokkos::parallel_for(
276  "Tpetra::CrsMatrix apply Dirichlet cols: Host",
277  range_type(0, numRows),
278  KOKKOS_LAMBDA(const LO i) {
279  for (auto j = rowptr(i); j < rowptr(i + 1); ++j) {
280  if (lclColFlags[colind[j]])
281  values(j) = KAT::zero();
282  }
283  });
284  }
285  if (wasFillComplete) {
286  A.fillComplete(A.getDomainMap(), A.getRangeMap());
287  }
288  }
289 };
290 
291 template <class SC, class LO, class GO, class NT>
292 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) {
293  using crs_matrix_type = ::Tpetra::CrsMatrix<SC, LO, GO, NT>;
294  using execution_space = typename crs_matrix_type::execution_space;
295  using memory_space = typename crs_matrix_type::device_type::memory_space;
296 
297  // Need a colMap
298  TEUCHOS_TEST_FOR_EXCEPTION(A.getColMap().get() == nullptr, std::invalid_argument, "The matrix must have a column Map.");
299 
300  // NOTE: We assume that the RowMap and DomainMap of the matrix match.
301  // This could get relaxed at a later date, if we need that functionality
302  TEUCHOS_TEST_FOR_EXCEPTION(!A.getRowMap()->isSameAs(*A.getDomainMap()), std::invalid_argument, "localRowsToColumns: Row/Domain maps do not match");
303 
304  // Assume that the dirichletColFlags array is the correct size
305  TEUCHOS_TEST_FOR_EXCEPTION(A.getColMap()->getLocalNumElements() != dirichletColFlags.size(), std::invalid_argument, "localRowsToColumns: dirichletColFlags must be the correct size");
306 
307  LO numDirichletRows = (LO)dirichletRowIds.size();
308  LO LO_INVALID = Teuchos::OrdinalTraits<LO>::invalid();
309 
310  if (A.getCrsGraph()->getImporter().is_null()) {
311  // Serial case: If A doesn't have an importer, just set the flags from the dirichletRowIds
312  using range_type = Kokkos::RangePolicy<execution_space, LO>;
313  auto lclRowMap = A.getRowMap()->getLocalMap();
314  auto lclColMap = A.getColMap()->getLocalMap();
315 
316  Kokkos::deep_copy(execSpace, dirichletColFlags, false);
317  using range_type = Kokkos::RangePolicy<execution_space, LO>;
318  Kokkos::parallel_for(
319  "Tpetra::CrsMatrix flag Dirichlet cols",
320  range_type(execSpace, 0, numDirichletRows),
321  KOKKOS_LAMBDA(const LO i) {
322  GO row_gid = lclRowMap.getGlobalElement(dirichletRowIds[i]);
323  LO col_lid = lclColMap.getLocalElement(row_gid);
324  if (col_lid != LO_INVALID)
325  dirichletColFlags[col_lid] = true;
326  });
327  } else {
328  // Parallel case: Use A's importer to halo-exchange Dirichlet information
329  auto Importer = A.getCrsGraph()->getImporter();
330  auto lclRowMap = A.getRowMap()->getLocalMap();
331  auto lclColMap = A.getColMap()->getLocalMap();
332  ::Tpetra::Vector<LO, LO, GO, NT> domainDirichlet(A.getDomainMap());
333  ::Tpetra::Vector<LO, LO, GO, NT> colDirichlet(A.getColMap());
334  const LO one = Teuchos::OrdinalTraits<LO>::one();
335  using range_type = Kokkos::RangePolicy<execution_space, LO>;
336  {
337  auto domain_data = domainDirichlet.template getLocalView<memory_space>(Access::ReadWrite);
338  Kokkos::parallel_for(
339  "Tpetra::CrsMatrix flag Dirichlet domains",
340  range_type(execSpace, 0, numDirichletRows),
341  KOKKOS_LAMBDA(const LO i) {
342  GO row_gid = lclRowMap.getGlobalElement(dirichletRowIds[i]);
343  LO col_lid = lclColMap.getLocalElement(row_gid);
344  if (col_lid != LO_INVALID)
345  domain_data(col_lid, 0) = one;
346  });
347  }
348  colDirichlet.doImport(domainDirichlet, *Importer, ::Tpetra::INSERT);
349  LO numCols = (LO)A.getColMap()->getLocalNumElements();
350  {
351  auto col_data = colDirichlet.template getLocalView<memory_space>(Access::ReadOnly);
352  Kokkos::parallel_for(
353  "Tpetra::CrsMatrix flag Dirichlet cols",
354  range_type(execSpace, 0, numCols),
355  KOKKOS_LAMBDA(const LO i) {
356  dirichletColFlags[i] = (col_data(i, 0) == one) ? true : false;
357  });
358  }
359  }
360 }
361 
362 } // namespace Details
363 
364 template <class CrsMatrixType>
365 void applyDirichletBoundaryConditionToLocalMatrixRows(const typename CrsMatrixType::execution_space& execSpace,
366  CrsMatrixType& A,
367  const Kokkos::View<
368  typename CrsMatrixType::local_ordinal_type*,
369  typename CrsMatrixType::device_type>& lclRowInds) {
370  using SC = typename CrsMatrixType::scalar_type;
371  using LO = typename CrsMatrixType::local_ordinal_type;
372  using GO = typename CrsMatrixType::global_ordinal_type;
373  using NT = typename CrsMatrixType::node_type;
374 
375  using local_row_indices_type =
376  Kokkos::View<const LO*, Kokkos::AnonymousSpace>;
377  const local_row_indices_type lclRowInds_a(lclRowInds);
378 
379  using Details::ApplyDirichletBoundaryConditionToLocalMatrixRows;
380  using impl_type =
381  ApplyDirichletBoundaryConditionToLocalMatrixRows<SC, LO, GO, NT>;
382  const bool runOnHost = false;
383  impl_type::run(execSpace, A, lclRowInds_a, runOnHost);
384 }
385 
386 template <class CrsMatrixType>
388  const Kokkos::View<
389  typename CrsMatrixType::local_ordinal_type*,
390  typename CrsMatrixType::device_type>& lclRowInds) {
391  using execution_space = typename CrsMatrixType::execution_space;
392  applyDirichletBoundaryConditionToLocalMatrixRows(execution_space(), A, lclRowInds);
393 }
394 
395 template <class CrsMatrixType>
397  const Kokkos::View<
398  typename CrsMatrixType::local_ordinal_type*,
399  Kokkos::HostSpace>& lclRowInds) {
400  using SC = typename CrsMatrixType::scalar_type;
401  using LO = typename CrsMatrixType::local_ordinal_type;
402  using GO = typename CrsMatrixType::global_ordinal_type;
403  using NT = typename CrsMatrixType::node_type;
404  using crs_matrix_type = ::Tpetra::CrsMatrix<SC, LO, GO, NT>;
405  using execution_space = typename crs_matrix_type::execution_space;
406  using memory_space = typename crs_matrix_type::device_type::memory_space;
407 
408  using Details::ApplyDirichletBoundaryConditionToLocalMatrixRows;
409  using impl_type =
410  ApplyDirichletBoundaryConditionToLocalMatrixRows<SC, LO, GO, NT>;
411 
412  // Only run on host if we can access the data
413  const bool runOnHost = Kokkos::SpaceAccessibility<Kokkos::Serial, memory_space>::accessible;
414  if (runOnHost) {
415  using local_row_indices_type = Kokkos::View<const LO*, Kokkos::AnonymousSpace>;
416  const local_row_indices_type lclRowInds_a(lclRowInds);
417  impl_type::run(execution_space(), A, lclRowInds_a, true);
418  } else {
419  auto lclRowInds_a = Kokkos::create_mirror_view_and_copy(execution_space(), lclRowInds);
420  impl_type::run(execution_space(), A, lclRowInds_a, false);
421  }
422 }
423 
424 template <class CrsMatrixType>
426  const Kokkos::View<
427  typename CrsMatrixType::local_ordinal_type*,
428  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;
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  TEUCHOS_TEST_FOR_EXCEPTION(A.getColMap().get() == nullptr, std::invalid_argument, "The matrix must have a column Map.");
438 
439  // Copy the Host array to device
440  auto lclRowInds_d = Kokkos::create_mirror_view_and_copy(execution_space(), lclRowInds);
441 
442  Kokkos::View<bool*, memory_space> dirichletColFlags("dirichletColFlags", A.getColMap()->getLocalNumElements());
443  Kokkos::View<bool*, Kokkos::AnonymousSpace> dirichletColFlags_a(dirichletColFlags);
444  Details::localRowsToColumns<SC, LO, GO, NT>(execution_space(), A, lclRowInds_d, dirichletColFlags_a);
445 
446  Details::ApplyDirichletBoundaryConditionToLocalMatrixColumns<SC, LO, GO, NT>::run(execution_space(), A, dirichletColFlags, false);
447  Details::ApplyDirichletBoundaryConditionToLocalMatrixRows<SC, LO, GO, NT>::run(execution_space(), A, lclRowInds_d, false);
448 }
449 
450 template <class CrsMatrixType>
452  const Kokkos::View<
453  typename CrsMatrixType::local_ordinal_type*,
454  typename CrsMatrixType::device_type>& lclRowInds_d) {
455  using SC = typename CrsMatrixType::scalar_type;
456  using LO = typename CrsMatrixType::local_ordinal_type;
457  using GO = typename CrsMatrixType::global_ordinal_type;
458  using NT = typename CrsMatrixType::node_type;
459  using crs_matrix_type = ::Tpetra::CrsMatrix<SC, LO, GO, NT>;
460  using execution_space = typename crs_matrix_type::execution_space;
461  using memory_space = typename crs_matrix_type::device_type::memory_space;
462 
463  TEUCHOS_TEST_FOR_EXCEPTION(A.getColMap().get() == nullptr, std::invalid_argument, "The matrix must have a column Map.");
464 
465  Kokkos::View<bool*, memory_space> dirichletColFlags("dirichletColFlags", A.getColMap()->getLocalNumElements());
466  Kokkos::View<bool*, Kokkos::AnonymousSpace> dirichletColFlags_a(dirichletColFlags);
467  Details::localRowsToColumns<SC, LO, GO, NT>(execution_space(), A, lclRowInds_d, dirichletColFlags_a);
468 
469  Details::ApplyDirichletBoundaryConditionToLocalMatrixColumns<SC, LO, GO, NT>::run(execution_space(), A, dirichletColFlags, false);
470  Details::ApplyDirichletBoundaryConditionToLocalMatrixRows<SC, LO, GO, NT>::run(execution_space(), A, lclRowInds_d, false);
471 }
472 
473 template <class CrsMatrixType>
474 void applyDirichletBoundaryConditionToLocalMatrixRowsAndColumns(const typename CrsMatrixType::execution_space& execSpace,
475  CrsMatrixType& A,
476  const Kokkos::View<
477  typename CrsMatrixType::local_ordinal_type*,
478  typename CrsMatrixType::device_type>& lclRowInds_d) {
479  using SC = typename CrsMatrixType::scalar_type;
480  using LO = typename CrsMatrixType::local_ordinal_type;
481  using GO = typename CrsMatrixType::global_ordinal_type;
482  using NT = typename CrsMatrixType::node_type;
483  using crs_matrix_type = ::Tpetra::CrsMatrix<SC, LO, GO, NT>;
484  // using execution_space = typename crs_matrix_type::execution_space;
485  using memory_space = typename crs_matrix_type::device_type::memory_space;
486 
487  TEUCHOS_TEST_FOR_EXCEPTION(A.getColMap().get() == nullptr, std::invalid_argument, "The matrix must have a column Map.");
488 
489  Kokkos::View<bool*, memory_space> dirichletColFlags("dirichletColFlags", A.getColMap()->getLocalNumElements());
490  Kokkos::View<bool*, Kokkos::AnonymousSpace> dirichletColFlags_a(dirichletColFlags);
491  Details::localRowsToColumns<SC, LO, GO, NT>(execSpace, A, lclRowInds_d, dirichletColFlags_a);
492 
493  Details::ApplyDirichletBoundaryConditionToLocalMatrixColumns<SC, LO, GO, NT>::run(execSpace, A, dirichletColFlags, false);
494  Details::ApplyDirichletBoundaryConditionToLocalMatrixRows<SC, LO, GO, NT>::run(execSpace, A, lclRowInds_d, false);
495 }
496 
497 } // namespace Tpetra
498 
499 #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...