Zoltan2
Loading...
Searching...
No Matches
Zoltan2_TpetraCrsColorer_Zoltan.hpp
Go to the documentation of this file.
1#pragma once
2
4
5#include "Teuchos_Array.hpp"
6#include "Teuchos_ArrayView.hpp"
7#include "Teuchos_TestForException.hpp"
8#include "Teuchos_DefaultMpiComm.hpp"
9
10#include "zoltan_cpp.h"
11
12namespace Zoltan2
13{
14
15namespace Impl {
16
17 template <typename SC, typename LO, typename GO, typename NO>
18 Teuchos::RCP<const typename Tpetra::CrsMatrix<SC, LO, GO, NO>::crs_graph_type>
19 get_graph(const Teuchos::RCP<Tpetra::CrsMatrix<SC, LO, GO, NO> > &matrix) {
20 return matrix->getCrsGraph();
21 }
22
23 template <typename SC, typename LO, typename GO, typename NO>
24 Teuchos::RCP<const typename Tpetra::BlockCrsMatrix<SC, LO, GO, NO>::crs_graph_type>
25 get_graph(const Teuchos::RCP<Tpetra::BlockCrsMatrix<SC, LO, GO, NO> > &matrix) {
26 using crs_graph_t = typename Tpetra::BlockCrsMatrix<SC, LO, GO, NO>::crs_graph_type;
27 return Teuchos::rcp(new crs_graph_t(matrix->getCrsGraph()));
28 }
29
30}
31
32// Implementation of CrsColorer<> using Zoltan partial distance-2 coloring.
33// This is distributed-parallel, but not shared.
34template <typename CrsMatrixType>
36{
37public:
38
39 typedef CrsMatrixType matrix_t;
40 typedef typename matrix_t::crs_graph_type graph_t;
41 typedef typename matrix_t::scalar_type scalar_t;
42 typedef typename matrix_t::local_ordinal_type lno_t;
43 typedef typename matrix_t::global_ordinal_type gno_t;
44 typedef typename matrix_t::node_type node_t;
45 typedef typename node_t::device_type device_t;
46 typedef typename device_t::execution_space execution_space;
47 typedef Kokkos::View<int *, device_t> list_of_colors_t;
48 typedef typename list_of_colors_t::HostMirror list_of_colors_host_t;
49
50 // Constructor
51 ZoltanCrsColorer(const Teuchos::RCP<matrix_t> &matrix_)
52 : matrix(matrix_), graph(Impl::get_graph(matrix_))
53 {}
54
55 // Destructor
57
58 // Compute coloring data
59 void
61 Teuchos::ParameterList &coloring_params,
62 int &num_colors,
63 list_of_colors_host_t &list_of_colors_host,
64 list_of_colors_t &list_of_colors) const;
65
66private:
67
68 const Teuchos::RCP<const matrix_t> matrix;
69 const Teuchos::RCP<const graph_t> graph;
70
71 //
72 // Call-back functions for Zoltan interface
73 //
74
75 // Data for call-back functions
76 struct ZoltanData
77 {
78 Teuchos::RCP<const graph_t> graph, trans_graph;
79 Teuchos::Array<int> col_procs, trans_col_procs;
80
81 // Set matrix graphs (trans_graph_ may be null for symmetric/symmetrized
82 // problems). Computes the remote processor ID's for each entry in the
83 // graph's/trans_graph's column maps (involves communication, and so can't
84 // be done inside the Zoltan call-back functions).
85 void
86 setGraphs(
87 const Teuchos::RCP<const graph_t> &graph_,
88 const Teuchos::RCP<const graph_t> &trans_graph_ = Teuchos::null)
89 {
90 graph = graph_;
91 trans_graph = trans_graph_;
92 col_procs.resize(graph->getColMap()->getLocalNumElements());
93 auto gids = graph->getColMap()->getLocalElementList();
94
95 Tpetra::LookupStatus ret =
96 graph->getRowMap()->getRemoteIndexList(gids, col_procs());
97 TEUCHOS_TEST_FOR_EXCEPTION(ret != Tpetra::AllIDsPresent, std::logic_error,
98 "Zoltan2::CrsColorer: getRemoteIndexList() "
99 "failed!");
100
101 if (trans_graph != Teuchos::null)
102 {
103 trans_col_procs.resize(trans_graph->getColMap()->getLocalNumElements());
104 gids = trans_graph->getColMap()->getLocalElementList();
105 ret = trans_graph->getRowMap()->getRemoteIndexList(gids,
106 trans_col_procs());
107 TEUCHOS_TEST_FOR_EXCEPTION(ret != Tpetra::AllIDsPresent,
108 std::logic_error,
109 "Zoltan2::CrsColorer getRemoteIndexList() "
110 "failed!");
111 }
112 }
113 };
114
115 // Number of vertices
116 static int
117 get_number_of_vertices(void *data, int *ierr);
118
119 // Vertex IDs on this processor
120 static void
121 get_vertex_list(
122 void *data,
123 int sizeGID,
124 int sizeLID,
125 ZOLTAN_ID_PTR global_ids,
126 ZOLTAN_ID_PTR local_ids,
127 int wgt_dim,
128 float *obj_wgts,
129 int *ierr);
130
131 // Get number of edges on this processor
132 static int
133 get_number_of_edges(
134 void *data,
135 int sizeGID,
136 int sizeLID,
137 ZOLTAN_ID_PTR global_id,
138 ZOLTAN_ID_PTR local_id,
139 int *ierr);
140
141 // Get edge ids on this processor
142 static void
143 get_edge_list(
144 void *data,
145 int sizeGID,
146 int sizeLID,
147 ZOLTAN_ID_PTR global_id,
148 ZOLTAN_ID_PTR local_id,
149 ZOLTAN_ID_PTR nbor_global_ids,
150 int *nbor_procs,
151 int wgt_dim,
152 float *ewgts,
153 int *ierr);
154
155 //
156 // Call-back functions for Zoltan interface with symmetric graph
157 //
158
159 // Number of vertices
160 static int
161 sym_get_number_of_vertices(void *data, int *ierr);
162
163 // Vertex IDs on this processor
164 static void
165 sym_get_vertex_list(
166 void *data,
167 int sizeGID,
168 int sizeLID,
169 ZOLTAN_ID_PTR global_ids,
170 ZOLTAN_ID_PTR local_ids,
171 int wgt_dim,
172 float *obj_wgts,
173 int *ierr);
174
175 // Get number of edges on this processor
176 static int
177 sym_get_number_of_edges(
178 void *data,
179 int sizeGID,
180 int sizeLID,
181 ZOLTAN_ID_PTR global_id,
182 ZOLTAN_ID_PTR local_id,
183 int *ierr);
184
185 // Get edge ids on this processor
186 static void
187 sym_get_edge_list(
188 void *data,
189 int sizeGID,
190 int sizeLID,
191 ZOLTAN_ID_PTR global_id,
192 ZOLTAN_ID_PTR local_id,
193 ZOLTAN_ID_PTR nbor_global_ids,
194 int *nbor_procs,
195 int wgt_dim,
196 float *ewgts,
197 int *ierr);
198};
199
201template <typename CrsMatrixType>
202void
204 Teuchos::ParameterList &coloring_params,
205 int &num_colors,
206 list_of_colors_host_t &list_of_colors_host,
207 list_of_colors_t &list_of_colors
208) const
209{
210 // User can tell us that the matrix is symmetric;
211 // otherwise, guess based on the matrix type
212 const std::string matrixType = coloring_params.get("matrixType", "Jacobian");
213 const bool symmetric = coloring_params.get("symmetric",
214 (matrixType == "Jacobian" ? false
215 : true));
216
217 // User request to use Zoltan's TRANSPOSE symmetrization, if needed
218 const bool symmetrize = coloring_params.get<bool>("symmetrize", false);
219
220 // Get MPI communicator, and throw an exception if our comm isn't MPI
221 Teuchos::RCP<const Teuchos::Comm<int>> comm =
222 this->graph->getRowMap()->getComm();
223#ifdef HAVE_ZOLTAN2_MPI
224 Teuchos::RCP<const Teuchos::MpiComm<int>> tmpicomm =
225 Teuchos::rcp_dynamic_cast<const Teuchos::MpiComm<int>>(comm, true);
226 MPI_Comm mpicomm = *tmpicomm->getRawMpiComm();
227#else
228 // Zoltan's siMPI will be used here
229 {
230 int flag;
231 MPI_Initialized(&flag);
232 if (!flag) {
233 int narg = 0;
234 char **argv = NULL;
235 MPI_Init(&narg, &argv);
236 }
237 }
238 MPI_Comm mpicomm = MPI_COMM_WORLD; // Will get MPI_COMM_WORLD from siMPI
239#endif
240
241 // Create Zoltan for coloring
242 Zoltan *zz = new Zoltan(mpicomm);
243 if (symmetric || symmetrize) {
244 zz->Set_Param("COLORING_PROBLEM", "DISTANCE-2");
245 }
246 else {
247 zz->Set_Param("COLORING_PROBLEM", "PARTIAL-DISTANCE-2");
248 }
249
250 if (!symmetric && symmetrize)
251 zz->Set_Param("GRAPH_SYMMETRIZE", "TRANSPOSE");
252
253 zz->Set_Param("DEBUG_LEVEL", "0");
254
255 // Set Zoltan params
256 Teuchos::ParameterList &zoltan_params = coloring_params.sublist("Zoltan");
257 for (auto p : zoltan_params)
258 zz->Set_Param(p.first, Teuchos::getValue<std::string>(p.second));
259
260 // Compute transpose graph
261 Teuchos::RCP<const graph_t> transpose_graph;
262 if (!symmetric && !symmetrize)
263 {
264 transpose_graph = Impl::compute_transpose_graph(*(this->matrix));
265 }
266
267 // Setup interface functions
268 ZoltanData zd;
269 if (symmetric || symmetrize)
270 {
271 zd.setGraphs(this->graph);
272 zz->Set_Num_Obj_Fn(sym_get_number_of_vertices, &zd);
273 zz->Set_Obj_List_Fn(sym_get_vertex_list, &zd);
274 zz->Set_Num_Edges_Fn(sym_get_number_of_edges, &zd);
275 zz->Set_Edge_List_Fn(sym_get_edge_list, &zd);
276 }
277 else
278 {
279 zd.setGraphs(this->graph, transpose_graph);
280 zz->Set_Num_Obj_Fn(get_number_of_vertices, &zd);
281 zz->Set_Obj_List_Fn(get_vertex_list, &zd);
282 zz->Set_Num_Edges_Fn(get_number_of_edges, &zd);
283 zz->Set_Edge_List_Fn(get_edge_list, &zd);
284 }
285
286 // Do coloring of columns with Zoltan -- we can request colors for
287 // columns we don't own
288 const size_t num_local_cols = this->graph->getLocalNumCols();
289 const size_t num_global_rows = std::max(
290 static_cast<typename CrsMatrixType::global_ordinal_type>(
291 this->graph->getGlobalNumRows()),
292 this->graph->getRowMap()->getMaxAllGlobalIndex()+1);
293
294 Teuchos::Array<ZOLTAN_ID_TYPE> col_gids(num_local_cols);
295 auto gids = this->graph->getColMap()->getLocalElementList();
296
297 if (symmetric || symmetrize)
298 for (size_t i = 0; i < num_local_cols; ++i)
299 col_gids[i] = gids[i];
300 else
301 for (size_t i = 0; i < num_local_cols; ++i)
302 col_gids[i] = gids[i] + num_global_rows;
303
304 list_of_colors_t my_list_of_colors("ZoltanCrsColorer::list_of_colors",
305 num_local_cols);
306 list_of_colors_host = Kokkos::create_mirror_view(my_list_of_colors);
307
308 int num_gid_entries = 1;
309 int ret = zz->Color(num_gid_entries, num_local_cols, col_gids.getRawPtr(),
310 list_of_colors_host.data());
311
312 TEUCHOS_TEST_FOR_EXCEPTION(ret != ZOLTAN_OK, std::logic_error,
313 "Zoltan::Color returned " << ret << std::endl);
314
315 Kokkos::deep_copy(my_list_of_colors, list_of_colors_host);
316 list_of_colors = my_list_of_colors;
317
318 const bool dump_zoltan = coloring_params.get("Dump Zoltan Data", false);
319 if (dump_zoltan)
320 {
321 std::string zoltan_dump_file =
322 coloring_params.get("Zoltan Dump File Name", "zoltan_graph.txt");
323 zz->Generate_Files(zoltan_dump_file, 0, 0, 1, 0);
324 }
325
326 delete zz;
327
328 // Compute global number of colors
329 int local_num_colors = 0;
330 Kokkos::parallel_reduce("ZoltanCrsColorer::find_num_colors",
331 Kokkos::RangePolicy<execution_space>(0, num_local_cols),
332 KOKKOS_LAMBDA(const size_t i, int &lcl_max) {
333 if (my_list_of_colors[i] > lcl_max)
334 lcl_max = my_list_of_colors[i];
335 },
336 Kokkos::Max<int>(local_num_colors));
337
338 Teuchos::reduceAll(*comm, Teuchos::REDUCE_MAX, 1, &local_num_colors,
339 &num_colors);
340}
341
343template <typename CrsMatrixType>
344int
346{
347 ZoltanData *zoltan_data = static_cast<ZoltanData *>(data);
348 *ierr = ZOLTAN_OK;
349 return zoltan_data->graph->getLocalNumRows() +
350 zoltan_data->trans_graph->getLocalNumRows();
351}
352
354template <typename CrsMatrixType>
355void
356ZoltanCrsColorer<CrsMatrixType>::get_vertex_list(
357 void *data,
358 int sizeGID,
359 int sizeLID,
360 ZOLTAN_ID_PTR global_ids,
361 ZOLTAN_ID_PTR local_ids,
362 int wgt_dim,
363 float *obj_wgts,
364 int *ierr)
365{
366 assert(sizeGID == 1 && sizeLID == 1);
367
368 ZoltanData *zoltan_data = static_cast<ZoltanData *>(data);
369 *ierr = ZOLTAN_OK;
370
371 const size_t num_local_rows = zoltan_data->graph->getLocalNumRows();
372 const size_t num_local_cols = zoltan_data->trans_graph->getLocalNumRows();
373 const size_t num_global_rows = std::max(
374 static_cast<typename CrsMatrixType::global_ordinal_type>(
375 zoltan_data->graph->getGlobalNumRows()),
376 zoltan_data->graph->getRowMap()->getMaxAllGlobalIndex()+1);
377 auto row_gids = zoltan_data->graph->getRowMap()->getLocalElementList();
378 auto col_gids = zoltan_data->trans_graph->getRowMap()->getLocalElementList();
379
380 for (size_t i = 0; i < num_local_rows; ++i)
381 {
382 local_ids[i] = i;
383 global_ids[i] = row_gids[i];
384 }
385 for (size_t i = 0; i < num_local_cols; ++i)
386 {
387 local_ids[num_local_rows + i] = num_local_rows + i;
388 global_ids[num_local_rows + i] = num_global_rows + col_gids[i];
389 }
390}
391
393template <typename CrsMatrixType>
394int
395ZoltanCrsColorer<CrsMatrixType>::get_number_of_edges(
396 void *data,
397 int sizeGID,
398 int sizeLID,
399 ZOLTAN_ID_PTR global_id,
400 ZOLTAN_ID_PTR local_id,
401 int *ierr)
402{
403 assert(sizeGID == 1 && sizeLID == 1);
404
405 ZoltanData *zoltan_data = static_cast<ZoltanData *>(data);
406 *ierr = ZOLTAN_OK;
407
408 const size_t num_local_rows = zoltan_data->graph->getLocalNumRows();
409 const ZOLTAN_ID_TYPE lid = *local_id;
410 int num_edges = 0;
411
412 if (lid < num_local_rows)
413 num_edges = zoltan_data->graph->getNumEntriesInLocalRow(lid);
414 else
415 num_edges =
416 zoltan_data->trans_graph->getNumEntriesInLocalRow(lid - num_local_rows);
417
418 return num_edges;
419}
420
422template <typename CrsMatrixType>
423void
424ZoltanCrsColorer<CrsMatrixType>::get_edge_list(
425 void *data,
426 int sizeGID,
427 int sizeLID,
428 ZOLTAN_ID_PTR global_id,
429 ZOLTAN_ID_PTR local_id,
430 ZOLTAN_ID_PTR nbor_global_ids,
431 int *nbor_procs,
432 int wgt_dim,
433 float *ewgts,
434 int *ierr)
435{
436 using Teuchos::Array;
437 using Teuchos::ArrayView;
438 using Teuchos::arrayView;
439
440 ZoltanData *zoltan_data = static_cast<ZoltanData *>(data);
441 *ierr = ZOLTAN_OK;
442
443 const size_t num_local_rows = zoltan_data->graph->getLocalNumRows();
444 const size_t num_global_rows = std::max(
445 static_cast<typename CrsMatrixType::global_ordinal_type>(
446 zoltan_data->graph->getGlobalNumRows()),
447 zoltan_data->graph->getRowMap()->getMaxAllGlobalIndex()+1);
448 const ZOLTAN_ID_TYPE lid = *local_id;
449
450 if (lid < num_local_rows)
451 {
452 const int num_nbr = zoltan_data->graph->getNumEntriesInLocalRow(lid);
453 const auto colMap = zoltan_data->graph->getColMap();
454
455 typename CrsMatrixType::local_inds_host_view_type lcl_ids;
456 zoltan_data->graph->getLocalRowView(lid, lcl_ids);
457
458 for (int j = 0; j < num_nbr; ++j) {
459 nbor_global_ids[j] = num_global_rows
460 + colMap->getGlobalElement(lcl_ids[j]);
461 nbor_procs[j] = zoltan_data->col_procs[lcl_ids[j]];
462 }
463 }
464 else
465 {
466 const int num_nbr =
467 zoltan_data->trans_graph->getNumEntriesInLocalRow(lid-num_local_rows);
468 const auto colMap = zoltan_data->trans_graph->getColMap();
469
470 typename CrsMatrixType::local_inds_host_view_type lcl_ids;
471 zoltan_data->trans_graph->getLocalRowView(lid - num_local_rows, lcl_ids);
472 for (int j = 0; j < num_nbr; ++j)
473 {
474 nbor_global_ids[j] = colMap->getGlobalElement(lcl_ids[j]);
475 nbor_procs[j] = zoltan_data->trans_col_procs[lcl_ids[j]];
476 }
477 }
478}
479
481template <typename CrsMatrixType>
482int
483ZoltanCrsColorer<CrsMatrixType>::sym_get_number_of_vertices(
484 void *data,
485 int *ierr)
486{
487 ZoltanData *zoltan_data = static_cast<ZoltanData *>(data);
488 *ierr = ZOLTAN_OK;
489 return zoltan_data->graph->getLocalNumRows();
490}
491
493template <typename CrsMatrixType>
494void
495ZoltanCrsColorer<CrsMatrixType>::sym_get_vertex_list(
496 void *data,
497 int sizeGID,
498 int sizeLID,
499 ZOLTAN_ID_PTR global_ids,
500 ZOLTAN_ID_PTR local_ids,
501 int wgt_dim,
502 float *obj_wgts,
503 int *ierr)
504{
505 ZoltanData *zoltan_data = static_cast<ZoltanData *>(data);
506 *ierr = ZOLTAN_OK;
507
508 const size_t num_local_rows = zoltan_data->graph->getLocalNumRows();
509 auto row_gids = zoltan_data->graph->getRowMap()->getLocalElementList();
510 for (size_t i = 0; i < num_local_rows; ++i)
511 {
512 local_ids[i] = i;
513 global_ids[i] = row_gids[i];
514 }
515}
516
518template <typename CrsMatrixType>
519int
520ZoltanCrsColorer<CrsMatrixType>::sym_get_number_of_edges(
521 void *data,
522 int sizeGID,
523 int sizeLID,
524 ZOLTAN_ID_PTR global_id,
525 ZOLTAN_ID_PTR local_id,
526 int *ierr)
527{
528 ZoltanData *zoltan_data = static_cast<ZoltanData *>(data);
529 *ierr = ZOLTAN_OK;
530
531 const ZOLTAN_ID_TYPE lid = *local_id;
532 int num_edges = zoltan_data->graph->getNumEntriesInLocalRow(lid);
533 return num_edges;
534}
535
537template <typename CrsMatrixType>
538void
539ZoltanCrsColorer<CrsMatrixType>::sym_get_edge_list(
540 void *data,
541 int sizeGID,
542 int sizeLID,
543 ZOLTAN_ID_PTR global_id,
544 ZOLTAN_ID_PTR local_id,
545 ZOLTAN_ID_PTR nbor_global_ids,
546 int *nbor_procs,
547 int wgt_dim,
548 float *ewgts,
549 int *ierr)
550{
551 using Teuchos::Array;
552 using Teuchos::ArrayView;
553 using Teuchos::arrayView;
554
555 ZoltanData *zoltan_data = static_cast<ZoltanData *>(data);
556 *ierr = ZOLTAN_OK;
557
558 const ZOLTAN_ID_TYPE lid = *local_id;
559 const int num_nbr = zoltan_data->graph->getNumEntriesInLocalRow(lid);
560
561 typename CrsMatrixType::local_inds_host_view_type lcl_ids;
562 zoltan_data->graph->getLocalRowView(lid, lcl_ids);
563 const auto colMap = zoltan_data->graph->getColMap();
564
565 for (int j = 0; j < num_nbr; ++j)
566 {
567 nbor_global_ids[j] = colMap->getGlobalElement(lcl_ids[j]);
568 nbor_procs[j] = zoltan_data->col_procs[lcl_ids[j]];
569 }
570}
571} // namespace Zoltan2
ZoltanCrsColorer(const Teuchos::RCP< matrix_t > &matrix_)
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
Kokkos::View< int *, device_t > list_of_colors_t
matrix_t::global_ordinal_type gno_t
list_of_colors_t::HostMirror list_of_colors_host_t
Teuchos::RCP< Tpetra::CrsGraph< LO, GO, NO > > compute_transpose_graph(const Tpetra::CrsGraph< LO, GO, NO > &graph)
Teuchos::RCP< const typename Tpetra::CrsMatrix< SC, LO, GO, NO >::crs_graph_type > get_graph(const Teuchos::RCP< Tpetra::CrsMatrix< SC, LO, GO, NO > > &matrix)
Created by mbenlioglu on Aug 31, 2020.