Zoltan2
Loading...
Searching...
No Matches
Zoltan2_TpetraCrsColorer_Zoltan2.hpp
Go to the documentation of this file.
1#pragma once
2
3#include "Teuchos_ArrayRCP.hpp"
4#include "Teuchos_TestForException.hpp"
5
6#include "TpetraExt_MatrixMatrix.hpp"
7
11
12namespace Zoltan2
13{
14
15// Implementation of CrsColorer<> using Zoltan2. Currently this is a distance-1
16// algorithm, which requires forming A^T*A, and only works in serial
17template <typename CrsMatrixType>
19{
20public:
21 typedef CrsMatrixType matrix_t;
22 typedef typename matrix_t::crs_graph_type graph_t;
23 typedef typename matrix_t::node_type node_t;
24 typedef typename node_t::device_type device_t;
25 typedef Kokkos::View<int *, device_t> list_of_colors_t;
26 typedef typename list_of_colors_t::HostMirror list_of_colors_host_t;
27
28 // Constructor
29 Zoltan2CrsColorer(const Teuchos::RCP<matrix_t> &matrix_)
30 : matrix(matrix_), graph(matrix_->getCrsGraph())
31 {}
32
33 // Destructor
35
36 // Compute coloring data
37 void
39 Teuchos::ParameterList &coloring_params,
40 int &num_colors,
41 list_of_colors_host_t &list_of_colors_host,
42 list_of_colors_t &list_of_colors)
43 {
44 // TODO: logic for symmetric/Jacobian/Hessian a la ZoltanCrsColorer
45 // User can tell us that the matrix is symmetric;
46 // otherwise, guess based on the matrix type
47 // const std::string matrixType = coloring_params.get("matrixType","Jacobian");
48 // const bool symmetric = coloring_params.get("symmetric",
49 // (matrixType=="Jacobian" ? false
50 // : true));
51
52 // TODO: Until Ian's code is ready...
53 // TODO: Check the logic here: doing Partial Distance-2 via local
54 // TODO: distance-1 on J^T*J
55
56 // Compute A^T*A where A = jac
57 // We have to do this because Zoltan2 can only do distance-1 coloring
58 this->matrix->setAllToScalar(1.0);
59
60 if (!this->matrix->isFillComplete())
61 this->matrix->fillComplete();
62
63 const size_t nzpr = this->matrix->getGlobalMaxNumRowEntries();
64 matrix_t C(this->matrix->getRowMap(), nzpr * nzpr); // TODO: Check this
65
66 Tpetra::MatrixMatrix::Multiply(*(this->matrix), true,
67 *(this->matrix), false, C);
68
69 // Create Zoltan2 coloring problem and solve
71 Z2Adapter_t z2_adapter(rcp(&C, false));
72
73 Teuchos::ParameterList z2_params = coloring_params.sublist("Zoltan2");
74 Zoltan2::ColoringProblem<Z2Adapter_t> z2_problem(&z2_adapter, &z2_params);
75 z2_problem.solve();
76
77 // Extract colors
79 z2_problem.getSolution();
80 int local_num_colors = z2_solution->getNumColors();
81
82 Teuchos::ArrayRCP<int> local_list_of_colors = z2_solution->getColorsRCP();
83 const size_t len = local_list_of_colors.size();
84
85 TEUCHOS_TEST_FOR_EXCEPTION(
86 len != this->graph->getColMap()->getLocalNumElements(), std::logic_error,
87 "Incorrect length of color list!");
88
89 list_of_colors_host_t list_of_colors_tmp(
90 local_list_of_colors.getRawPtr(), len);
91
92 list_of_colors = list_of_colors_t("list_of_colors", len);
93 list_of_colors_host = Kokkos::create_mirror_view(list_of_colors);
94
95 Kokkos::deep_copy(list_of_colors_host, list_of_colors_tmp);
96 Kokkos::deep_copy(list_of_colors, list_of_colors_host);
97
98 // Compute global number of colors
99 Teuchos::RCP<const Teuchos::Comm<int>> comm =
100 this->graph->getRowMap()->getComm();
101 Teuchos::reduceAll(*comm, Teuchos::REDUCE_MAX, 1,
102 &local_num_colors, &num_colors);
103 }
104
105private:
106
107 const Teuchos::RCP<matrix_t> matrix;
108 const Teuchos::RCP<const graph_t> graph;
109};
110
112// Specialization of Tpetra::Zoltan2CrsColorer for BlockCrsMatrix
113// Zoltan2 does not directly support BlockCrs, so this implementation
114// creates a point matrix from the graph of the BlockCrs matrix
115template <typename SC, typename LO, typename GO, typename NO>
116class Zoltan2CrsColorer<Tpetra::BlockCrsMatrix<SC, LO, GO, NO> >
117{
118public:
119 typedef Tpetra::BlockCrsMatrix<SC, LO, GO, NO> matrix_t;
120 typedef typename matrix_t::crs_graph_type graph_t;
121 typedef typename matrix_t::node_type node_t;
122 typedef typename node_t::device_type device_t;
123 typedef Kokkos::View<int *, device_t> list_of_colors_t;
124 typedef typename list_of_colors_t::HostMirror list_of_colors_host_t;
125
126 // Constructor
127 Zoltan2CrsColorer(const Teuchos::RCP<matrix_t> &matrix_)
128 : matrix(matrix_), graph(Teuchos::rcp(&(matrix->getCrsGraph()), false))
129 {}
130
131 // Destructor
133
134 // Compute coloring data
135 void
137 Teuchos::ParameterList &coloring_params,
138 int &num_colors,
139 list_of_colors_host_t &list_of_colors_host,
140 list_of_colors_t &list_of_colors)
141 {
142 using point_matrix_t = Tpetra::CrsMatrix<SC,LO,GO,NO>;
143 Teuchos::RCP<point_matrix_t> point_matrix = Teuchos::rcp(new point_matrix_t(graph));
144 point_matrix->setAllToScalar(1.0);
145 point_matrix->fillComplete();
146 Zoltan2CrsColorer<point_matrix_t> point_colorer(point_matrix);
147 point_colorer.computeColoring(coloring_params, num_colors,
148 list_of_colors_host, list_of_colors);
149 }
150
151private:
152
153 const Teuchos::RCP<matrix_t> matrix;
154 const Teuchos::RCP<const graph_t> graph;
155};
156} // namespace Zoltan2
Defines the ColoringProblem class.
Defines the ColoringSolution class.
Defines the XpetraCrsMatrixAdapter class.
ColoringProblem sets up coloring problems for the user.
The class containing coloring solution.
int getNumColors()
Get local number of colors. This is computed from the coloring each time, as this is cheap.
Provides access for Zoltan2 to Xpetra::CrsMatrix data.
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)
Kokkos::View< int *, device_t > list_of_colors_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)
list_of_colors_t::HostMirror list_of_colors_host_t
Zoltan2CrsColorer(const Teuchos::RCP< matrix_t > &matrix_)
Created by mbenlioglu on Aug 31, 2020.