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