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