Zoltan2
Loading...
Searching...
No Matches
Zoltan2_TpetraCrsColorer.hpp
Go to the documentation of this file.
1#pragma once
2
3#include <stdexcept>
4
5#include "Teuchos_RCP.hpp"
6#include "Teuchos_ParameterList.hpp"
7#include "Teuchos_TestForException.hpp"
8
9#include "Tpetra_Map.hpp"
10#include "Tpetra_MultiVector.hpp"
11#include "Tpetra_CrsMatrix.hpp"
12#include "Tpetra_CrsGraph.hpp"
13#include "Tpetra_BlockMultiVector.hpp"
14#include "Tpetra_BlockCrsMatrix.hpp"
15
19
20namespace Zoltan2
21{
22
23// Base class for coloring Tpetra::CrsMatrix for use in column compression
24template <typename CrsMatrixType>
26{
27public:
28 typedef CrsMatrixType matrix_t;
29 typedef typename matrix_t::crs_graph_type graph_t;
30 typedef typename matrix_t::scalar_type scalar_t;
31 typedef typename matrix_t::local_ordinal_type lno_t;
32 typedef typename matrix_t::global_ordinal_type gno_t;
33 typedef typename matrix_t::node_type node_t;
34 typedef typename node_t::device_type device_t;
35 typedef typename device_t::execution_space execution_space;
36 typedef Kokkos::View<int *, device_t> list_of_colors_t;
37 typedef typename list_of_colors_t::HostMirror list_of_colors_host_t;
38
39 // Constructor
40 TpetraCrsColorer(const Teuchos::RCP<matrix_t> &matrix_);
41
42 // Destructor
44
45 // Compute coloring data
46 void
47 computeColoring(Teuchos::ParameterList &coloring_params);
48
49 // Compute seed matrix
50 template <typename MultiVectorType>
51 void
52 computeSeedMatrix(MultiVectorType &V, const int color_beg = 0) const;
53
54 // Compute seed matrix with distribution fitted to the graph's column map
55 template <typename MultiVectorType>
56 void
57 computeSeedMatrixFitted(MultiVectorType &V, const int color_beg = 0) const;
58
59 // Reconstruct matrix from supplied compressed vector
60 template <typename MultiVectorType>
61 void
62 reconstructMatrix(MultiVectorType &W, const int color_beg = 0) const;
63 template <typename MultiVectorType>
64 void
65 reconstructMatrix(MultiVectorType &W, matrix_t &mat, const int color_beg = 0) const;
66
67 // Reconstruct matrix from supplied compressed vector fitted to the graph's
68 // row map
69 template <typename MultiVectorType>
70 void
71 reconstructMatrixFitted(MultiVectorType &W, const int color_beg = 0) const;
72 template <typename MultiVectorType>
73 void
74 reconstructMatrixFitted(MultiVectorType &W, matrix_t &mat, const int color_beg = 0) const;
75
76 // Return number of colors
77 // KDD should num_colors be int or size_t?
78 int
80 {
81 return num_colors;
82 }
83
84 // Get color given a column index
85 int
86 getColor(const size_t col) const
87 {
88 return list_of_colors_host(col) - 1;
89 }
90
91 // Check coloring is valid, i.e., no row has two nonzero columns with
92 // same color
93 bool
95 {
97 }
98
99protected:
100 Teuchos::RCP<matrix_t> matrix;
101 Teuchos::RCP<const graph_t> graph;
105};
106
108// Specialization of TpetraCrsColorer for BlockCrsMatrix
109template <typename SC, typename LO, typename GO, typename NO>
110class TpetraCrsColorer<Tpetra::BlockCrsMatrix<SC,LO,GO,NO> >
111{
112public:
113 typedef Tpetra::BlockCrsMatrix<SC, LO, GO, NO> matrix_t;
114 typedef Tpetra::BlockMultiVector<SC, LO, GO, NO> multivector_t;
115 typedef typename matrix_t::crs_graph_type graph_t;
116 typedef typename matrix_t::scalar_type scalar_t;
117 typedef typename matrix_t::local_ordinal_type lno_t;
118 typedef typename matrix_t::global_ordinal_type gno_t;
119 typedef typename matrix_t::node_type node_t;
120 typedef typename node_t::device_type device_t;
121 typedef typename device_t::execution_space execution_space;
122 typedef Kokkos::View<int *, device_t> list_of_colors_t;
123 typedef typename list_of_colors_t::HostMirror list_of_colors_host_t;
124
125 // Constructor
126 TpetraCrsColorer(const Teuchos::RCP<matrix_t> &matrix_);
127
128 // Destructor
130
131 // Compute coloring data
132 void
133 computeColoring(Teuchos::ParameterList &coloring_params);
134
135 // Compute seed matrix
136 void
137 computeSeedMatrix(multivector_t &V, const int color_beg = 0) const;
138
139 // Compute seed matrix with distribution fitted to the graph's column map
140 void
141 computeSeedMatrixFitted(multivector_t &V, const int color_beg = 0) const;
142
143 // Reconstruct matrix from supplied compressed vector
144 void
145 reconstructMatrix(multivector_t &W, const int color_beg = 0) const;
146 void
147 reconstructMatrix(multivector_t &W, matrix_t &mat, const int color_beg = 0) const;
148
149 // Reconstruct matrix from supplied compressed vector fitted to the graph's
150 // row map
151 void
152 reconstructMatrixFitted(multivector_t &W, const int color_beg = 0) const;
153 void
154 reconstructMatrixFitted(multivector_t &W, matrix_t &mat, const int color_beg = 0) const;
155
156 // Return number of colors
157 int
159 {
160 return num_colors * matrix->getBlockSize();
161 }
162
163 // Get color given a column index
164 int
165 getColor(const size_t col) const
166 {
167 const int block_size = matrix->getBlockSize();
168 const size_t block_col = col / block_size;
169 const int offset = col - block_col * block_size;
170 return (list_of_colors_host(block_col) - 1) * block_size + offset;
171 }
172
173 // Check coloring is valid, i.e., no row has two nonzero columns with
174 // same color
175 bool
177 {
179 }
180
181protected:
182 Teuchos::RCP<matrix_t> matrix;
183 Teuchos::RCP<const graph_t> graph;
187};
188
190
191template <typename CrsMatrixType>
193 const Teuchos::RCP<matrix_t> &matrix_
194)
195 : matrix(matrix_),
196 graph(matrix->getCrsGraph()),
197 list_of_colors(),
198 list_of_colors_host(),
199 num_colors(0)
200{
201
202}
203
205template <typename CrsMatrixType>
206void
208 Teuchos::ParameterList &coloring_params
209)
210{
211 const std::string library = coloring_params.get("library", "zoltan");
212
213 if (library == "zoltan") {
214 // Use Zoltan's coloring
215 ZoltanCrsColorer<matrix_t> zz(matrix);
216 zz.computeColoring(coloring_params,
217 num_colors, list_of_colors_host, list_of_colors);
218 }
219 else {
220 // Use Zoltan2's coloring
221 Zoltan2CrsColorer<matrix_t> zz2(matrix);
222 zz2.computeColoring(coloring_params,
223 num_colors, list_of_colors_host, list_of_colors);
224 }
225}
226
228template <typename CrsMatrixType>
229template <typename MultiVectorType>
230void
232 MultiVectorType &V, const int color_beg) const
233{
234 MultiVectorType V_fitted(graph->getColMap(), V.getNumVectors());
235
236 computeSeedMatrixFitted(V_fitted, color_beg);
237
238 Tpetra::Import<lno_t, gno_t, node_t>
239 importer(graph->getColMap(), V.getMap());
240
241 V.doImport(V_fitted, importer, Tpetra::INSERT);
242}
243
245template <typename CrsMatrixType>
246template <typename MultiVectorType>
247void
249 MultiVectorType &V, const int color_beg) const
250{
251 // Check V's map is locally fitted to the graph's column map
252 TEUCHOS_TEST_FOR_EXCEPTION(
253 !(graph->getColMap()->isLocallyFitted(*(V.getMap()))), std::domain_error,
254 "Map for supplied vector is not locally fitted to the column map of the"
255 " Jacobian graph. "
256 "You must call the non-fitted version of this function.");
257
258 auto V_view_dev = V.getLocalViewDevice(Tpetra::Access::OverwriteAll);
259 const size_t num_local_cols = graph->getLocalNumCols();
260 const int chunk_size = V.getNumVectors();
261 list_of_colors_t my_list_of_colors = list_of_colors;
262
263 Kokkos::parallel_for(
264 "TpetraCrsColorer::computeSeedMatrixFitted()",
265 Kokkos::RangePolicy<execution_space>(0, num_local_cols),
266 KOKKOS_LAMBDA(const size_t i) {
267 const int color = my_list_of_colors[i] - 1 - color_beg;
268 if (color >= 0 && color < chunk_size)
269 V_view_dev(i, color) = scalar_t(1.0);
270 });
271
272}
273
275template <typename CrsMatrixType>
276template <typename MultiVectorType>
277void
279 MultiVectorType &W, const int color_beg) const
280{
281 reconstructMatrix(W, *matrix, color_beg);
282}
283
285template <typename CrsMatrixType>
286template <typename MultiVectorType>
287void
289 MultiVectorType &W,
290 matrix_t &mat,
291 const int color_beg) const
292{
293 MultiVectorType W_fitted(graph->getRowMap(), W.getNumVectors());
294
295 Tpetra::Import<lno_t, gno_t, node_t>
296 importer(W.getMap(), graph->getRowMap());
297
298 W_fitted.doImport(W, importer, Tpetra::INSERT);
299
300 reconstructMatrixFitted(W_fitted, mat, color_beg);
301}
302
304template <typename CrsMatrixType>
305template <typename MultiVectorType>
306void
308 MultiVectorType &W, const int color_beg) const
309{
310 reconstructMatrixFitted(W, *matrix, color_beg);
311}
312
314template <typename CrsMatrixType>
315template <typename MultiVectorType>
316void
318 MultiVectorType &W,
319 matrix_t &mat,
320 const int color_beg) const
321{
322 // Check the graph's row map is locally fitted to W's map
323 TEUCHOS_TEST_FOR_EXCEPTION(
324 !(W.getMap()->isLocallyFitted(*(graph->getRowMap()))), std::domain_error,
325 "Row map of the Jacobian graph is not locally fitted to the vector's map."
326 " You must call the non-fitted version of this function.");
327
328 auto W_view_dev = W.getLocalViewDevice(Tpetra::Access::ReadOnly);
329 auto local_matrix = mat.getLocalMatrixDevice();
330 auto local_graph = local_matrix.graph;
331 const size_t num_local_rows = graph->getLocalNumRows();
332 const int chunk_size = W.getNumVectors();
333 list_of_colors_t my_list_of_colors = list_of_colors;
334
335 Kokkos::parallel_for(
336 "TpetraCrsColorer::reconstructMatrixFitted()",
337 Kokkos::RangePolicy<execution_space>(0, num_local_rows),
338 KOKKOS_LAMBDA(const size_t row) {
339 const size_t entry_begin = local_graph.row_map(row);
340 const size_t entry_end = local_graph.row_map(row + 1);
341 for (size_t entry = entry_begin; entry < entry_end; entry++)
342 {
343 const size_t col = local_graph.entries(entry);
344 const int color = my_list_of_colors[col] - 1 - color_beg;
345 if (color >= 0 && color < chunk_size)
346 local_matrix.values(entry) = W_view_dev(row, color);
347 }
348 });
349}
350
352template <typename SC, typename LO, typename GO, typename NO>
353TpetraCrsColorer<Tpetra::BlockCrsMatrix<SC,LO,GO,NO> >::TpetraCrsColorer(
354 const Teuchos::RCP<matrix_t> &matrix_
355)
356 : matrix(matrix_)
357 , graph(Teuchos::rcp(&(matrix->getCrsGraph()), false))
358 , list_of_colors()
359 , list_of_colors_host()
360 , num_colors(0)
361{
362}
363
365template <typename SC, typename LO, typename GO, typename NO>
366void
367TpetraCrsColorer<Tpetra::BlockCrsMatrix<SC,LO,GO,NO> >::computeColoring(
368 Teuchos::ParameterList &coloring_params
369)
370{
371 const std::string library = coloring_params.get("library", "zoltan");
372
373 if (library == "zoltan") {
374 // Use Zoltan's coloring
375 ZoltanCrsColorer<matrix_t> zz(matrix);
376 zz.computeColoring(coloring_params,
377 num_colors, list_of_colors_host, list_of_colors);
378 }
379 else {
380 // Use Zoltan2's coloring
381 Zoltan2CrsColorer<matrix_t> zz2(matrix);
382 zz2.computeColoring(coloring_params,
383 num_colors, list_of_colors_host, list_of_colors);
384 }
385}
386
388template <typename SC, typename LO, typename GO, typename NO>
389void
390TpetraCrsColorer<Tpetra::BlockCrsMatrix<SC,LO,GO,NO> >::computeSeedMatrix(
391 multivector_t &block_V, const int color_beg) const
392{
393 multivector_t block_V_fitted(*(graph->getColMap()), matrix->getBlockSize(),
394 block_V.getNumVectors());
395
396 computeSeedMatrixFitted(block_V_fitted, color_beg);
397
398 const lno_t block_size = block_V.getBlockSize();
399 auto col_point_map = multivector_t::makePointMap(*graph->getColMap(),
400 block_size);
401 auto blockV_point_map = multivector_t::makePointMap(*block_V.getMap(),
402 block_size);
403 const auto col_point_map_rcp =
404 Teuchos::rcp(new Tpetra::Map<LO, GO, NO>(col_point_map));
405 const auto blockV_point_map_rcp =
406 Teuchos::rcp(new Tpetra::Map<LO, GO, NO>(blockV_point_map));
407
408 Tpetra::Import<LO, GO, NO> importer_point(col_point_map_rcp,
409 blockV_point_map_rcp);
410
411 block_V.getMultiVectorView().doImport(block_V_fitted.getMultiVectorView(),
412 importer_point, Tpetra::INSERT);
413
414 // Tpetra::Import<LO,GO,NO> importer(graph->getColMap(), block_V.getMap());
415 // block_V.doImport(block_V_fitted, importer, Tpetra::INSERT);
416}
417
419template <typename SC, typename LO, typename GO, typename NO>
420void
421TpetraCrsColorer<Tpetra::BlockCrsMatrix<SC,LO,GO,NO> >::computeSeedMatrixFitted(
422 multivector_t &blockV, const int color_beg) const
423{
424 // Check blockV's map is locally fitted to the graph's column map
425 TEUCHOS_TEST_FOR_EXCEPTION(
426 !(graph->getColMap()->isLocallyFitted(*(blockV.getMap()))),
427 std::domain_error,
428 "Map for supplied vector is not locally fitted to the column map of the"
429 " Jacobian graph. "
430 "You must call the non-fitted version of this function.");
431
432 const lno_t block_size = blockV.getBlockSize();
433 auto V = blockV.getMultiVectorView();
434 auto V_view_dev = V.getLocalViewDevice(Tpetra::Access::ReadWrite);
435 const size_t num_local_cols = V_view_dev.extent(0) / block_size;
436 const int chunk_size = V.getNumVectors() / block_size;
437 list_of_colors_t my_list_of_colors = list_of_colors;
438
439 Kokkos::parallel_for(
440 "TpetraCrsColorer::computeSeedMatrixFitted()",
441 Kokkos::RangePolicy<execution_space>(0, num_local_cols),
442 KOKKOS_LAMBDA(const size_t i) {
443 const int color = my_list_of_colors[i] - 1 - color_beg;
444 if (color >= 0 && color < chunk_size)
445 for (lno_t j = 0; j < block_size; ++j)
446 V_view_dev(i*block_size+j, color*block_size+j) = scalar_t(1.0);
447 });
448
449}
450
452template <typename SC, typename LO, typename GO, typename NO>
453void
454TpetraCrsColorer<Tpetra::BlockCrsMatrix<SC,LO,GO,NO> >::reconstructMatrix(
455 multivector_t &W, const int color_beg) const
456{
457 reconstructMatrix(W, *matrix, color_beg);
458}
459
461template <typename SC, typename LO, typename GO, typename NO>
462void
463TpetraCrsColorer<Tpetra::BlockCrsMatrix<SC,LO,GO,NO> >::reconstructMatrix(
464 multivector_t &block_W,
465 matrix_t &mat,
466 const int color_beg) const
467{
468 multivector_t block_W_fitted(*(graph->getRowMap()), matrix->getBlockSize(),
469 block_W.getNumVectors());
470
471 const lno_t block_size = block_W.getBlockSize();
472 auto row_point_map = multivector_t::makePointMap(*graph->getRowMap(),
473 block_size);
474 auto blockW_point_map = multivector_t::makePointMap(*block_W.getMap(),
475 block_size);
476 const auto row_point_map_rcp =
477 Teuchos::rcp(new Tpetra::Map<LO, GO, NO>(row_point_map));
478 const auto blockW_point_map_rcp =
479 Teuchos::rcp(new Tpetra::Map<LO, GO, NO>(blockW_point_map));
480
481 Tpetra::Import<lno_t, gno_t, node_t>
482 importer_point(blockW_point_map_rcp, row_point_map_rcp);
483
484 block_W_fitted.getMultiVectorView().doImport(block_W.getMultiVectorView(),
485 importer_point, Tpetra::INSERT);
486 reconstructMatrixFitted(block_W_fitted, mat, color_beg);
487}
488
490template <typename SC, typename LO, typename GO, typename NO>
491void
492TpetraCrsColorer<Tpetra::BlockCrsMatrix<SC,LO,GO,NO> >::reconstructMatrixFitted(
493 multivector_t &W, const int color_beg) const
494{
495 reconstructMatrixFitted(W, *matrix, color_beg);
496}
497
499template <typename SC, typename LO, typename GO, typename NO>
500void
501TpetraCrsColorer<Tpetra::BlockCrsMatrix<SC,LO,GO,NO> >::reconstructMatrixFitted(
502 multivector_t &block_W,
503 matrix_t &mat,
504 const int color_beg) const
505{
506 // Check the graph's row map is locally fitted to W's map
507 TEUCHOS_TEST_FOR_EXCEPTION(
508 !(block_W.getMap()->isLocallyFitted(*(graph->getRowMap()))),
509 std::domain_error,
510 "Row map of the Jacobian graph is not locally fitted to the vector's map."
511 " You must call the non-fitted version of this function.");
512
513 // Blocks are block_size x block_size stored with LayoutRight
514 const lno_t block_size = block_W.getBlockSize();
515 const lno_t block_stride = block_size * block_size;
516 const lno_t block_row_stride = block_size;
517
518 auto W = block_W.getMultiVectorView();
519 auto W_view_dev = W.getLocalViewDevice(Tpetra::Access::ReadOnly);
520 auto matrix_vals = mat.getValuesDeviceNonConst();
521 auto local_graph = graph->getLocalGraphDevice();
522 const size_t num_local_rows = graph->getLocalNumRows();
523 const int chunk_size = W.getNumVectors() / block_size;
524 list_of_colors_t my_list_of_colors = list_of_colors;
525
526 Kokkos::parallel_for(
527 "TpetraCrsColorer::reconstructMatrix()",
528 Kokkos::RangePolicy<execution_space>(0, num_local_rows),
529 KOKKOS_LAMBDA(const size_t block_row) {
530 const size_t entry_begin = local_graph.row_map(block_row);
531 const size_t entry_end = local_graph.row_map(block_row + 1);
532 for (size_t block_entry = entry_begin; block_entry < entry_end;
533 block_entry++)
534 {
535 const size_t block_col = local_graph.entries(block_entry);
536 const size_t block_offset = block_stride * block_entry;
537 const int block_color = my_list_of_colors[block_col] - 1 - color_beg;
538 if (block_color >= 0 && block_color < chunk_size)
539 {
540 for (lno_t i = 0; i < block_size; ++i)
541 {
542 const size_t row = block_row * block_size + i;
543 for (lno_t j = 0; j < block_size; ++j)
544 {
545 const size_t entry = block_offset + block_row_stride * i + j;
546 matrix_vals(entry) = W_view_dev(row, block_color*block_size+j);
547 }
548 }
549 }
550 }
551 });
552}
553} // namespace Tpetra
list_of_colors_t::HostMirror list_of_colors_host_t
void computeSeedMatrix(MultiVectorType &V, const int color_beg=0) const
TpetraCrsColorer(const Teuchos::RCP< matrix_t > &matrix_)
Kokkos::View< int *, device_t > list_of_colors_t
void computeSeedMatrixFitted(MultiVectorType &V, const int color_beg=0) const
device_t::execution_space execution_space
Teuchos::RCP< const graph_t > graph
void computeColoring(Teuchos::ParameterList &coloring_params)
void reconstructMatrix(MultiVectorType &W, const int color_beg=0) const
list_of_colors_host_t list_of_colors_host
void reconstructMatrixFitted(MultiVectorType &W, const int color_beg=0) const
matrix_t::local_ordinal_type lno_t
int getColor(const size_t col) const
matrix_t::global_ordinal_type gno_t
bool check_coloring(const Tpetra::CrsGraph< LO, GO, NO > &graph, const list_of_colors_t &list_of_colors)
Created by mbenlioglu on Aug 31, 2020.