Tpetra parallel linear algebra Version of the Day
Loading...
Searching...
No Matches
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
27namespace Tpetra {
28
38template <class CrsMatrixType>
39void 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
54template <class CrsMatrixType>
56 const Kokkos::View<
57 typename CrsMatrixType::local_ordinal_type*,
58 typename CrsMatrixType::device_type>& lclRowInds);
59
69template <class CrsMatrixType>
71 const Kokkos::View<
72 typename CrsMatrixType::local_ordinal_type*,
73 Kokkos::HostSpace>& lclRowInds);
74
84template <class CrsMatrixType>
85void 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
101template <class CrsMatrixType>
103 const Kokkos::View<
104 typename CrsMatrixType::local_ordinal_type*,
105 Kokkos::HostSpace>& lclRowInds);
106
116template <class CrsMatrixType>
118 const Kokkos::View<
119 typename CrsMatrixType::local_ordinal_type*,
120 typename CrsMatrixType::device_type>& lclRowInds);
121
122namespace Details {
123
124template <class SC, class LO, class GO, class NT>
125struct 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
210template <class SC, class LO, class GO, class NT>
211struct 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
291template <class SC, class LO, class GO, class NT>
292void 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
364template <class CrsMatrixType>
365void applyDirichletBoundaryConditionToLocalMatrixRows(const typename CrsMatrixType::execution_space& execSpace,
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 =
382 const bool runOnHost = false;
383 impl_type::run(execSpace, A, lclRowInds_a, runOnHost);
384}
385
386template <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;
393}
394
395template <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 =
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
424template <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
450template <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
473template <class CrsMatrixType>
474void applyDirichletBoundaryConditionToLocalMatrixRowsAndColumns(const typename CrsMatrixType::execution_space& execSpace,
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
Struct that holds views of the contents of a CrsMatrix.
typename device_type::execution_space execution_space
The Kokkos execution space.
typename row_matrix_type::impl_scalar_type impl_scalar_type
The type used internally in place of Scalar.
Implementation details of Tpetra.
Namespace Tpetra contains the class and methods constituting the Tpetra library.
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...
@ INSERT
Insert new values that don't currently exist.
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 ...