Zoltan2
 All Classes Namespaces Files Functions Variables Typedefs Enumerations Enumerator Macros Pages
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 
20 namespace Zoltan2
21 {
22 
23 // Base class for coloring Tpetra::CrsMatrix for use in column compression
24 template <typename CrsMatrixType>
26 {
27 public:
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
79  getNumColors() const
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
94  checkColoring() const
95  {
97  }
98 
99 protected:
100  Teuchos::RCP<matrix_t> matrix;
101  Teuchos::RCP<const graph_t> graph;
105 };
106 
108 // Specialization of TpetraCrsColorer for BlockCrsMatrix
109 template <typename SC, typename LO, typename GO, typename NO>
110 class TpetraCrsColorer<Tpetra::BlockCrsMatrix<SC,LO,GO,NO> >
111 {
112 public:
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
158  getNumColors() const
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 
181 protected:
182  Teuchos::RCP<matrix_t> matrix;
183  Teuchos::RCP<const graph_t> graph;
187 };
188 
190 
191 template <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 
205 template <typename CrsMatrixType>
206 void
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 
228 template <typename CrsMatrixType>
229 template <typename MultiVectorType>
230 void
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 
245 template <typename CrsMatrixType>
246 template <typename MultiVectorType>
247 void
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 
275 template <typename CrsMatrixType>
276 template <typename MultiVectorType>
277 void
279  MultiVectorType &W, const int color_beg) const
280 {
281  reconstructMatrix(W, *matrix, color_beg);
282 }
283 
285 template <typename CrsMatrixType>
286 template <typename MultiVectorType>
287 void
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 
304 template <typename CrsMatrixType>
305 template <typename MultiVectorType>
306 void
308  MultiVectorType &W, const int color_beg) const
309 {
310  reconstructMatrixFitted(W, *matrix, color_beg);
311 }
312 
314 template <typename CrsMatrixType>
315 template <typename MultiVectorType>
316 void
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 
352 template <typename SC, typename LO, typename GO, typename NO>
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 
365 template <typename SC, typename LO, typename GO, typename NO>
366 void
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 
388 template <typename SC, typename LO, typename GO, typename NO>
389 void
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 
419 template <typename SC, typename LO, typename GO, typename NO>
420 void
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  const int block_color_beg = color_beg / block_size;
438  list_of_colors_t my_list_of_colors = list_of_colors;
439 
440  Kokkos::parallel_for(
441  "TpetraCrsColorer::computeSeedMatrixFitted()",
442  Kokkos::RangePolicy<execution_space>(0, num_local_cols),
443  KOKKOS_LAMBDA(const size_t i) {
444  const int color = my_list_of_colors[i] - 1 - block_color_beg;
445  if (color >= 0 && color < chunk_size)
446  for (lno_t j = 0; j < block_size; ++j)
447  V_view_dev(i*block_size+j, color*block_size+j) = scalar_t(1.0);
448  });
449 
450 }
451 
453 template <typename SC, typename LO, typename GO, typename NO>
454 void
456  multivector_t &W, const int color_beg) const
457 {
458  reconstructMatrix(W, *matrix, color_beg);
459 }
460 
462 template <typename SC, typename LO, typename GO, typename NO>
463 void
465  multivector_t &block_W,
466  matrix_t &mat,
467  const int color_beg) const
468 {
469  multivector_t block_W_fitted(*(graph->getRowMap()), matrix->getBlockSize(),
470  block_W.getNumVectors());
471 
472  const lno_t block_size = block_W.getBlockSize();
473  auto row_point_map = multivector_t::makePointMap(*graph->getRowMap(),
474  block_size);
475  auto blockW_point_map = multivector_t::makePointMap(*block_W.getMap(),
476  block_size);
477  const auto row_point_map_rcp =
478  Teuchos::rcp(new Tpetra::Map<LO, GO, NO>(row_point_map));
479  const auto blockW_point_map_rcp =
480  Teuchos::rcp(new Tpetra::Map<LO, GO, NO>(blockW_point_map));
481 
482  Tpetra::Import<lno_t, gno_t, node_t>
483  importer_point(blockW_point_map_rcp, row_point_map_rcp);
484 
485  block_W_fitted.getMultiVectorView().doImport(block_W.getMultiVectorView(),
486  importer_point, Tpetra::INSERT);
487  reconstructMatrixFitted(block_W_fitted, mat, color_beg);
488 }
489 
491 template <typename SC, typename LO, typename GO, typename NO>
492 void
494  multivector_t &W, const int color_beg) const
495 {
496  reconstructMatrixFitted(W, *matrix, color_beg);
497 }
498 
500 template <typename SC, typename LO, typename GO, typename NO>
501 void
503  multivector_t &block_W,
504  matrix_t &mat,
505  const int color_beg) const
506 {
507  // Check the graph's row map is locally fitted to W's map
508  TEUCHOS_TEST_FOR_EXCEPTION(
509  !(block_W.getMap()->isLocallyFitted(*(graph->getRowMap()))),
510  std::domain_error,
511  "Row map of the Jacobian graph is not locally fitted to the vector's map."
512  " You must call the non-fitted version of this function.");
513 
514  // Blocks are block_size x block_size stored with LayoutRight
515  const lno_t block_size = block_W.getBlockSize();
516  const lno_t block_stride = block_size * block_size;
517  const lno_t block_row_stride = block_size;
518 
519  auto W = block_W.getMultiVectorView();
520  auto W_view_dev = W.getLocalViewDevice(Tpetra::Access::ReadOnly);
521  auto matrix_vals = mat.getValuesDeviceNonConst();
522  auto local_graph = graph->getLocalGraphDevice();
523  const size_t num_local_rows = graph->getLocalNumRows();
524  const int chunk_size = W.getNumVectors() / block_size;
525  const int block_color_beg = color_beg / block_size;
526  list_of_colors_t my_list_of_colors = list_of_colors;
527 
528  Kokkos::parallel_for(
529  "TpetraCrsColorer::reconstructMatrix()",
530  Kokkos::RangePolicy<execution_space>(0, num_local_rows),
531  KOKKOS_LAMBDA(const size_t block_row) {
532  const size_t entry_begin = local_graph.row_map(block_row);
533  const size_t entry_end = local_graph.row_map(block_row + 1);
534  for (size_t block_entry = entry_begin; block_entry < entry_end;
535  block_entry++)
536  {
537  const size_t block_col = local_graph.entries(block_entry);
538  const size_t block_offset = block_stride * block_entry;
539  const int block_color = my_list_of_colors[block_col] - 1 - block_color_beg;
540  if (block_color >= 0 && block_color < chunk_size)
541  {
542  for (lno_t i = 0; i < block_size; ++i)
543  {
544  const size_t row = block_row * block_size + i;
545  for (lno_t j = 0; j < block_size; ++j)
546  {
547  const size_t entry = block_offset + block_row_stride * i + j;
548  matrix_vals(entry) = W_view_dev(row, block_color*block_size+j);
549  }
550  }
551  }
552  }
553  });
554 }
555 } // namespace Tpetra
void computeColoring(Teuchos::ParameterList &coloring_params, int &num_colors, list_of_colors_host_t &list_of_colors_host, list_of_colors_t &list_of_colors) const
list_of_colors_host_t list_of_colors_host
void computeColoring(Teuchos::ParameterList &coloring_params)
void computeColoring(Teuchos::ParameterList &coloring_params, int &num_colors, list_of_colors_host_t &list_of_colors_host, list_of_colors_t &list_of_colors)
list_of_colors_t::HostMirror list_of_colors_host_t
matrix_t::local_ordinal_type lno_t
void computeSeedMatrix(MultiVectorType &V, const int color_beg=0) const
void reconstructMatrix(MultiVectorType &W, const int color_beg=0) const
void reconstructMatrixFitted(MultiVectorType &W, const int color_beg=0) const
bool check_coloring(const Tpetra::CrsGraph< LO, GO, NO > &graph, const list_of_colors_t &list_of_colors)
Kokkos::View< int *, device_t > list_of_colors_t
TpetraCrsColorer(const Teuchos::RCP< matrix_t > &matrix_)
matrix_t::global_ordinal_type gno_t
int getColor(const size_t col) const
Teuchos::RCP< const graph_t > graph
void computeSeedMatrixFitted(MultiVectorType &V, const int color_beg=0) const
device_t::execution_space execution_space
matrix_t::crs_graph_type graph_t