Teko Version of the Day
Loading...
Searching...
No Matches
Teko_ProbingPreconditionerFactory.cpp
1// @HEADER
2// *****************************************************************************
3// Teko: A package for block and physics based preconditioning
4//
5// Copyright 2010 NTESS and the Teko contributors.
6// SPDX-License-Identifier: BSD-3-Clause
7// *****************************************************************************
8// @HEADER
9
10#include "Teko_ProbingPreconditionerFactory.hpp"
11
12#include "Teko_TpetraOperatorWrapper.hpp"
13#include "Teko_PreconditionerInverseFactory.hpp"
14#include "Teko_RequestHandler.hpp"
15#include "Teko_RequestMesg.hpp"
16
17#include "Thyra_TpetraLinearOp.hpp"
18#include "Thyra_DefaultLinearOpSource.hpp"
19#include "Thyra_TpetraThyraWrappers.hpp"
20
21#include "KokkosGraph_Distance2ColorHandle.hpp"
22#include "KokkosGraph_Distance2Color.hpp"
23#include "KokkosKernels_Handle.hpp"
24
25#include <Kokkos_Sort.hpp>
26#include <cstdint>
27#include <vector>
28
29using Teuchos::rcp;
30using Teuchos::RCP;
31using Teuchos::rcp_dynamic_cast;
32
33namespace {
34
35using ST = Teko::ST;
36using LO = Teko::LO;
37using GO = Teko::GO;
38using NT = Teko::NT;
39
40using exec_space = typename NT::device_type::execution_space;
41
42template <class RowMapType, class EntriesType>
43struct LocalCrsGraphViews {
44 RowMapType row_map;
45 EntriesType entries;
46};
47
48template <class MVType, class HostColorViewType>
49void set_probe_by_color_host(const Teuchos::RCP<MVType>& probeVec,
50 const HostColorViewType& h_colors,
51 const std::vector<LO>& colLidToDomainLid, const LO color) {
52 const LO invalid = Teuchos::OrdinalTraits<LO>::invalid();
53
54 probeVec->putScalar(Teuchos::ScalarTraits<ST>::zero());
55
56 auto hostView = probeVec->getLocalViewHost(Tpetra::Access::ReadWrite);
57 const LO numLocalCols = static_cast<LO>(colLidToDomainLid.size());
58
59 for (LO lcol = 0; lcol < numLocalCols; ++lcol) {
60 if (static_cast<LO>(h_colors(lcol)) == color) {
61 const LO domainLid = colLidToDomainLid[static_cast<size_t>(lcol)];
62 if (domainLid != invalid) {
63 hostView(domainLid, 0) = Teuchos::ScalarTraits<ST>::one();
64 }
65 }
66 }
67}
68
69template <class RowPtrViewType, class ColIndViewType, class ColorViewType, class ResponseViewType,
70 class ValuesViewType>
71void decode_probe_by_color(const RowPtrViewType& rowPtrs, const ColIndViewType& localColInds,
72 const ColorViewType& colors, const ResponseViewType& responseView,
73 const LO numRows, const LO color, const ValuesViewType& values) {
74 Kokkos::parallel_for(
75 "Teko::ProbingPreconditionerFactory::decode_probe",
76 Kokkos::RangePolicy<exec_space>(0, numRows), KOKKOS_LAMBDA(const LO lrow) {
77 const auto rowStart = rowPtrs(lrow);
78 const auto rowEnd = rowPtrs(lrow + 1);
79 const ST rowValue = responseView(lrow, 0);
80
81 for (auto entry = rowStart; entry < rowEnd; ++entry) {
82 const LO lcol = static_cast<LO>(localColInds(entry));
83 if (colors(lcol) == color) {
84 values(entry) = rowValue;
85 }
86 }
87 });
88}
89
90template <class RowPtrViewType>
91LO compute_max_row_length(const RowPtrViewType& rowPtrs, const LO numRows) {
92 LO maxNumEntriesPerRow = 0;
93
94 Kokkos::parallel_reduce(
95 "Teko::ProbingPreconditionerFactory::max_row_length",
96 Kokkos::RangePolicy<exec_space>(0, numRows),
97 KOKKOS_LAMBDA(const LO r, LO& localMax) {
98 const LO rowLen = static_cast<LO>(rowPtrs(r + 1) - rowPtrs(r));
99 if (rowLen > localMax) localMax = rowLen;
100 },
101 Kokkos::Max<LO>(maxNumEntriesPerRow));
102
103 return maxNumEntriesPerRow;
104}
105
106template <class RowMapViewType>
107typename RowMapViewType::non_const_value_type count_candidate_col_graph_edges(
108 const RowMapViewType& row_map_copy, const LO numRows) {
109 using row_map_value_type = typename RowMapViewType::non_const_value_type;
110
111 row_map_value_type candidateEdgeCount = 0;
112 Kokkos::parallel_reduce(
113 "Teko::ProbingPreconditionerFactory::count_candidate_col_graph_edges",
114 Kokkos::RangePolicy<exec_space>(0, numRows),
115 KOKKOS_LAMBDA(const LO row, row_map_value_type& update) {
116 const row_map_value_type d = row_map_copy(row + 1) - row_map_copy(row);
117 if (d > 1) update += d * (d - 1);
118 },
119 candidateEdgeCount);
120
121 return candidateEdgeCount;
122}
123
124template <class RowMapDeviceView, class EntriesDeviceView>
125LocalCrsGraphViews<Kokkos::View<typename RowMapDeviceView::non_const_value_type*,
126 typename RowMapDeviceView::device_type>,
127 Kokkos::View<typename EntriesDeviceView::non_const_value_type*,
128 typename EntriesDeviceView::device_type> >
129build_column_intersection_graph_device(const RowMapDeviceView& row_map_in,
130 const EntriesDeviceView& entries_in, const LO numRows,
131 const LO numLocalCols) {
132 using row_offset_type = typename RowMapDeviceView::non_const_value_type;
133 using col_index_type = typename EntriesDeviceView::non_const_value_type;
134 using device_type = typename RowMapDeviceView::device_type;
135 using key_type = std::uint64_t;
136
137 using row_map_out_type = Kokkos::View<row_offset_type*, device_type>;
138 using entries_out_type = Kokkos::View<col_index_type*, device_type>;
139
140 const row_offset_type numCandidateKeys = count_candidate_col_graph_edges(row_map_in, numRows);
141
142 Kokkos::View<row_offset_type*, device_type> candidateOffsets(
143 Kokkos::ViewAllocateWithoutInitializing("teko_probe_candidate_offsets"),
144 static_cast<size_t>(numRows) + 1);
145
146 Kokkos::parallel_for(
147 "Teko::ProbingPreconditionerFactory::fill_candidate_offsets",
148 Kokkos::RangePolicy<exec_space>(0, numRows), KOKKOS_LAMBDA(const LO row) {
149 const row_offset_type d = row_map_in(row + 1) - row_map_in(row);
150 candidateOffsets(row) = (d > 1 ? d * (d - 1) : 0);
151 });
152
153 Kokkos::parallel_scan(
154 "Teko::ProbingPreconditionerFactory::scan_candidate_offsets",
155 Kokkos::RangePolicy<exec_space>(0, numRows + 1),
156 KOKKOS_LAMBDA(const LO i, row_offset_type& update, const bool final_pass) {
157 const row_offset_type val = (i < numRows ? candidateOffsets(i) : row_offset_type(0));
158 if (final_pass) candidateOffsets(i) = update;
159 update += val;
160 });
161
162 Kokkos::View<key_type*, device_type> candidateKeys(
163 Kokkos::ViewAllocateWithoutInitializing("teko_probe_candidate_keys"),
164 static_cast<size_t>(numCandidateKeys));
165
166 const key_type numColsKey = static_cast<key_type>(numLocalCols);
167
168 Kokkos::parallel_for(
169 "Teko::ProbingPreconditionerFactory::fill_candidate_keys",
170 Kokkos::RangePolicy<exec_space>(0, numRows), KOKKOS_LAMBDA(const LO row) {
171 const auto rowStart = row_map_in(row);
172 const auto rowEnd = row_map_in(row + 1);
173 const row_offset_type d = rowEnd - rowStart;
174 row_offset_type pos = candidateOffsets(row);
175
176 for (row_offset_type i = 0; i < d; ++i) {
177 const key_type src = static_cast<key_type>(entries_in(rowStart + i));
178 for (row_offset_type j = 0; j < d; ++j) {
179 if (i == j) continue;
180 const key_type dst = static_cast<key_type>(entries_in(rowStart + j));
181 candidateKeys(pos++) = src * numColsKey + dst;
182 }
183 }
184 });
185
186 Kokkos::sort(candidateKeys);
187
188 Kokkos::View<row_offset_type*, device_type> uniqueFlags(
189 Kokkos::ViewAllocateWithoutInitializing("teko_probe_unique_flags"),
190 static_cast<size_t>(numCandidateKeys));
191
192 Kokkos::parallel_for(
193 "Teko::ProbingPreconditionerFactory::mark_unique_candidate_keys",
194 Kokkos::RangePolicy<exec_space>(0, numCandidateKeys), KOKKOS_LAMBDA(const row_offset_type i) {
195 if (i == 0) {
196 uniqueFlags(i) = 1;
197 } else {
198 uniqueFlags(i) = (candidateKeys(i) != candidateKeys(i - 1) ? 1 : 0);
199 }
200 });
201
202 Kokkos::View<row_offset_type*, device_type> uniqueOffsets(
203 Kokkos::ViewAllocateWithoutInitializing("teko_probe_unique_offsets"),
204 static_cast<size_t>(numCandidateKeys) + 1);
205
206 Kokkos::parallel_scan(
207 "Teko::ProbingPreconditionerFactory::scan_unique_flags",
208 Kokkos::RangePolicy<exec_space>(0, numCandidateKeys + 1),
209 KOKKOS_LAMBDA(const row_offset_type i, row_offset_type& update, const bool final_pass) {
210 const row_offset_type val = (i < numCandidateKeys ? uniqueFlags(i) : row_offset_type(0));
211 if (final_pass) uniqueOffsets(i) = update;
212 update += val;
213 });
214
215 auto h_uniqueCount = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), uniqueOffsets);
216 const row_offset_type numUniqueKeys = h_uniqueCount(numCandidateKeys);
217
218 Kokkos::View<key_type*, device_type> uniqueKeys(
219 Kokkos::ViewAllocateWithoutInitializing("teko_probe_unique_keys"),
220 static_cast<size_t>(numUniqueKeys));
221
222 Kokkos::parallel_for(
223 "Teko::ProbingPreconditionerFactory::pack_unique_keys",
224 Kokkos::RangePolicy<exec_space>(0, numCandidateKeys), KOKKOS_LAMBDA(const row_offset_type i) {
225 if (uniqueFlags(i)) {
226 uniqueKeys(uniqueOffsets(i)) = candidateKeys(i);
227 }
228 });
229
230 Kokkos::View<row_offset_type*, device_type> rowCounts(
231 Kokkos::ViewAllocateWithoutInitializing("teko_probe_row_counts"),
232 static_cast<size_t>(numLocalCols));
233 Kokkos::deep_copy(rowCounts, row_offset_type(0));
234
235 Kokkos::parallel_for(
236 "Teko::ProbingPreconditionerFactory::count_unique_keys_per_row",
237 Kokkos::RangePolicy<exec_space>(0, numUniqueKeys), KOKKOS_LAMBDA(const row_offset_type i) {
238 const key_type key = uniqueKeys(i);
239 const row_offset_type src = static_cast<row_offset_type>(key / numColsKey);
240 Kokkos::atomic_fetch_add(&rowCounts(src), row_offset_type(1));
241 });
242
243 row_map_out_type row_map_out(
244 Kokkos::ViewAllocateWithoutInitializing("teko_probe_col_graph_row_map_dev"),
245 static_cast<size_t>(numLocalCols) + 1);
246
247 Kokkos::parallel_scan(
248 "Teko::ProbingPreconditionerFactory::scan_col_graph_row_map",
249 Kokkos::RangePolicy<exec_space>(0, numLocalCols + 1),
250 KOKKOS_LAMBDA(const LO i, row_offset_type& update, const bool final_pass) {
251 const row_offset_type val = (i < numLocalCols ? rowCounts(i) : row_offset_type(0));
252 if (final_pass) row_map_out(i) = update;
253 update += val;
254 });
255
256 entries_out_type entries_out(
257 Kokkos::ViewAllocateWithoutInitializing("teko_probe_col_graph_entries_dev"),
258 static_cast<size_t>(numUniqueKeys));
259
260 Kokkos::View<row_offset_type*, device_type> fillOffsets(
261 Kokkos::ViewAllocateWithoutInitializing("teko_probe_fill_offsets"),
262 static_cast<size_t>(numLocalCols));
263
264 Kokkos::parallel_for(
265 "Teko::ProbingPreconditionerFactory::init_fill_offsets",
266 Kokkos::RangePolicy<exec_space>(0, numLocalCols),
267 KOKKOS_LAMBDA(const LO col) { fillOffsets(col) = row_map_out(col); });
268
269 Kokkos::parallel_for(
270 "Teko::ProbingPreconditionerFactory::fill_col_graph_entries",
271 Kokkos::RangePolicy<exec_space>(0, numUniqueKeys), KOKKOS_LAMBDA(const row_offset_type i) {
272 const key_type key = uniqueKeys(i);
273 const row_offset_type src = static_cast<row_offset_type>(key / numColsKey);
274 const col_index_type dst = static_cast<col_index_type>(key % numColsKey);
275 const row_offset_type pos = Kokkos::atomic_fetch_add(&fillOffsets(src), row_offset_type(1));
276 entries_out(pos) = dst;
277 });
278
279 LocalCrsGraphViews<row_map_out_type, entries_out_type> out = {row_map_out, entries_out};
280 return out;
281}
282
283} // namespace
284
285namespace Teko {
286
287ProbingPreconditionerFactory::ProbingPreconditionerFactory() {}
288
289LinearOp ProbingPreconditionerFactory::buildPreconditionerOperator(
290 LinearOp& lo, PreconditionerState& state) const {
291 RCP<Tpetra::CrsMatrix<ST, LO, GO, NT> > probedMat = probe(lo);
292
293 LinearOp probedOp = Thyra::tpetraLinearOp<ST, LO, GO, NT>(
294 Thyra::tpetraVectorSpace<ST, LO, GO, NT>(probedMat->getDomainMap()),
295 Thyra::tpetraVectorSpace<ST, LO, GO, NT>(probedMat->getRangeMap()), probedMat);
296
297 return Teko::buildInverse(*invFactory_, probedOp);
298}
299
300void ProbingPreconditionerFactory::initializeFromParameterList(const Teuchos::ParameterList& pl) {
301 RCP<const InverseLibrary> invLib = getInverseLibrary();
302
303 const std::string inverse_type = "Inverse Type";
304 const std::string probing_graph_operator = "Probing Graph Operator";
305 const std::string probing_graph = "Probing Graph";
306 const std::string user_graph = "User Will Set Probing Graph";
307
308 std::string invStr = "Ifpack2";
309 if (pl.isParameter(inverse_type)) invStr = pl.get<std::string>(inverse_type);
310
311 if (pl.isParameter(probing_graph_operator))
312 setGraphOperator(pl.get<Teko::LinearOp>(probing_graph_operator));
313 else if (pl.isParameter(probing_graph))
314 setGraph(pl.get<RCP<const Tpetra::CrsGraph<LO, GO, NT> > >(probing_graph));
315 else if (pl.isParameter(user_graph) && pl.get<bool>(user_graph)) {
316 // noop
317 } else {
318 Teuchos::RCP<Teko::RequestHandler> rh = getRequestHandler();
319 rh->preRequest<RCP<const Tpetra::CrsGraph<LO, GO, NT> > >(Teko::RequestMesg("Probing Graph"));
320 setGraph(
321 rh->request<RCP<const Tpetra::CrsGraph<LO, GO, NT> > >(Teko::RequestMesg("Probing Graph")));
322 }
323
324 setInverseFactory(invLib->getInverseFactory(invStr));
325}
326
327void ProbingPreconditionerFactory::setGraphOperator(const Teko::LinearOp& graphOp) {
328 RCP<const Thyra::TpetraLinearOp<ST, LO, GO, NT> > tOp =
329 rcp_dynamic_cast<const Thyra::TpetraLinearOp<ST, LO, GO, NT> >(graphOp, true);
330 RCP<const Tpetra::CrsMatrix<ST, LO, GO, NT> > crsMatrix =
331 rcp_dynamic_cast<const Tpetra::CrsMatrix<ST, LO, GO, NT> >(tOp->getConstTpetraOperator(),
332 true);
333 setGraph(crsMatrix->getCrsGraph());
334}
335
336void ProbingPreconditionerFactory::setGraph(
337 const Teuchos::RCP<const Tpetra::CrsGraph<LO, GO, NT> >& graph) {
338 graph_ = graph;
339}
340
341RCP<Tpetra::CrsMatrix<ST, LO, GO, NT> > ProbingPreconditionerFactory::probe(
342 const LinearOp& lo) const {
343 TEUCHOS_TEST_FOR_EXCEPTION(graph_ == Teuchos::null, std::runtime_error,
344 "ProbingPreconditionerFactory::probe: probing graph is null");
345
346 RCP<const Thyra::LinearOpBase<ST> > thyraOp = lo;
347 RCP<const Thyra::TpetraLinearOp<ST, LO, GO, NT> > tOp =
348 rcp_dynamic_cast<const Thyra::TpetraLinearOp<ST, LO, GO, NT> >(thyraOp, true);
349 RCP<const Tpetra::Operator<ST, LO, GO, NT> > tpetraOp = tOp->getConstTpetraOperator();
350
351 RCP<const Tpetra::Map<LO, GO, NT> > domainMap = tpetraOp->getDomainMap();
352 RCP<const Tpetra::Map<LO, GO, NT> > rangeMap = tpetraOp->getRangeMap();
353 RCP<const Tpetra::Map<LO, GO, NT> > rowMap = graph_->getRowMap();
354 RCP<const Tpetra::Map<LO, GO, NT> > colMap = graph_->getColMap();
355
356 using local_graph_device_type = typename Tpetra::CrsGraph<LO, GO, NT>::local_graph_device_type;
357 using row_map_type = typename local_graph_device_type::row_map_type;
358 using index_type = typename local_graph_device_type::entries_type;
359 using local_matrix_type = typename Tpetra::CrsMatrix<ST, LO, GO, NT>::local_matrix_device_type;
360 using values_type = typename local_matrix_type::values_type::non_const_type;
361 using exec_space = typename NT::device_type::execution_space;
362 using memory_space = typename NT::device_type::memory_space;
363 using row_map_nonc_type = typename row_map_type::non_const_type;
364 using entries_nonc_type = typename index_type::non_const_type;
365
366 using kernel_handle_type =
367 KokkosKernels::Experimental::KokkosKernelsHandle<typename row_map_type::value_type,
368 typename index_type::value_type, ST,
369 exec_space, memory_space, memory_space>;
370
371 const LO numRows = static_cast<LO>(graph_->getLocalNumRows());
372 const LO numLocalCols = static_cast<LO>(colMap->getLocalNumElements());
373
374 row_map_nonc_type row_map_copy;
375 entries_nonc_type entries_copy;
376 LO totalNumEntries = 0;
377 LO maxNumEntriesPerRow = 0;
378
379 {
380 auto localGraphDevice = graph_->getLocalGraphDevice();
381
382 row_map_copy =
383 row_map_nonc_type(Kokkos::ViewAllocateWithoutInitializing("teko_probe_row_map_copy"),
384 localGraphDevice.row_map.extent(0));
385 entries_copy =
386 entries_nonc_type(Kokkos::ViewAllocateWithoutInitializing("teko_probe_entries_copy"),
387 localGraphDevice.entries.extent(0));
388
389 Kokkos::deep_copy(row_map_copy, localGraphDevice.row_map);
390 Kokkos::deep_copy(entries_copy, localGraphDevice.entries);
391
392 totalNumEntries = static_cast<LO>(entries_copy.extent(0));
393 maxNumEntriesPerRow = compute_max_row_length(row_map_copy, numRows);
394 }
395
396 auto colGraph =
397 build_column_intersection_graph_device(row_map_copy, entries_copy, numRows, numLocalCols);
398
399 kernel_handle_type kh;
400 kh.create_distance2_graph_coloring_handle(KokkosGraph::COLORING_D2_DEFAULT);
401
402 KokkosGraph::Experimental::graph_color_distance2(&kh, numLocalCols, colGraph.row_map,
403 colGraph.entries);
404
405 auto coloringHandle = kh.get_distance2_graph_coloring_handle();
406 auto colors = coloringHandle->get_vertex_colors();
407 const LO numColors = static_cast<LO>(coloringHandle->get_num_colors());
408
409 TEUCHOS_TEST_FOR_EXCEPTION(static_cast<LO>(colors.extent(0)) < numLocalCols, std::runtime_error,
410 "ProbingPreconditionerFactory::probe: colors view extent ("
411 << colors.extent(0) << ") is smaller than local column count ("
412 << numLocalCols << ").");
413
414 std::vector<LO> colLidToDomainLid(static_cast<size_t>(numLocalCols));
415 {
416 const LO invalid = Teuchos::OrdinalTraits<LO>::invalid();
417 for (LO lcol = 0; lcol < numLocalCols; ++lcol) {
418 const GO gid = colMap->getGlobalElement(lcol);
419 const LO dlid =
420 (gid == Teuchos::OrdinalTraits<GO>::invalid() ? invalid
421 : domainMap->getLocalElement(gid));
422 colLidToDomainLid[static_cast<size_t>(lcol)] = dlid;
423 }
424 }
425
426 values_type values(Kokkos::ViewAllocateWithoutInitializing("probed_values"), totalNumEntries);
427 Kokkos::deep_copy(values, Teuchos::ScalarTraits<ST>::zero());
428
429 auto h_colors = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), colors);
430
431 RCP<Tpetra::MultiVector<ST, LO, GO, NT> > probeVec =
432 rcp(new Tpetra::MultiVector<ST, LO, GO, NT>(domainMap, 1));
433 RCP<Tpetra::MultiVector<ST, LO, GO, NT> > response =
434 rcp(new Tpetra::MultiVector<ST, LO, GO, NT>(rangeMap, 1));
435
436 for (LO color = 1; color <= numColors; ++color) {
437 set_probe_by_color_host(probeVec, h_colors, colLidToDomainLid, color);
438
439 tpetraOp->apply(*probeVec, *response, Teuchos::NO_TRANS, Teuchos::ScalarTraits<ST>::one(),
440 Teuchos::ScalarTraits<ST>::zero());
441
442 auto responseView = response->getLocalViewDevice(Tpetra::Access::ReadOnly);
443 decode_probe_by_color(row_map_copy, entries_copy, colors, responseView, numRows, color, values);
444 }
445
446 kh.destroy_graph_coloring_handle();
447
448 auto lclMat = local_matrix_type("probed_local_matrix", numRows, maxNumEntriesPerRow,
449 totalNumEntries, values, row_map_copy, entries_copy);
450
451 return Teuchos::rcp(
452 new Tpetra::CrsMatrix<ST, LO, GO, NT>(lclMat, rowMap, colMap, domainMap, rangeMap));
453}
454
455} // namespace Teko
Teko::buildInverse
InverseLinearOp buildInverse(const InverseFactory &factory, const LinearOp &A)
Build an inverse operator using a factory and a linear operator.
Definition Teko_InverseFactory.cpp:81
Teko::PreconditionerFactory::getInverseLibrary
Teuchos::RCP< const InverseLibrary > getInverseLibrary() const
Get the inverse library used by this preconditioner factory.
Definition Teko_PreconditionerFactory.cpp:129
Teko::PreconditionerFactory::getRequestHandler
Teuchos::RCP< RequestHandler > getRequestHandler() const
Get the request handler with pointers to the appropriate callbacks.
Definition Teko_PreconditionerFactory.hpp:189
Teko::PreconditionerState
An implementation of a state object preconditioners.
Definition Teko_PreconditionerState.hpp:34
Teko::ProbingPreconditionerFactory::buildPreconditionerOperator
LinearOp buildPreconditionerOperator(LinearOp &lo, PreconditionerState &state) const override
Definition Teko_ProbingPreconditionerFactory.cpp:289
Teko::ProbingPreconditionerFactory::initializeFromParameterList
void initializeFromParameterList(const Teuchos::ParameterList &pl) override
Initialize from a parameter list.
Definition Teko_ProbingPreconditionerFactory.cpp:300
Generated by doxygen 1.9.8