5 #include "Teuchos_RCP.hpp"
6 #include "Teuchos_ParameterList.hpp"
7 #include "Teuchos_TestForException.hpp"
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"
24 template <
typename CrsMatrixType>
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;
50 template <
typename MultiVectorType>
55 template <
typename MultiVectorType>
60 template <
typename MultiVectorType>
63 template <
typename MultiVectorType>
69 template <
typename MultiVectorType>
72 template <
typename MultiVectorType>
109 template <
typename SC,
typename LO,
typename GO,
typename NO>
113 typedef Tpetra::BlockCrsMatrix<SC, LO, GO, NO>
matrix_t;
115 typedef typename matrix_t::crs_graph_type
graph_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;
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;
191 template <
typename CrsMatrixType>
193 const Teuchos::RCP<matrix_t> &matrix_
196 graph(matrix->getCrsGraph()),
198 list_of_colors_host(),
205 template <
typename CrsMatrixType>
208 Teuchos::ParameterList &coloring_params
211 const std::string library = coloring_params.get(
"library",
"zoltan");
213 if (library ==
"zoltan") {
217 num_colors, list_of_colors_host, list_of_colors);
223 num_colors, list_of_colors_host, list_of_colors);
228 template <
typename CrsMatrixType>
229 template <
typename MultiVectorType>
232 MultiVectorType &V,
const int color_beg)
const
234 MultiVectorType V_fitted(graph->getColMap(), V.getNumVectors());
236 computeSeedMatrixFitted(V_fitted, color_beg);
238 Tpetra::Import<lno_t, gno_t, node_t>
239 importer(graph->getColMap(), V.getMap());
241 V.doImport(V_fitted, importer, Tpetra::INSERT);
245 template <
typename CrsMatrixType>
246 template <
typename MultiVectorType>
249 MultiVectorType &V,
const int color_beg)
const
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"
256 "You must call the non-fitted version of this function.");
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();
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);
275 template <
typename CrsMatrixType>
276 template <
typename MultiVectorType>
279 MultiVectorType &W,
const int color_beg)
const
281 reconstructMatrix(W, *matrix, color_beg);
285 template <
typename CrsMatrixType>
286 template <
typename MultiVectorType>
291 const int color_beg)
const
293 MultiVectorType W_fitted(graph->getRowMap(), W.getNumVectors());
295 Tpetra::Import<lno_t, gno_t, node_t>
296 importer(W.getMap(), graph->getRowMap());
298 W_fitted.doImport(W, importer, Tpetra::INSERT);
300 reconstructMatrixFitted(W_fitted, mat, color_beg);
304 template <
typename CrsMatrixType>
305 template <
typename MultiVectorType>
308 MultiVectorType &W,
const int color_beg)
const
310 reconstructMatrixFitted(W, *matrix, color_beg);
314 template <
typename CrsMatrixType>
315 template <
typename MultiVectorType>
320 const int color_beg)
const
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.");
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();
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++)
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);
352 template <
typename SC,
typename LO,
typename GO,
typename NO>
354 const Teuchos::RCP<matrix_t> &matrix_
357 , graph(Teuchos::rcp(&(matrix->getCrsGraph()), false))
359 , list_of_colors_host()
365 template <
typename SC,
typename LO,
typename GO,
typename NO>
368 Teuchos::ParameterList &coloring_params
371 const std::string library = coloring_params.get(
"library",
"zoltan");
373 if (library ==
"zoltan") {
377 num_colors, list_of_colors_host, list_of_colors);
383 num_colors, list_of_colors_host, list_of_colors);
388 template <
typename SC,
typename LO,
typename GO,
typename NO>
393 multivector_t block_V_fitted(*(graph->getColMap()), matrix->getBlockSize(),
394 block_V.getNumVectors());
396 computeSeedMatrixFitted(block_V_fitted, color_beg);
398 const lno_t block_size = block_V.getBlockSize();
399 auto col_point_map = multivector_t::makePointMap(*graph->getColMap(),
401 auto blockV_point_map = multivector_t::makePointMap(*block_V.getMap(),
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));
408 Tpetra::Import<LO, GO, NO> importer_point(col_point_map_rcp,
409 blockV_point_map_rcp);
411 block_V.getMultiVectorView().doImport(block_V_fitted.getMultiVectorView(),
412 importer_point, Tpetra::INSERT);
419 template <
typename SC,
typename LO,
typename GO,
typename NO>
425 TEUCHOS_TEST_FOR_EXCEPTION(
426 !(graph->getColMap()->isLocallyFitted(*(blockV.getMap()))),
428 "Map for supplied vector is not locally fitted to the column map of the"
430 "You must call the non-fitted version of this function.");
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;
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);
453 template <
typename SC,
typename LO,
typename GO,
typename NO>
458 reconstructMatrix(W, *matrix, color_beg);
462 template <
typename SC,
typename LO,
typename GO,
typename NO>
467 const int color_beg)
const
469 multivector_t block_W_fitted(*(graph->getRowMap()), matrix->getBlockSize(),
470 block_W.getNumVectors());
472 const lno_t block_size = block_W.getBlockSize();
473 auto row_point_map = multivector_t::makePointMap(*graph->getRowMap(),
475 auto blockW_point_map = multivector_t::makePointMap(*block_W.getMap(),
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));
482 Tpetra::Import<lno_t, gno_t, node_t>
483 importer_point(blockW_point_map_rcp, row_point_map_rcp);
485 block_W_fitted.getMultiVectorView().doImport(block_W.getMultiVectorView(),
486 importer_point, Tpetra::INSERT);
487 reconstructMatrixFitted(block_W_fitted, mat, color_beg);
491 template <
typename SC,
typename LO,
typename GO,
typename NO>
496 reconstructMatrixFitted(W, *matrix, color_beg);
500 template <
typename SC,
typename LO,
typename GO,
typename NO>
505 const int color_beg)
const
508 TEUCHOS_TEST_FOR_EXCEPTION(
509 !(block_W.getMap()->isLocallyFitted(*(graph->getRowMap()))),
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.");
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;
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;
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;
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)
542 for (
lno_t i = 0; i < block_size; ++i)
544 const size_t row = block_row * block_size + i;
545 for (
lno_t j = 0; j < block_size; ++j)
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);
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
matrix_t::global_ordinal_type gno_t
Tpetra::BlockCrsMatrix< SC, LO, GO, NO > matrix_t
Teuchos::RCP< matrix_t > matrix
list_of_colors_host_t list_of_colors_host
node_t::device_type device_t
void computeColoring(Teuchos::ParameterList &coloring_params)
matrix_t::scalar_type scalar_t
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)
matrix_t::local_ordinal_type lno_t
Kokkos::View< int *, device_t > list_of_colors_t
list_of_colors_t::HostMirror list_of_colors_host_t
matrix_t::local_ordinal_type lno_t
Teuchos::RCP< const graph_t > graph
void computeSeedMatrix(MultiVectorType &V, const int color_beg=0) const
matrix_t::node_type node_t
matrix_t::crs_graph_type graph_t
list_of_colors_t list_of_colors
list_of_colors_t list_of_colors
node_t::device_type device_t
list_of_colors_host_t list_of_colors_host
matrix_t::node_type node_t
matrix_t::scalar_type scalar_t
Tpetra::BlockMultiVector< SC, LO, GO, NO > multivector_t
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
list_of_colors_t::HostMirror list_of_colors_host_t
bool checkColoring() const
TpetraCrsColorer(const Teuchos::RCP< matrix_t > &matrix_)
device_t::execution_space execution_space
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
bool checkColoring() const
int getColor(const size_t col) const
device_t::execution_space execution_space
matrix_t::crs_graph_type graph_t