Tpetra parallel linear algebra Version of the Day
Loading...
Searching...
No Matches
TpetraExt_MatrixMatrix_def.hpp
Go to the documentation of this file.
1// @HEADER
2// *****************************************************************************
3// Tpetra: Templated Linear Algebra Services Package
4//
5// Copyright 2008 NTESS and the Tpetra contributors.
6// SPDX-License-Identifier: BSD-3-Clause
7// *****************************************************************************
8// @HEADER
9
10#ifndef TPETRA_MATRIXMATRIX_DEF_HPP
11#define TPETRA_MATRIXMATRIX_DEF_HPP
13#include "KokkosSparse_Utils.hpp"
14#include "Tpetra_ConfigDefs.hpp"
16#include "Teuchos_VerboseObject.hpp"
17#include "Teuchos_Array.hpp"
18#include "Tpetra_Util.hpp"
19#include "Tpetra_CrsMatrix.hpp"
20#include "Tpetra_BlockCrsMatrix.hpp"
22#include "Tpetra_RowMatrixTransposer.hpp"
25#include "Tpetra_Details_makeColMap.hpp"
26#include "Tpetra_ConfigDefs.hpp"
27#include "Tpetra_Map.hpp"
28#include "Tpetra_Export.hpp"
32
33#include <algorithm>
34#include <type_traits>
35#include "Teuchos_FancyOStream.hpp"
36
37#include "TpetraExt_MatrixMatrix_ExtraKernels_def.hpp"
39
40#include "KokkosSparse_spgemm.hpp"
41#include "KokkosSparse_spadd.hpp"
42#include "Kokkos_Bitset.hpp"
43
44#include <MatrixMarket_Tpetra.hpp>
45
50
51/*********************************************************************************************************/
52// Include the architecture-specific kernel partial specializations here
53// NOTE: This needs to be outside all namespaces
54#include "TpetraExt_MatrixMatrix_OpenMP.hpp"
55#include "TpetraExt_MatrixMatrix_Cuda.hpp"
56#include "TpetraExt_MatrixMatrix_HIP.hpp"
57#include "TpetraExt_MatrixMatrix_SYCL.hpp"
58
59namespace Tpetra {
60
61namespace MatrixMatrix {
62
63//
64// This method forms the matrix-matrix product C = op(A) * op(B), where
65// op(A) == A if transposeA is false,
66// op(A) == A^T if transposeA is true,
67// and similarly for op(B).
68//
69template <class Scalar,
70 class LocalOrdinal,
71 class GlobalOrdinal,
72 class Node>
75 bool transposeA,
77 bool transposeB,
79 bool call_FillComplete_on_result,
80 const std::string& label,
81 const Teuchos::RCP<Teuchos::ParameterList>& params) {
82 using Teuchos::null;
83 using Teuchos::RCP;
84 using Teuchos::rcp;
85 typedef Scalar SC;
86 typedef LocalOrdinal LO;
87 typedef GlobalOrdinal GO;
88 typedef Node NO;
89 typedef CrsMatrix<SC, LO, GO, NO> crs_matrix_type;
91 typedef CrsMatrixStruct<SC, LO, GO, NO> crs_matrix_struct_type;
93 typedef RowMatrixTransposer<SC, LO, GO, NO> transposer_type;
94
95 RCP<Tpetra::Details::ProfilingRegion> MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: MMM: All Setup"));
96
97 const std::string prefix = "TpetraExt::MatrixMatrix::Multiply(): ";
98
99 // TEUCHOS_FUNC_TIME_MONITOR_DIFF("My Matrix Mult", mmm_multiply);
100
101 // The input matrices A and B must both be fillComplete.
102 TEUCHOS_TEST_FOR_EXCEPTION(!A.isFillComplete(), std::runtime_error, prefix << "Matrix A is not fill complete.");
103 TEUCHOS_TEST_FOR_EXCEPTION(!B.isFillComplete(), std::runtime_error, prefix << "Matrix B is not fill complete.");
104
105 // If transposeA is true, then Aprime will be the transpose of A
106 // (computed explicitly via RowMatrixTransposer). Otherwise, Aprime
107 // will just be a pointer to A.
108 RCP<const crs_matrix_type> Aprime = null;
109 // If transposeB is true, then Bprime will be the transpose of B
110 // (computed explicitly via RowMatrixTransposer). Otherwise, Bprime
111 // will just be a pointer to B.
112 RCP<const crs_matrix_type> Bprime = null;
113
114 // Is this a "clean" matrix?
115 //
116 // mfh 27 Sep 2016: Historically, if Epetra_CrsMatrix was neither
117 // locally nor globally indexed, then it was empty. I don't like
118 // this, because the most straightforward implementation presumes
119 // lazy allocation of indices. However, historical precedent
120 // demands that we keep around this predicate as a way to test
121 // whether the matrix is empty.
122 const bool newFlag = !C.getGraph()->isLocallyIndexed() && !C.getGraph()->isGloballyIndexed();
123
124 bool use_optimized_ATB = false;
125 if (transposeA && !transposeB && call_FillComplete_on_result && newFlag)
126 use_optimized_ATB = true;
127
128#ifdef USE_OLD_TRANSPOSE // NOTE: For Grey Ballard's use. Remove this later.
129 use_optimized_ATB = false;
130#endif
131
132 using Teuchos::ParameterList;
133 RCP<ParameterList> transposeParams(new ParameterList);
134#if KOKKOSKERNELS_VERSION >= 50299
135 transposeParams->set("sort", false);
136#else
137 transposeParams->set("sort", true); // Kokkos Kernels spgemm requires inputs to be sorted
138#endif
139
140 if (!use_optimized_ATB && transposeA) {
141 transposer_type transposer(rcpFromRef(A));
142 Aprime = transposer.createTranspose(transposeParams);
143 } else {
144 Aprime = rcpFromRef(A);
145 }
146
147 if (transposeB) {
148 transposer_type transposer(rcpFromRef(B));
149 Bprime = transposer.createTranspose(transposeParams);
150 } else {
151 Bprime = rcpFromRef(B);
152 }
153
154 // Check size compatibility
155 global_size_t numACols = A.getDomainMap()->getGlobalNumElements();
156 global_size_t numBCols = B.getDomainMap()->getGlobalNumElements();
157 global_size_t Aouter = transposeA ? numACols : A.getGlobalNumRows();
158 global_size_t Bouter = transposeB ? B.getGlobalNumRows() : numBCols;
159 global_size_t Ainner = transposeA ? A.getGlobalNumRows() : numACols;
160 global_size_t Binner = transposeB ? numBCols : B.getGlobalNumRows();
161 TEUCHOS_TEST_FOR_EXCEPTION(Ainner != Binner, std::runtime_error,
162 prefix << "ERROR, inner dimensions of op(A) and op(B) "
163 "must match for matrix-matrix product. op(A) is "
164 << Aouter << "x" << Ainner << ", op(B) is " << Binner << "x" << Bouter);
165
166 // The result matrix C must at least have a row-map that reflects the correct
167 // row-size. Don't check the number of columns because rectangular matrices
168 // which were constructed with only one map can still end up having the
169 // correct capacity and dimensions when filled.
170 TEUCHOS_TEST_FOR_EXCEPTION(Aouter > C.getGlobalNumRows(), std::runtime_error,
171 prefix << "ERROR, dimensions of result C must "
172 "match dimensions of op(A) * op(B). C has "
173 << C.getGlobalNumRows()
174 << " rows, should have at least " << Aouter << std::endl);
175
176 // It doesn't matter whether C is already Filled or not. If it is already
177 // Filled, it must have space allocated for the positions that will be
178 // referenced in forming C = op(A)*op(B). If it doesn't have enough space,
179 // we'll error out later when trying to store result values.
180
181 // CGB: However, matrix must be in active-fill
182 if (!C.isFillActive()) C.resumeFill();
183
184 // We're going to need to import remotely-owned sections of A and/or B if
185 // more than one processor is performing this run, depending on the scenario.
186 int numProcs = A.getComm()->getSize();
187
188 // Declare a couple of structs that will be used to hold views of the data
189 // of A and B, to be used for fast access during the matrix-multiplication.
190 crs_matrix_struct_type Aview;
191 crs_matrix_struct_type Bview;
192
193 RCP<const map_type> targetMap_A = Aprime->getRowMap();
194 RCP<const map_type> targetMap_B = Bprime->getRowMap();
195
196 {
197 Tpetra::Details::ProfilingRegion r("TpetraExt: MMM: All I&X");
198
199 // Now import any needed remote rows and populate the Aview struct
200 // NOTE: We assert that an import isn't needed --- since we do the transpose
201 // above to handle that.
202 if (!use_optimized_ATB) {
203 RCP<const import_type> dummyImporter;
204 MMdetails::import_and_extract_views(*Aprime, targetMap_A, Aview, dummyImporter, true, label, params);
205 }
206
207 // We will also need local access to all rows of B that correspond to the
208 // column-map of op(A).
209 if (numProcs > 1)
210 targetMap_B = Aprime->getColMap();
211
212 // Import any needed remote rows and populate the Bview struct.
213 if (!use_optimized_ATB)
214 MMdetails::import_and_extract_views(*Bprime, targetMap_B, Bview, Aprime->getGraph()->getImporter(), Aprime->getGraph()->getImporter().is_null(), label, params);
215
216 } // stop MM_importExtract here
217
218 // stop the setup timer, and start the multiply timer
219 MM = Teuchos::null;
220 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: MMM: All Multiply"));
221
222 // Call the appropriate method to perform the actual multiplication.
223 if (use_optimized_ATB) {
224 MMdetails::mult_AT_B_newmatrix(A, B, C, label, params);
225
226 } else if (call_FillComplete_on_result && newFlag) {
227 MMdetails::mult_A_B_newmatrix(Aview, Bview, C, label, params);
228
229 } else if (call_FillComplete_on_result) {
230 MMdetails::mult_A_B_reuse(Aview, Bview, C, label, params);
231
232 } else {
233 // mfh 27 Sep 2016: Is this the "slow" case? This
234 // "CrsWrapper_CrsMatrix" thing could perhaps be made to support
235 // thread-parallel inserts, but that may take some effort.
236 CrsWrapper_CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node> crsmat(C);
237
238 MMdetails::mult_A_B(Aview, Bview, crsmat, label, params);
239 }
240
241 MM = Teuchos::null;
242 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: MMM: All FillComplete"));
243 if (call_FillComplete_on_result && !C.isFillComplete()) {
244 // We'll call FillComplete on the C matrix before we exit, and give it a
245 // domain-map and a range-map.
246 // The domain-map will be the domain-map of B, unless
247 // op(B)==transpose(B), in which case the range-map of B will be used.
248 // The range-map will be the range-map of A, unless op(A)==transpose(A),
249 // in which case the domain-map of A will be used.
250 C.fillComplete(Bprime->getDomainMap(), Aprime->getRangeMap());
251 }
252}
253
254//
255// This method forms the matrix-matrix product C = op(A) * op(B), where
256// op(A) == A and similarly for op(B). op(A) = A^T is not yet implemented.
257//
258template <class Scalar,
259 class LocalOrdinal,
260 class GlobalOrdinal,
261 class Node>
264 bool transposeA,
266 bool transposeB,
268 const std::string& label) {
269 using Teuchos::null;
270 using Teuchos::RCP;
271 using Teuchos::rcp;
272 typedef Scalar SC;
273 typedef LocalOrdinal LO;
274 typedef GlobalOrdinal GO;
275 typedef Node NO;
276 typedef BlockCrsMatrixStruct<SC, LO, GO, NO> blockcrs_matrix_struct_type;
279
280 std::string prefix = std::string("TpetraExt ") + label + std::string(": ");
281
282 TEUCHOS_TEST_FOR_EXCEPTION(transposeA == true, std::runtime_error, prefix << "Matrix A cannot be transposed.");
283 TEUCHOS_TEST_FOR_EXCEPTION(transposeB == true, std::runtime_error, prefix << "Matrix B cannot be transposed.");
284
285 // Check size compatibility
286 global_size_t numACols = A->getGlobalNumCols();
287 global_size_t numBCols = B->getGlobalNumCols();
288 global_size_t numARows = A->getGlobalNumRows();
289 global_size_t numBRows = B->getGlobalNumRows();
290
291 global_size_t Aouter = numARows;
292 global_size_t Bouter = numBCols;
293 global_size_t Ainner = numACols;
294 global_size_t Binner = numBRows;
295 TEUCHOS_TEST_FOR_EXCEPTION(Ainner != Binner, std::runtime_error,
296 prefix << "ERROR, inner dimensions of op(A) and op(B) "
297 "must match for matrix-matrix product. op(A) is "
298 << Aouter << "x" << Ainner << ", op(B) is " << Binner << "x" << Bouter);
299
300 // We're going to need to import remotely-owned sections of A and/or B if
301 // more than one processor is performing this run, depending on the scenario.
302 int numProcs = A->getComm()->getSize();
303
304 const LO blocksize = A->getBlockSize();
305 TEUCHOS_TEST_FOR_EXCEPTION(blocksize != B->getBlockSize(), std::runtime_error,
306 prefix << "ERROR, Blocksizes do not match. A.blocksize = " << blocksize << ", B.blocksize = " << B->getBlockSize());
307
308 // Declare a couple of structs that will be used to hold views of the data
309 // of A and B, to be used for fast access during the matrix-multiplication.
310 blockcrs_matrix_struct_type Aview(blocksize);
311 blockcrs_matrix_struct_type Bview(blocksize);
312
313 RCP<const map_type> targetMap_A = A->getRowMap();
314 RCP<const map_type> targetMap_B = B->getRowMap();
315
316 // Populate the Aview struct. No remotes are needed.
317 RCP<const import_type> dummyImporter;
318 MMdetails::import_and_extract_views(*A, targetMap_A, Aview, dummyImporter, true);
319
320 // We will also need local access to all rows of B that correspond to the
321 // column-map of op(A).
322 if (numProcs > 1)
323 targetMap_B = A->getColMap();
324
325 // Import any needed remote rows and populate the Bview struct.
326 MMdetails::import_and_extract_views(*B, targetMap_B, Bview, A->getGraph()->getImporter(),
327 A->getGraph()->getImporter().is_null());
328
329 // Call the appropriate method to perform the actual multiplication.
330 MMdetails::mult_A_B_newmatrix(Aview, Bview, C);
331}
332
333template <class Scalar,
334 class LocalOrdinal,
335 class GlobalOrdinal,
336 class Node>
337void Jacobi(typename Teuchos::ScalarTraits<Scalar>::magnitudeType omega,
342 bool call_FillComplete_on_result,
343 const std::string& label,
344 const Teuchos::RCP<Teuchos::ParameterList>& params) {
345 using Teuchos::RCP;
346 using Teuchos::rcp;
347 typedef Scalar SC;
348 typedef LocalOrdinal LO;
349 typedef GlobalOrdinal GO;
350 typedef Node NO;
352 typedef CrsMatrixStruct<SC, LO, GO, NO> crs_matrix_struct_type;
354 typedef CrsMatrix<SC, LO, GO, NO> crs_matrix_type;
355
356 RCP<Tpetra::Details::ProfilingRegion> MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: Jacobi: All Setup"));
357
358 const std::string prefix = "TpetraExt::MatrixMatrix::Jacobi(): ";
359
360 // A and B should already be Filled.
361 // Should we go ahead and call FillComplete() on them if necessary or error
362 // out? For now, we choose to error out.
363 TEUCHOS_TEST_FOR_EXCEPTION(!A.isFillComplete(), std::runtime_error, prefix << "Matrix A is not fill complete.");
364 TEUCHOS_TEST_FOR_EXCEPTION(!B.isFillComplete(), std::runtime_error, prefix << "Matrix B is not fill complete.");
365
366 RCP<const crs_matrix_type> Aprime = rcpFromRef(A);
367 RCP<const crs_matrix_type> Bprime = rcpFromRef(B);
368
369 // Now check size compatibility
370 global_size_t numACols = A.getDomainMap()->getGlobalNumElements();
371 global_size_t numBCols = B.getDomainMap()->getGlobalNumElements();
372 global_size_t Aouter = A.getGlobalNumRows();
373 global_size_t Bouter = numBCols;
374 global_size_t Ainner = numACols;
375 global_size_t Binner = B.getGlobalNumRows();
376 TEUCHOS_TEST_FOR_EXCEPTION(Ainner != Binner, std::runtime_error,
377 prefix << "ERROR, inner dimensions of op(A) and op(B) "
378 "must match for matrix-matrix product. op(A) is "
379 << Aouter << "x" << Ainner << ", op(B) is " << Binner << "x" << Bouter);
380
381 // The result matrix C must at least have a row-map that reflects the correct
382 // row-size. Don't check the number of columns because rectangular matrices
383 // which were constructed with only one map can still end up having the
384 // correct capacity and dimensions when filled.
385 TEUCHOS_TEST_FOR_EXCEPTION(Aouter > C.getGlobalNumRows(), std::runtime_error,
386 prefix << "ERROR, dimensions of result C must "
387 "match dimensions of op(A) * op(B). C has "
388 << C.getGlobalNumRows()
389 << " rows, should have at least " << Aouter << std::endl);
390
391 // It doesn't matter whether C is already Filled or not. If it is already
392 // Filled, it must have space allocated for the positions that will be
393 // referenced in forming C = op(A)*op(B). If it doesn't have enough space,
394 // we'll error out later when trying to store result values.
395
396 // CGB: However, matrix must be in active-fill
397 TEUCHOS_TEST_FOR_EXCEPT(C.isFillActive() == false);
398
399 // We're going to need to import remotely-owned sections of A and/or B if
400 // more than one processor is performing this run, depending on the scenario.
401 int numProcs = A.getComm()->getSize();
402
403 // Declare a couple of structs that will be used to hold views of the data of
404 // A and B, to be used for fast access during the matrix-multiplication.
405 crs_matrix_struct_type Aview;
406 crs_matrix_struct_type Bview;
407
408 RCP<const map_type> targetMap_A = Aprime->getRowMap();
409 RCP<const map_type> targetMap_B = Bprime->getRowMap();
410
411 {
412 Tpetra::Details::ProfilingRegion r("TpetraExt: Jacobi: All I&X");
413 // Disable globalConstants by default
414 // NOTE: the I&X routine sticks an importer on the paramlist as output, so we have to use a unique guy here
415 RCP<Teuchos::ParameterList> importParams = Teuchos::rcp(new Teuchos::ParameterList);
416 importParams->set("compute global constants", false);
417 if (!params.is_null()) {
418 importParams->setParameters(*params);
419 importParams->set("compute global constants", params->get("compute global constants: temporaries", false));
420 }
421
422 // Now import any needed remote rows and populate the Aview struct.
423 RCP<const import_type> dummyImporter;
424 MMdetails::import_and_extract_views(*Aprime, targetMap_A, Aview, dummyImporter, true, label, importParams);
425
426 // We will also need local access to all rows of B that correspond to the
427 // column-map of op(A).
428 if (numProcs > 1)
429 targetMap_B = Aprime->getColMap();
430
431 // Now import any needed remote rows and populate the Bview struct.
432 MMdetails::import_and_extract_views(*Bprime, targetMap_B, Bview, Aprime->getGraph()->getImporter(), Aprime->getGraph()->getImporter().is_null(), label, importParams);
433 }
434 MM = Teuchos::null;
435 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: Jacobi All Multiply"));
436
437 // Now call the appropriate method to perform the actual multiplication.
438 CrsWrapper_CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node> crsmat(C);
439
440 // Is this a "clean" matrix
441 bool newFlag = !C.getGraph()->isLocallyIndexed() && !C.getGraph()->isGloballyIndexed();
442
443 if (call_FillComplete_on_result && newFlag) {
444 MMdetails::jacobi_A_B_newmatrix(omega, Dinv, Aview, Bview, C, label, params);
445
446 } else if (call_FillComplete_on_result) {
447 MMdetails::jacobi_A_B_reuse(omega, Dinv, Aview, Bview, C, label, params);
448
449 } else {
450 TEUCHOS_TEST_FOR_EXCEPTION(true, std::runtime_error, "jacobi_A_B_general not implemented");
451 }
452
453 if (!params.is_null()) {
454 bool removeZeroEntries = params->get("remove zeros", false);
455 if (removeZeroEntries) {
456 typedef Teuchos::ScalarTraits<Scalar> STS;
457 typename STS::magnitudeType threshold = params->get("remove zeros threshold", STS::magnitude(STS::zero()));
458 removeCrsMatrixZeros(C, threshold);
459 }
460 }
461}
462
463template <class Scalar,
464 class LocalOrdinal,
465 class GlobalOrdinal,
466 class Node>
467void Add(
469 bool transposeA,
470 Scalar scalarA,
472 Scalar scalarB) {
473 using Teuchos::Array;
474 using Teuchos::null;
475 using Teuchos::RCP;
476 typedef Scalar SC;
477 typedef LocalOrdinal LO;
478 typedef GlobalOrdinal GO;
479 typedef Node NO;
480 typedef CrsMatrix<SC, LO, GO, NO> crs_matrix_type;
481 typedef RowMatrixTransposer<SC, LO, GO, NO> transposer_type;
482
483 const std::string prefix = "TpetraExt::MatrixMatrix::Add(): ";
484
485 TEUCHOS_TEST_FOR_EXCEPTION(!A.isFillComplete(), std::runtime_error,
486 prefix << "ERROR, input matrix A.isFillComplete() is false; it is required to be true. "
487 "(Result matrix B is not required to be isFillComplete()).");
488 TEUCHOS_TEST_FOR_EXCEPTION(B.isFillComplete(), std::runtime_error,
489 prefix << "ERROR, input matrix B must not be fill complete!");
490 TEUCHOS_TEST_FOR_EXCEPTION(B.isStaticGraph(), std::runtime_error,
491 prefix << "ERROR, input matrix B must not have static graph!");
492 TEUCHOS_TEST_FOR_EXCEPTION(B.isLocallyIndexed(), std::runtime_error,
493 prefix << "ERROR, input matrix B must not be locally indexed!");
494
495 using Teuchos::ParameterList;
496 RCP<ParameterList> transposeParams(new ParameterList);
497 transposeParams->set("sort", false);
498
499 RCP<const crs_matrix_type> Aprime = null;
500 if (transposeA) {
501 transposer_type transposer(rcpFromRef(A));
502 Aprime = transposer.createTranspose(transposeParams);
503 } else {
504 Aprime = rcpFromRef(A);
505 }
506
507 size_t a_numEntries;
508 typename crs_matrix_type::nonconst_global_inds_host_view_type a_inds("a_inds", A.getLocalMaxNumRowEntries());
509 typename crs_matrix_type::nonconst_values_host_view_type a_vals("a_vals", A.getLocalMaxNumRowEntries());
510 GO row;
511
512 if (scalarB != Teuchos::ScalarTraits<SC>::one())
513 B.scale(scalarB);
514
515 size_t numMyRows = B.getLocalNumRows();
516 if (scalarA != Teuchos::ScalarTraits<SC>::zero()) {
517 for (LO i = 0; (size_t)i < numMyRows; ++i) {
518 row = B.getRowMap()->getGlobalElement(i);
519 Aprime->getGlobalRowCopy(row, a_inds, a_vals, a_numEntries);
520
521 if (scalarA != Teuchos::ScalarTraits<SC>::one()) {
522 for (size_t j = 0; j < a_numEntries; ++j)
523 a_vals[j] *= scalarA;
524 }
525 B.insertGlobalValues(row, a_numEntries, reinterpret_cast<Scalar*>(a_vals.data()), a_inds.data());
526 }
527 }
528}
529
530template <class Scalar,
531 class LocalOrdinal,
532 class GlobalOrdinal,
533 class Node>
534Teuchos::RCP<CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>>
535add(const Scalar& alpha,
536 const bool transposeA,
538 const Scalar& beta,
539 const bool transposeB,
541 const Teuchos::RCP<const Map<LocalOrdinal, GlobalOrdinal, Node>>& domainMap,
542 const Teuchos::RCP<const Map<LocalOrdinal, GlobalOrdinal, Node>>& rangeMap,
543 const Teuchos::RCP<Teuchos::ParameterList>& params) {
544 using Teuchos::ParameterList;
545 using Teuchos::RCP;
546 using Teuchos::rcp;
547 using Teuchos::rcpFromRef;
549 if (!params.is_null()) {
550 TEUCHOS_TEST_FOR_EXCEPTION(
551 params->isParameter("Call fillComplete") && !params->get<bool>("Call fillComplete"),
552 std::invalid_argument,
553 "Tpetra::MatrixMatrix::add(): this version of add() always calls fillComplete\n"
554 "on the result, but you explicitly set 'Call fillComplete' = false in the parameter list. Don't set this explicitly.");
555 params->set("Call fillComplete", true);
556 }
557 // If transposeB, must compute B's explicit transpose to
558 // get the correct row map for C.
559 RCP<const crs_matrix_type> Brcp = rcpFromRef(B);
560 if (transposeB) {
562 Brcp = transposer.createTranspose();
563 }
564 // Check that A,B are fillComplete before getting B's column map
565 TEUCHOS_TEST_FOR_EXCEPTION(!A.isFillComplete() || !Brcp->isFillComplete(), std::invalid_argument,
566 "TpetraExt::MatrixMatrix::add(): A and B must both be fill complete.");
567 RCP<crs_matrix_type> C = rcp(new crs_matrix_type(Brcp->getRowMap(), 0));
568 // this version of add() always fill completes the result, no matter what is in params on input
569 add(alpha, transposeA, A, beta, false, *Brcp, *C, domainMap, rangeMap, params);
570 return C;
571}
572
573// This functor does the same thing as CrsGraph::convertColumnIndicesFromGlobalToLocal,
574// but since the spadd() output is always packed there is no need for a separate
575// numRowEntries here.
576//
577template <class LO, class GO, class LOView, class GOView, class LocalMap>
578struct ConvertGlobalToLocalFunctor {
579 ConvertGlobalToLocalFunctor(LOView& lids_, const GOView& gids_, const LocalMap localColMap_)
580 : lids(lids_)
581 , gids(gids_)
582 , localColMap(localColMap_) {}
583
584 KOKKOS_FUNCTION void operator()(const GO i) const {
585 lids(i) = localColMap.getLocalElement(gids(i));
586 }
587
588 LOView lids;
589 const GOView gids;
590 const LocalMap localColMap;
591};
592
593template <class Scalar,
594 class LocalOrdinal,
595 class GlobalOrdinal,
596 class Node>
597void add(const Scalar& alpha,
598 const bool transposeA,
600 const Scalar& beta,
601 const bool transposeB,
604 const Teuchos::RCP<const Map<LocalOrdinal, GlobalOrdinal, Node>>& domainMap,
605 const Teuchos::RCP<const Map<LocalOrdinal, GlobalOrdinal, Node>>& rangeMap,
606 const Teuchos::RCP<Teuchos::ParameterList>& params) {
607 using Teuchos::RCP;
608 using Teuchos::rcp;
609 using Teuchos::rcp_dynamic_cast;
610 using Teuchos::rcp_implicit_cast;
611 using Teuchos::rcpFromRef;
612 using Teuchos::TimeMonitor;
613 using SC = Scalar;
614 using LO = LocalOrdinal;
615 using GO = GlobalOrdinal;
616 using NO = Node;
617 using crs_matrix_type = CrsMatrix<SC, LO, GO, NO>;
620 using transposer_type = RowMatrixTransposer<SC, LO, GO, NO>;
623 using exec_space = typename crs_graph_type::execution_space;
624 using AddKern = MMdetails::AddKernels<SC, LO, GO, NO>;
625 const char* prefix_mmm = "TpetraExt::MatrixMatrix::add: ";
626 constexpr bool debug = false;
627
628 RCP<Tpetra::Details::ProfilingRegion> MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: add: Transpose"));
629
630 if (debug) {
631 std::ostringstream os;
632 os << "Proc " << A.getMap()->getComm()->getRank() << ": "
633 << "TpetraExt::MatrixMatrix::add" << std::endl;
634 std::cerr << os.str();
635 }
636
637 TEUCHOS_TEST_FOR_EXCEPTION(C.isLocallyIndexed() || C.isGloballyIndexed(), std::invalid_argument,
638 prefix_mmm << "C must be a 'new' matrix (neither locally nor globally indexed).");
639 TEUCHOS_TEST_FOR_EXCEPTION(!A.isFillComplete() || !B.isFillComplete(), std::invalid_argument,
640 prefix_mmm << "A and B must both be fill complete.");
641#ifdef HAVE_TPETRA_DEBUG
642 // The matrices don't have domain or range Maps unless they are fill complete.
643 if (A.isFillComplete() && B.isFillComplete()) {
644 const bool domainMapsSame =
645 (!transposeA && !transposeB &&
646 !A.getDomainMap()->locallySameAs(*B.getDomainMap())) ||
647 (!transposeA && transposeB &&
648 !A.getDomainMap()->isSameAs(*B.getRangeMap())) ||
649 (transposeA && !transposeB &&
650 !A.getRangeMap()->isSameAs(*B.getDomainMap()));
651 TEUCHOS_TEST_FOR_EXCEPTION(domainMapsSame, std::invalid_argument,
652 prefix_mmm << "The domain Maps of Op(A) and Op(B) are not the same.");
653
654 const bool rangeMapsSame =
655 (!transposeA && !transposeB &&
656 !A.getRangeMap()->isSameAs(*B.getRangeMap())) ||
657 (!transposeA && transposeB &&
658 !A.getRangeMap()->isSameAs(*B.getDomainMap())) ||
659 (transposeA && !transposeB &&
660 !A.getDomainMap()->isSameAs(*B.getRangeMap()));
661 TEUCHOS_TEST_FOR_EXCEPTION(rangeMapsSame, std::invalid_argument,
662 prefix_mmm << "The range Maps of Op(A) and Op(B) are not the same.");
663 }
664#endif // HAVE_TPETRA_DEBUG
665
666 using Teuchos::ParameterList;
667 // Form the explicit transpose of A if necessary.
668 RCP<const crs_matrix_type> Aprime = rcpFromRef(A);
669 if (transposeA) {
670 transposer_type transposer(Aprime);
671 Aprime = transposer.createTranspose();
672 }
673
674#ifdef HAVE_TPETRA_DEBUG
675 TEUCHOS_TEST_FOR_EXCEPTION(Aprime.is_null(), std::logic_error,
676 prefix_mmm << "Failed to compute Op(A). "
677 "Please report this bug to the Tpetra developers.");
678#endif // HAVE_TPETRA_DEBUG
679
680 // Form the explicit transpose of B if necessary.
681 RCP<const crs_matrix_type> Bprime = rcpFromRef(B);
682 if (transposeB) {
683 if (debug) {
684 std::ostringstream os;
685 os << "Proc " << A.getMap()->getComm()->getRank() << ": "
686 << "Form explicit xpose of B" << std::endl;
687 std::cerr << os.str();
688 }
689 transposer_type transposer(Bprime);
690 Bprime = transposer.createTranspose();
691 }
692#ifdef HAVE_TPETRA_DEBUG
693 TEUCHOS_TEST_FOR_EXCEPTION(Bprime.is_null(), std::logic_error,
694 prefix_mmm << "Failed to compute Op(B). Please report this bug to the Tpetra developers.");
695 TEUCHOS_TEST_FOR_EXCEPTION(
696 !Aprime->isFillComplete() || !Bprime->isFillComplete(), std::invalid_argument,
697 prefix_mmm << "Aprime and Bprime must both be fill complete. "
698 "Please report this bug to the Tpetra developers.");
699#endif // HAVE_TPETRA_DEBUG
700 RCP<const map_type> CDomainMap = domainMap;
701 RCP<const map_type> CRangeMap = rangeMap;
702 if (CDomainMap.is_null()) {
703 CDomainMap = Bprime->getDomainMap();
704 }
705 if (CRangeMap.is_null()) {
706 CRangeMap = Bprime->getRangeMap();
707 }
708 assert(!(CDomainMap.is_null()));
709 assert(!(CRangeMap.is_null()));
710 typedef typename AddKern::values_array values_array;
711 typedef typename AddKern::row_ptrs_array row_ptrs_array;
712 typedef typename AddKern::col_inds_array col_inds_array;
713 bool AGraphSorted = Aprime->getCrsGraph()->isSorted();
714 bool BGraphSorted = Bprime->getCrsGraph()->isSorted();
715 values_array vals;
716 row_ptrs_array rowptrs;
717 col_inds_array colinds;
718
719 MM = Teuchos::null;
720 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: add: rowmap check/import"));
721
722 if (!(Aprime->getRowMap()->isSameAs(*(Bprime->getRowMap())))) {
723 // import Aprime into Bprime's row map so the local matrices have same # of rows
724 auto import = rcp(new import_type(Aprime->getRowMap(), Bprime->getRowMap()));
725 // cbl do not set
726 // parameterlist "isMatrixMatrix_TransferAndFillComplete" true here as
727 // this import _may_ take the form of a transfer. In practice it would be unlikely,
728 // but the general case is not so forgiving.
729 Aprime = importAndFillCompleteCrsMatrix<crs_matrix_type>(Aprime, *import, Bprime->getDomainMap(), Bprime->getRangeMap());
730 }
731 bool matchingColMaps = Aprime->getColMap()->isSameAs(*(Bprime->getColMap()));
732 bool sorted = AGraphSorted && BGraphSorted;
733 RCP<const import_type> Cimport = Teuchos::null;
734 RCP<export_type> Cexport = Teuchos::null;
735 bool doFillComplete = true;
736 if (Teuchos::nonnull(params) && params->isParameter("Call fillComplete")) {
737 doFillComplete = params->get<bool>("Call fillComplete");
738 }
739 auto Alocal = Aprime->getLocalMatrixDevice();
740 auto Blocal = Bprime->getLocalMatrixDevice();
741 LO numLocalRows = Alocal.numRows();
742 if (numLocalRows == 0) {
743 // KokkosKernels spadd assumes rowptrs.extent(0) + 1 == nrows,
744 // but an empty Tpetra matrix is allowed to have rowptrs.extent(0) == 0.
745 // Handle this case now
746 //(without interfering with collective operations, since it's possible for
747 // some ranks to have 0 local rows and others not).
748 rowptrs = row_ptrs_array("C rowptrs", 0);
749 }
750 auto Acolmap = Aprime->getColMap();
751 auto Bcolmap = Bprime->getColMap();
752 if (!matchingColMaps) {
753 using global_col_inds_array = typename AddKern::global_col_inds_array;
754 MM = Teuchos::null;
755 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: add: mismatched col map full kernel"));
756
757 // use kernel that converts col indices in both A and B to common domain map before adding
758 auto AlocalColmap = Acolmap->getLocalMap();
759 auto BlocalColmap = Bcolmap->getLocalMap();
760 global_col_inds_array globalColinds;
761 if (debug) {
762 std::ostringstream os;
763 os << "Proc " << A.getMap()->getComm()->getRank() << ": "
764 << "Call AddKern::convertToGlobalAndAdd(...)" << std::endl;
765 std::cerr << os.str();
766 }
767 AddKern::convertToGlobalAndAdd(
768 Alocal, alpha, Blocal, beta, AlocalColmap, BlocalColmap,
769 vals, rowptrs, globalColinds);
770 if (debug) {
771 std::ostringstream os;
772 os << "Proc " << A.getMap()->getComm()->getRank() << ": "
773 << "Finished AddKern::convertToGlobalAndAdd(...)" << std::endl;
774 std::cerr << os.str();
775 }
776 MM = Teuchos::null;
777 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: add: Constructing graph"));
778
779 RCP<const map_type> CcolMap;
780 Tpetra::Details::makeColMap<LocalOrdinal, GlobalOrdinal, Node>(CcolMap, CDomainMap, globalColinds);
781 C.replaceColMap(CcolMap);
782 col_inds_array localColinds("C colinds", globalColinds.extent(0));
783 Kokkos::parallel_for(Kokkos::RangePolicy<exec_space>(0, globalColinds.extent(0)),
784 ConvertGlobalToLocalFunctor<LocalOrdinal, GlobalOrdinal,
785 col_inds_array, global_col_inds_array,
786 typename map_type::local_map_type>(localColinds, globalColinds, CcolMap->getLocalMap()));
787 Import_Util::sortCrsEntries(rowptrs, localColinds, vals);
788 C.setAllValues(rowptrs, localColinds, vals);
789 C.fillComplete(CDomainMap, CRangeMap, params);
790 if (!doFillComplete)
791 C.resumeFill();
792 } else {
793 // Aprime, Bprime and C all have the same column maps
794 auto Avals = Alocal.values;
795 auto Bvals = Blocal.values;
796 auto Arowptrs = Alocal.graph.row_map;
797 auto Browptrs = Blocal.graph.row_map;
798 auto Acolinds = Alocal.graph.entries;
799 auto Bcolinds = Blocal.graph.entries;
800 if (sorted) {
801 // use sorted kernel
802 MM = Teuchos::null;
803 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: add: sorted entries full kernel"));
804 if (debug) {
805 std::ostringstream os;
806 os << "Proc " << A.getMap()->getComm()->getRank() << ": "
807 << "Call AddKern::addSorted(...)" << std::endl;
808 std::cerr << os.str();
809 }
810 AddKern::addSorted(Avals, Arowptrs, Acolinds, alpha, Bvals, Browptrs, Bcolinds, beta, Aprime->getGlobalNumCols(), vals, rowptrs, colinds);
811 } else {
812 // use unsorted kernel
813 MM = Teuchos::null;
814 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: add: unsorted entries full kernel"));
815
816 if (debug) {
817 std::ostringstream os;
818 os << "Proc " << A.getMap()->getComm()->getRank() << ": "
819 << "Call AddKern::addUnsorted(...)" << std::endl;
820 std::cerr << os.str();
821 }
822 AddKern::addUnsorted(Avals, Arowptrs, Acolinds, alpha, Bvals, Browptrs, Bcolinds, beta, Aprime->getGlobalNumCols(), vals, rowptrs, colinds);
823 }
824 // Bprime col map works as C's row map, since Aprime and Bprime have the same colmaps.
825 RCP<const map_type> Ccolmap = Bcolmap;
826 C.replaceColMap(Ccolmap);
827 C.setAllValues(rowptrs, colinds, vals);
828 if (doFillComplete) {
829 MM = Teuchos::null;
830 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: add: expertStaticFillComplete"));
831 if (!CDomainMap->isSameAs(*Ccolmap)) {
832 if (debug) {
833 std::ostringstream os;
834 os << "Proc " << A.getMap()->getComm()->getRank() << ": "
835 << "Create Cimport" << std::endl;
836 std::cerr << os.str();
837 }
838 Cimport = rcp(new import_type(CDomainMap, Ccolmap));
839 }
840 if (!C.getRowMap()->isSameAs(*CRangeMap)) {
841 if (debug) {
842 std::ostringstream os;
843 os << "Proc " << A.getMap()->getComm()->getRank() << ": "
844 << "Create Cexport" << std::endl;
845 std::cerr << os.str();
846 }
847 Cexport = rcp(new export_type(C.getRowMap(), CRangeMap));
848 }
849
850 if (debug) {
851 std::ostringstream os;
852 os << "Proc " << A.getMap()->getComm()->getRank() << ": "
853 << "Call C->expertStaticFillComplete(...)" << std::endl;
854 std::cerr << os.str();
855 }
856 C.expertStaticFillComplete(CDomainMap, CRangeMap, Cimport, Cexport, params);
857 }
858 }
859}
860
861// This version of Add takes C as RCP&, so C may be null on input (in this case,
862// it is allocated and constructed in this function).
863template <class Scalar,
864 class LocalOrdinal,
865 class GlobalOrdinal,
866 class Node>
867void Add(
869 bool transposeA,
870 Scalar scalarA,
872 bool transposeB,
873 Scalar scalarB,
875 using std::endl;
876 using Teuchos::Array;
877 using Teuchos::ArrayRCP;
878 using Teuchos::ArrayView;
879 using Teuchos::RCP;
880 using Teuchos::rcp;
881 using Teuchos::rcp_dynamic_cast;
882 using Teuchos::rcpFromRef;
883 using Teuchos::tuple;
884 // typedef typename ArrayView<const Scalar>::size_type size_type;
885 typedef Teuchos::ScalarTraits<Scalar> STS;
887 // typedef Import<LocalOrdinal, GlobalOrdinal, Node> import_type;
888 // typedef RowGraph<LocalOrdinal, GlobalOrdinal, Node> row_graph_type;
889 // typedef CrsGraph<LocalOrdinal, GlobalOrdinal, Node> crs_graph_type;
892
893 std::string prefix = "TpetraExt::MatrixMatrix::Add(): ";
894
895 TEUCHOS_TEST_FOR_EXCEPTION(
896 !A.isFillComplete() || !B.isFillComplete(), std::invalid_argument,
897 prefix << "A and B must both be fill complete before calling this function.");
898
899 if (C.is_null()) {
900 TEUCHOS_TEST_FOR_EXCEPTION(!A.haveGlobalConstants(), std::logic_error,
901 prefix << "C is null (must be allocated), but A.haveGlobalConstants() is false. "
902 "Please report this bug to the Tpetra developers.");
903 TEUCHOS_TEST_FOR_EXCEPTION(!B.haveGlobalConstants(), std::logic_error,
904 prefix << "C is null (must be allocated), but B.haveGlobalConstants() is false. "
905 "Please report this bug to the Tpetra developers.");
906 }
907
908#ifdef HAVE_TPETRA_DEBUG
909 {
910 const bool domainMapsSame =
911 (!transposeA && !transposeB && !A.getDomainMap()->isSameAs(*(B.getDomainMap()))) ||
912 (!transposeA && transposeB && !A.getDomainMap()->isSameAs(*(B.getRangeMap()))) ||
913 (transposeA && !transposeB && !A.getRangeMap()->isSameAs(*(B.getDomainMap())));
914 TEUCHOS_TEST_FOR_EXCEPTION(domainMapsSame, std::invalid_argument,
915 prefix << "The domain Maps of Op(A) and Op(B) are not the same.");
916
917 const bool rangeMapsSame =
918 (!transposeA && !transposeB && !A.getRangeMap()->isSameAs(*(B.getRangeMap()))) ||
919 (!transposeA && transposeB && !A.getRangeMap()->isSameAs(*(B.getDomainMap()))) ||
920 (transposeA && !transposeB && !A.getDomainMap()->isSameAs(*(B.getRangeMap())));
921 TEUCHOS_TEST_FOR_EXCEPTION(rangeMapsSame, std::invalid_argument,
922 prefix << "The range Maps of Op(A) and Op(B) are not the same.");
923 }
924#endif // HAVE_TPETRA_DEBUG
925
926 using Teuchos::ParameterList;
927 RCP<ParameterList> transposeParams(new ParameterList);
928 transposeParams->set("sort", false);
929
930 // Form the explicit transpose of A if necessary.
931 RCP<const crs_matrix_type> Aprime;
932 if (transposeA) {
933 transposer_type theTransposer(rcpFromRef(A));
934 Aprime = theTransposer.createTranspose(transposeParams);
935 } else {
936 Aprime = rcpFromRef(A);
937 }
938
939#ifdef HAVE_TPETRA_DEBUG
940 TEUCHOS_TEST_FOR_EXCEPTION(Aprime.is_null(), std::logic_error,
941 prefix << "Failed to compute Op(A). Please report this bug to the Tpetra developers.");
942#endif // HAVE_TPETRA_DEBUG
943
944 // Form the explicit transpose of B if necessary.
945 RCP<const crs_matrix_type> Bprime;
946 if (transposeB) {
947 transposer_type theTransposer(rcpFromRef(B));
948 Bprime = theTransposer.createTranspose(transposeParams);
949 } else {
950 Bprime = rcpFromRef(B);
951 }
952
953#ifdef HAVE_TPETRA_DEBUG
954 TEUCHOS_TEST_FOR_EXCEPTION(Bprime.is_null(), std::logic_error,
955 prefix << "Failed to compute Op(B). Please report this bug to the Tpetra developers.");
956#endif // HAVE_TPETRA_DEBUG
957
958 bool CwasFillComplete = false;
959
960 // Allocate or zero the entries of the result matrix.
961 if (!C.is_null()) {
962 CwasFillComplete = C->isFillComplete();
963 if (CwasFillComplete)
964 C->resumeFill();
965 C->setAllToScalar(STS::zero());
966 } else {
967 // FIXME (mfh 08 May 2013) When I first looked at this method, I
968 // noticed that C was being given the row Map of Aprime (the
969 // possibly transposed version of A). Is this what we want?
970
971 // It is a precondition that Aprime and Bprime have the same domain and range maps.
972 // However, they may have different row maps. In this case, it's difficult to
973 // get a precise upper bound on the number of entries in each local row of C, so
974 // just use the looser upper bound based on the max number of entries in any row of Aprime and Bprime.
975 if (Aprime->getRowMap()->isSameAs(*Bprime->getRowMap())) {
976 LocalOrdinal numLocalRows = Aprime->getLocalNumRows();
977 Array<size_t> CmaxEntriesPerRow(numLocalRows);
978 for (LocalOrdinal i = 0; i < numLocalRows; i++) {
979 CmaxEntriesPerRow[i] = Aprime->getNumEntriesInLocalRow(i) + Bprime->getNumEntriesInLocalRow(i);
980 }
981 C = rcp(new crs_matrix_type(Aprime->getRowMap(), CmaxEntriesPerRow()));
982 } else {
983 // Note: above we checked that Aprime and Bprime have global constants, so it's safe to ask for max entries per row.
984 C = rcp(new crs_matrix_type(Aprime->getRowMap(), Aprime->getGlobalMaxNumRowEntries() + Bprime->getGlobalMaxNumRowEntries()));
985 }
986 }
987
988#ifdef HAVE_TPETRA_DEBUG
989 TEUCHOS_TEST_FOR_EXCEPTION(Aprime.is_null(), std::logic_error,
990 prefix << "At this point, Aprime is null. Please report this bug to the Tpetra developers.");
991 TEUCHOS_TEST_FOR_EXCEPTION(Bprime.is_null(), std::logic_error,
992 prefix << "At this point, Bprime is null. Please report this bug to the Tpetra developers.");
993 TEUCHOS_TEST_FOR_EXCEPTION(C.is_null(), std::logic_error,
994 prefix << "At this point, C is null. Please report this bug to the Tpetra developers.");
995#endif // HAVE_TPETRA_DEBUG
996
997 Array<RCP<const crs_matrix_type>> Mat =
998 tuple<RCP<const crs_matrix_type>>(Aprime, Bprime);
999 Array<Scalar> scalar = tuple<Scalar>(scalarA, scalarB);
1000
1001 // do a loop over each matrix to add: A reordering might be more efficient
1002 for (int k = 0; k < 2; ++k) {
1003 typename crs_matrix_type::nonconst_global_inds_host_view_type Indices;
1004 typename crs_matrix_type::nonconst_values_host_view_type Values;
1005
1006 // Loop over each locally owned row of the current matrix (either
1007 // Aprime or Bprime), and sum its entries into the corresponding
1008 // row of C. This works regardless of whether Aprime or Bprime
1009 // has the same row Map as C, because both sumIntoGlobalValues and
1010 // insertGlobalValues allow summing resp. inserting into nonowned
1011 // rows of C.
1012#ifdef HAVE_TPETRA_DEBUG
1013 TEUCHOS_TEST_FOR_EXCEPTION(Mat[k].is_null(), std::logic_error,
1014 prefix << "At this point, curRowMap is null. Please report this bug to the Tpetra developers.");
1015#endif // HAVE_TPETRA_DEBUG
1016 RCP<const map_type> curRowMap = Mat[k]->getRowMap();
1017#ifdef HAVE_TPETRA_DEBUG
1018 TEUCHOS_TEST_FOR_EXCEPTION(curRowMap.is_null(), std::logic_error,
1019 prefix << "At this point, curRowMap is null. Please report this bug to the Tpetra developers.");
1020#endif // HAVE_TPETRA_DEBUG
1021
1022 const size_t localNumRows = Mat[k]->getLocalNumRows();
1023 for (size_t i = 0; i < localNumRows; ++i) {
1024 const GlobalOrdinal globalRow = curRowMap->getGlobalElement(i);
1025 size_t numEntries = Mat[k]->getNumEntriesInGlobalRow(globalRow);
1026 if (numEntries > 0) {
1027 if (numEntries > Indices.extent(0)) {
1028 Kokkos::resize(Indices, numEntries);
1029 Kokkos::resize(Values, numEntries);
1030 }
1031 Mat[k]->getGlobalRowCopy(globalRow, Indices, Values, numEntries);
1032
1033 if (scalar[k] != STS::one()) {
1034 for (size_t j = 0; j < numEntries; ++j) {
1035 Values[j] *= scalar[k];
1036 }
1037 }
1038
1039 if (CwasFillComplete) {
1040 size_t result = C->sumIntoGlobalValues(globalRow, numEntries,
1041 reinterpret_cast<Scalar*>(Values.data()), Indices.data());
1042 TEUCHOS_TEST_FOR_EXCEPTION(result != numEntries, std::logic_error,
1043 prefix << "sumIntoGlobalValues failed to add entries from A or B into C.");
1044 } else {
1045 C->insertGlobalValues(globalRow, numEntries,
1046 reinterpret_cast<Scalar*>(Values.data()), Indices.data());
1047 }
1048 }
1049 }
1050 }
1051 if (CwasFillComplete) {
1052 C->fillComplete(C->getDomainMap(),
1053 C->getRangeMap());
1054 }
1055}
1056
1057// This version of Add takes C as const RCP&, so C must not be null on input. Otherwise, its behavior is identical
1058// to the above version where C is RCP&.
1059template <class Scalar,
1060 class LocalOrdinal,
1061 class GlobalOrdinal,
1062 class Node>
1063void Add(
1065 bool transposeA,
1066 Scalar scalarA,
1068 bool transposeB,
1069 Scalar scalarB,
1071 std::string prefix = "TpetraExt::MatrixMatrix::Add(): ";
1072
1073 TEUCHOS_TEST_FOR_EXCEPTION(C.is_null(), std::invalid_argument,
1074 prefix << "C must not be null");
1075
1076 Teuchos::RCP<CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>> C_ = C;
1077 Add(A, transposeA, scalarA, B, transposeB, scalarB, C_);
1078}
1079
1080} // End namespace MatrixMatrix
1081
1082namespace MMdetails {
1083
1084// nvcc 12.3 and 12.4 report getApplyHelper() as inaccessible even if
1085// CrsMatrix friends kokkos_kernels_mult_A_B_newmatrix. Use this non-template
1086// wrapper instead. Someone had the same problem in 2011 here:
1087// https://forums.developer.nvidia.com/t/24378.
1088struct CrsMatrixApplyHelperAccess {
1089 template <class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node>
1090 static auto get(const CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& matrix) {
1091 return matrix.getApplyHelper();
1092 }
1093};
1094
1095template <class Scalar,
1096 class LocalOrdinal,
1097 class GlobalOrdinal,
1098 class Node,
1099 class LocalOrdinalViewType>
1100void kokkos_kernels_mult_A_B_newmatrix(
1101 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
1102 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Bview,
1103 const LocalOrdinalViewType& Acol2Brow,
1104 const LocalOrdinalViewType& Acol2Irow,
1105 const LocalOrdinalViewType& Bcol2Ccol,
1106 const LocalOrdinalViewType& Icol2Ccol,
1107 CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& C,
1108 Teuchos::RCP<const Import<LocalOrdinal, GlobalOrdinal, Node>> Cimport,
1109 const std::string& label,
1110 const Teuchos::RCP<Teuchos::ParameterList>& params) {
1111 using backend_type = KokkosKernelsSPGEMMBackend<Node>;
1112 using Teuchos::RCP;
1113 using Teuchos::rcp;
1114
1116 using device_t = typename KCRS::device_type;
1117 using graph_t = typename KCRS::StaticCrsGraphType;
1118 using lno_view_t = typename graph_t::row_map_type::non_const_type;
1119 using int_view_t = Kokkos::View<int*, typename lno_view_t::array_layout, typename lno_view_t::memory_space, typename lno_view_t::memory_traits>;
1120 using lno_nnz_view_t = typename graph_t::entries_type::non_const_type;
1121 using scalar_view_t = typename KCRS::values_type::non_const_type;
1122 using KernelHandle = KokkosKernels::Experimental::KokkosKernelsHandle<
1123 typename lno_view_t::const_value_type, typename lno_nnz_view_t::const_value_type, typename scalar_view_t::const_value_type,
1124 typename device_t::execution_space, typename device_t::memory_space, typename device_t::memory_space>;
1125 using IntKernelHandle = KokkosKernels::Experimental::KokkosKernelsHandle<
1126 typename int_view_t::const_value_type, typename lno_nnz_view_t::const_value_type, typename scalar_view_t::const_value_type,
1127 typename device_t::execution_space, typename device_t::memory_space, typename device_t::memory_space>;
1128
1129 const bool debug = Tpetra::Details::Behavior::debug();
1130 const std::string wrapperLabel = "TpetraExt: MMM: Newmatrix " + backend_type::algorithm_label() + "Wrapper";
1131 RCP<Tpetra::Details::ProfilingRegion> MM =
1132 rcp(new Tpetra::Details::ProfilingRegion(wrapperLabel));
1133
1134 bool spgemm_options_customized = false;
1135 int team_work_size = 16;
1136 std::string myalg("SPGEMM_DEFAULT");
1137 if (!params.is_null()) {
1138 const std::string prefixedAlg = backend_type::parameter_prefix() + ": algorithm";
1139 const std::string prefixedTeam = backend_type::parameter_prefix() + ": team work size";
1140 if (params->isParameter(prefixedAlg))
1141 myalg = params->get(prefixedAlg, myalg);
1142 if (params->isParameter(prefixedTeam)) {
1143 team_work_size = params->get(prefixedTeam, team_work_size);
1144 spgemm_options_customized = true;
1145 }
1146 }
1147#if KOKKOSKERNELS_VERSION < 50299
1148 (void)spgemm_options_customized;
1149#endif
1150
1151 const KCRS Amat = Aview.origMatrix->getLocalMatrixDevice();
1152
1153 const std::string genericAlg = backend_type::algorithm_label() + " algorithm";
1154 if (!params.is_null() && params->isParameter(genericAlg))
1155 myalg = params->get(genericAlg, myalg);
1156 KokkosSparse::SPGEMMAlgorithm alg_enum = KokkosSparse::StringToSPGEMMAlgorithm(myalg);
1157
1158 KCRS Bmerged = Tpetra::MMdetails::merge_matrices(
1159 Aview, Bview, Acol2Brow, Acol2Irow, Bcol2Ccol, Icol2Ccol, C.getColMap()->getLocalNumElements());
1160
1161 KCRS Cmat;
1162
1163#if KOKKOSKERNELS_VERSION >= 50299
1164 // Determine whether the local spgemm inputs are known to be sorted.
1165 // For A, we can directly ask its CrsGraph if it's sorted.
1166 // For Bmerged, we don't have a CrsGraph. But if no nontrivial import+merge
1167 // took place (Bview.importMatrix.is_null()), and the unmerged B was sorted, then
1168 // Bmerged must also be sorted.
1169 const bool input_sorted =
1170 Aview.origMatrix->getCrsGraph()->isSorted() &&
1171 (Bview.importMatrix.is_null() && Bview.origMatrix->getCrsGraph()->isSorted());
1172 // We always want C to be sorted.
1173 const bool result_sorted = true;
1174
1175 // For backward compatibility, fall back to the reuse (handle-based) interface
1176 // if the user has customized algorithm parameters (e.g. team work size),
1177 // which the non-reuse interface cannot express. This is an unlikely case.
1178 if (!spgemm_options_customized) {
1179 const std::string coreLabel = "TpetraExt: MMM: Newmatrix " + backend_type::algorithm_label() + "Core";
1180 MM = Teuchos::null;
1181 MM = rcp(new Tpetra::Details::ProfilingRegion(coreLabel));
1182
1183 Tpetra::Details::IntRowPtrHelper<decltype(Bmerged)> irph(Bmerged.nnz(), Bmerged.graph.row_map);
1184 const bool useIntRowptrs =
1185 irph.shouldUseIntRowptrs() &&
1186 CrsMatrixApplyHelperAccess::get(*Aview.origMatrix)->shouldUseIntRowptrs();
1187
1188 if (useIntRowptrs) {
1189 // Run spgemm with int-typed rowptrs (helps TPLs), then copy the resulting
1190 // rowptrs back to the correct rowptr type.
1191 auto Aint = CrsMatrixApplyHelperAccess::get(*Aview.origMatrix)->getIntRowptrMatrix(Amat);
1192 auto Bint = irph.getIntRowptrMatrix(Bmerged);
1193 using int_matrix_t = decltype(Bint);
1194
1195 int_matrix_t Cint = KokkosSparse::spgemm<int_matrix_t>(
1196 alg_enum, Aint, false, Bint, false, input_sorted, result_sorted);
1197
1198 auto int_row_mapC = Cint.graph.row_map;
1199 lno_view_t row_mapC(Kokkos::ViewAllocateWithoutInitializing("non_const_lno_row"), int_row_mapC.extent(0));
1200 Kokkos::parallel_for(
1201 Kokkos::RangePolicy<typename device_t::execution_space>(0, int_row_mapC.extent(0)),
1202 KOKKOS_LAMBDA(const int i) { row_mapC(i) = int_row_mapC(i); });
1203 Cmat = KCRS("C", Cint.numRows(), Cint.numCols(), Cint.nnz(), Cint.values, row_mapC, Cint.graph.entries);
1204 } else {
1205 Cmat = KokkosSparse::spgemm<KCRS>(
1206 alg_enum, Amat, false, Bmerged, false, input_sorted, result_sorted);
1207 }
1208 } else
1209#endif
1210 {
1211 backend_type::pre_spgemm(*Aview.origMatrix);
1212 backend_type::pre_spgemm(Bmerged);
1213
1214 const std::string coreLabel = "TpetraExt: MMM: Newmatrix " + backend_type::algorithm_label() + "Core";
1215 MM = Teuchos::null;
1216 MM = rcp(new Tpetra::Details::ProfilingRegion(coreLabel));
1217
1218 Tpetra::Details::IntRowPtrHelper<decltype(Bmerged)> irph(Bmerged.nnz(), Bmerged.graph.row_map);
1219 const bool useIntRowptrs =
1220 irph.shouldUseIntRowptrs() &&
1221 CrsMatrixApplyHelperAccess::get(*Aview.origMatrix)->shouldUseIntRowptrs();
1222
1223 if (useIntRowptrs) {
1224 typename KernelHandle::nnz_lno_t AnumRows = Amat.numRows();
1225 typename KernelHandle::nnz_lno_t BnumRows = Bmerged.numRows();
1226 typename KernelHandle::nnz_lno_t BnumCols = Bmerged.numCols();
1227
1228 lno_view_t row_mapC(Kokkos::ViewAllocateWithoutInitializing("non_const_lno_row"), AnumRows + 1);
1229 lno_nnz_view_t entriesC;
1230 scalar_view_t valuesC;
1231
1232 IntKernelHandle kh;
1233#if KOKKOSKERNELS_VERSION >= 50299
1234 kh.create_spgemm_handle(alg_enum, input_sorted, result_sorted);
1235#else
1236 kh.create_spgemm_handle(alg_enum);
1237#endif
1238 kh.set_team_work_size(team_work_size);
1239
1240 int_view_t int_row_mapC(Kokkos::ViewAllocateWithoutInitializing("non_const_int_row"), AnumRows + 1);
1241
1242 auto Aint = CrsMatrixApplyHelperAccess::get(*Aview.origMatrix)->getIntRowptrMatrix(Amat);
1243 auto Bint = irph.getIntRowptrMatrix(Bmerged);
1244
1245 {
1246 Tpetra::Details::ProfilingRegion MM2("TpetraExt: MMM: Newmatrix KokkosKernels symbolic int");
1247 KokkosSparse::spgemm_symbolic(
1248 &kh, AnumRows, BnumRows, BnumCols, Aint.graph.row_map, Aint.graph.entries, false, Bint.graph.row_map, Bint.graph.entries, false, int_row_mapC);
1249 }
1250
1251 Tpetra::Details::ProfilingRegion MM2("TpetraExt: MMM: Newmatrix KokkosKernels numeric int");
1252 size_t c_nnz_size = kh.get_spgemm_handle()->get_c_nnz();
1253 if (c_nnz_size) {
1254 entriesC = lno_nnz_view_t(Kokkos::ViewAllocateWithoutInitializing("entriesC"), c_nnz_size);
1255 valuesC = scalar_view_t(Kokkos::ViewAllocateWithoutInitializing("valuesC"), c_nnz_size);
1256 }
1257 KokkosSparse::spgemm_numeric(
1258 &kh, AnumRows, BnumRows, BnumCols, Aint.graph.row_map, Aint.graph.entries, Aint.values, false,
1259 Bint.graph.row_map, Bint.graph.entries, Bint.values, false, int_row_mapC, entriesC, valuesC);
1260 Kokkos::parallel_for(
1261 Kokkos::RangePolicy<typename device_t::execution_space>(0, int_row_mapC.size()),
1262 KOKKOS_LAMBDA(const int i) { row_mapC(i) = int_row_mapC(i); });
1263 kh.destroy_spgemm_handle();
1264 Cmat = KCRS("C", AnumRows, BnumCols, c_nnz_size, valuesC, row_mapC, entriesC);
1265 } else {
1266 KernelHandle kh;
1267#if KOKKOSKERNELS_VERSION >= 50299
1268 kh.create_spgemm_handle(alg_enum, input_sorted, result_sorted);
1269#else
1270 kh.create_spgemm_handle(alg_enum);
1271#endif
1272 kh.set_team_work_size(team_work_size);
1273
1274 {
1275 Tpetra::Details::ProfilingRegion MM2("TpetraExt: MMM: Newmatrix KokkosKernels symbolic non-int");
1276 KokkosSparse::spgemm_symbolic(kh, Amat, false, Bmerged, false, Cmat);
1277 }
1278
1279 Tpetra::Details::ProfilingRegion MM2("TpetraExt: MMM: Newmatrix KokkosKernels numeric non-int");
1280 KokkosSparse::spgemm_numeric(kh, Amat, false, Bmerged, false, Cmat);
1281 kh.destroy_spgemm_handle();
1282 }
1283 }
1284
1285#if KOKKOSKERNELS_VERSION < 50299
1286 if (debug) {
1287 // In KokkosKernels < 5.3, we did not have the option to pass result_sorted=true to the spgemm handle,
1288 // but the spgemm_numeric implementation (native and all TPLs) did explicitly sort the output.
1289 // In debug mode, ensure that this worked as expected.
1290 TEUCHOS_TEST_FOR_EXCEPTION(!KokkosSparse::Impl::isCrsGraphSorted(Cmat.graph.row_map, Cmat.graph.entries),
1291 std::runtime_error,
1292 "KokkosKernels spgemm_numeric did not produce sorted output as expected!");
1293 }
1294#else
1295 (void)debug; // avoid -Wunused
1296#endif
1297
1298 C.setAllValues(Cmat);
1299
1300 const std::string esfcLabel = "TpetraExt: MMM: Newmatrix " + backend_type::algorithm_label() + "ESFC";
1301 MM = Teuchos::null;
1302 MM = rcp(new Tpetra::Details::ProfilingRegion(esfcLabel));
1303
1304 RCP<Teuchos::ParameterList> labelList = rcp(new Teuchos::ParameterList);
1305 labelList->set("Timer Label", label);
1306 if (!params.is_null())
1307 labelList->set("compute global constants", params->get("compute global constants", true));
1308 RCP<const Export<LocalOrdinal, GlobalOrdinal, Node>> dummyExport;
1309 C.expertStaticFillComplete(Bview.origMatrix->getDomainMap(), Aview.origMatrix->getRangeMap(), Cimport, dummyExport, labelList);
1310 TEUCHOS_TEST_FOR_EXCEPTION(!C.getCrsGraph()->isSorted(), std::runtime_error,
1311 "TpetraExt: kokkos_kernels_mult_A_B_newmatrix: result C does not have a sorted CrsGraph.");
1312}
1313
1314template <class Scalar,
1315 class LocalOrdinal,
1316 class GlobalOrdinal,
1317 class Node,
1318 class LocalOrdinalViewType>
1319void host_mult_A_B_reuse(
1320 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
1321 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Bview,
1322 const LocalOrdinalViewType& targetMapToOrigRow_dev,
1323 const LocalOrdinalViewType& targetMapToImportRow_dev,
1324 const LocalOrdinalViewType& Bcol2Ccol_dev,
1325 const LocalOrdinalViewType& Icol2Ccol_dev,
1326 CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& C,
1327 Teuchos::RCP<const Import<LocalOrdinal, GlobalOrdinal, Node>> Cimport,
1328 const std::string& label,
1329 const Teuchos::RCP<Teuchos::ParameterList>& params) {
1330 // FIXME: Right now, this is a cut-and-paste of the serial kernel
1331
1332 using Teuchos::RCP;
1333 using Teuchos::rcp;
1334
1335 // By default, if A*B results in an entry that is not supported by the graph of C, we throw.
1336 // This option allows to override this behavior and silently ignores such entries.
1337 bool throwOnInsert = true;
1338 if (!params.is_null() && params->isType<bool>("MM Throw For Non-Existent Entries"))
1339 throwOnInsert = params->get<bool>("MM Throw For Non-Existent Entries");
1340
1341 RCP<Tpetra::Details::ProfilingRegion> MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: MMM: Reuse SerialCore"));
1342
1343 // Lots and lots of typedefs
1344 typedef typename Tpetra::CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>::local_matrix_host_type KCRS;
1345 typedef typename KCRS::StaticCrsGraphType graph_t;
1346 typedef typename graph_t::row_map_type::const_type c_lno_view_t;
1347 typedef typename graph_t::entries_type::non_const_type lno_nnz_view_t;
1348 typedef typename KCRS::values_type::non_const_type scalar_view_t;
1349
1350 typedef Scalar SC;
1351 typedef LocalOrdinal LO;
1352 typedef GlobalOrdinal GO;
1353 typedef Node NO;
1354 typedef Map<LO, GO, NO> map_type;
1355 const size_t ST_INVALID = Teuchos::OrdinalTraits<LO>::invalid();
1356 const LO LO_INVALID = Teuchos::OrdinalTraits<LO>::invalid();
1357 const SC SC_ZERO = Teuchos::ScalarTraits<Scalar>::zero();
1358
1359 // Since this is being run on Cuda, we need to fence because the below code will use UVM
1360 // typename graph_t::execution_space().fence();
1361
1362 // KDDKDD UVM Without UVM, need to copy targetMap arrays to host.
1363 // KDDKDD UVM Ideally, this function would run on device and use
1364 // KDDKDD UVM KokkosKernels instead of this host implementation.
1365 auto targetMapToOrigRow =
1366 Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(),
1367 targetMapToOrigRow_dev);
1368 auto targetMapToImportRow =
1369 Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(),
1370 targetMapToImportRow_dev);
1371 auto Bcol2Ccol =
1372 Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(),
1373 Bcol2Ccol_dev);
1374 auto Icol2Ccol =
1375 Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(),
1376 Icol2Ccol_dev);
1377
1378 // Sizes
1379 RCP<const map_type> Ccolmap = C.getColMap();
1380 size_t m = Aview.origMatrix->getLocalNumRows();
1381 size_t n = Ccolmap->getLocalNumElements();
1382
1383 // Grab the Kokkos::SparseCrsMatrices & inner stuff
1384 const KCRS Amat = Aview.origMatrix->getLocalMatrixHost();
1385 const KCRS Bmat = Bview.origMatrix->getLocalMatrixHost();
1386 const KCRS Cmat = C.getLocalMatrixHost();
1387
1388 c_lno_view_t Arowptr = Amat.graph.row_map,
1389 Browptr = Bmat.graph.row_map,
1390 Crowptr = Cmat.graph.row_map;
1391 const lno_nnz_view_t Acolind = Amat.graph.entries,
1392 Bcolind = Bmat.graph.entries,
1393 Ccolind = Cmat.graph.entries;
1394 const scalar_view_t Avals = Amat.values, Bvals = Bmat.values;
1395 scalar_view_t Cvals = Cmat.values;
1396
1397 c_lno_view_t Irowptr;
1398 lno_nnz_view_t Icolind;
1399 scalar_view_t Ivals;
1400 if (!Bview.importMatrix.is_null()) {
1401 auto lclB = Bview.importMatrix->getLocalMatrixHost();
1402 Irowptr = lclB.graph.row_map;
1403 Icolind = lclB.graph.entries;
1404 Ivals = lclB.values;
1405 }
1406
1407 // Classic csr assembly (low memory edition)
1408 // mfh 27 Sep 2016: The c_status array is an implementation detail
1409 // of the local sparse matrix-matrix multiply routine.
1410
1411 // The status array will contain the index into colind where this entry was last deposited.
1412 // c_status[i] < CSR_ip - not in the row yet
1413 // c_status[i] >= CSR_ip - this is the entry where you can find the data
1414 // We start with this filled with INVALID's indicating that there are no entries yet.
1415 // Sadly, this complicates the code due to the fact that size_t's are unsigned.
1416 std::vector<size_t> c_status(n, ST_INVALID);
1417
1418 // For each row of A/C
1419 size_t CSR_ip = 0, OLD_ip = 0;
1420 for (size_t i = 0; i < m; i++) {
1421 // First fill the c_status array w/ locations where we're allowed to
1422 // generate nonzeros for this row
1423 OLD_ip = Crowptr[i];
1424 CSR_ip = Crowptr[i + 1];
1425 for (size_t k = OLD_ip; k < CSR_ip; k++) {
1426 c_status[Ccolind[k]] = k;
1427
1428 // Reset values in the row of C
1429 Cvals[k] = SC_ZERO;
1430 }
1431
1432 for (size_t k = Arowptr[i]; k < Arowptr[i + 1]; k++) {
1433 LO Aik = Acolind[k];
1434 const SC Aval = Avals[k];
1435 if (Aval == SC_ZERO)
1436 continue;
1437
1438 if (targetMapToOrigRow[Aik] != LO_INVALID) {
1439 // Local matrix
1440 size_t Bk = Teuchos::as<size_t>(targetMapToOrigRow[Aik]);
1441
1442 for (size_t j = Browptr[Bk]; j < Browptr[Bk + 1]; ++j) {
1443 LO Bkj = Bcolind[j];
1444 LO Cij = Bcol2Ccol[Bkj];
1445
1446 const bool badInsert = (Cij == LO_INVALID) || (c_status[Cij] < OLD_ip) || (c_status[Cij] >= CSR_ip);
1447 if (!badInsert)
1448 Cvals[c_status[Cij]] += Aval * Bvals[j];
1449 else if (throwOnInsert)
1450 TEUCHOS_TEST_FOR_EXCEPTION(badInsert,
1451 std::runtime_error, "Trying to insert a new entry (" << i << "," << Cij << ") into a static graph "
1452 << "(c_status = " << c_status[Cij] << " of [" << OLD_ip << "," << CSR_ip << "))");
1453 }
1454
1455 } else {
1456 // Remote matrix
1457 size_t Ik = Teuchos::as<size_t>(targetMapToImportRow[Aik]);
1458 for (size_t j = Irowptr[Ik]; j < Irowptr[Ik + 1]; ++j) {
1459 LO Ikj = Icolind[j];
1460 LO Cij = Icol2Ccol[Ikj];
1461
1462 const bool badInsert = (Cij == LO_INVALID) || (c_status[Cij] < OLD_ip) || (c_status[Cij] >= CSR_ip);
1463 if (!badInsert)
1464 Cvals[c_status[Cij]] += Aval * Ivals[j];
1465 else if (throwOnInsert)
1466 TEUCHOS_TEST_FOR_EXCEPTION(badInsert,
1467 std::runtime_error, "Trying to insert a new entry (" << i << "," << Cij << ") into a static graph "
1468 << "(c_status = " << c_status[Cij] << " of [" << OLD_ip << "," << CSR_ip << "))");
1469 }
1470 }
1471 }
1472 }
1473
1474 C.fillComplete(C.getDomainMap(), C.getRangeMap());
1475}
1476
1477template <class Scalar,
1478 class LocalOrdinal,
1479 class GlobalOrdinal,
1480 class Node,
1481 class LocalOrdinalViewType>
1482void kokkos_kernels_jacobi_A_B_newmatrix(typename Teuchos::ScalarTraits<Scalar>::magnitudeType omega,
1483 const Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Dinv,
1484 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
1485 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Bview,
1486 const LocalOrdinalViewType& Acol2Brow,
1487 const LocalOrdinalViewType& Acol2Irow,
1488 const LocalOrdinalViewType& Bcol2Ccol,
1489 const LocalOrdinalViewType& Icol2Ccol,
1490 CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& C,
1491 Teuchos::RCP<const Import<LocalOrdinal, GlobalOrdinal, Node>> Cimport,
1492 const std::string& label,
1493 const Teuchos::RCP<Teuchos::ParameterList>& params) {
1494 using backend_type = KokkosKernelsSPGEMMBackend<Node>;
1495 // Check if the diagonal entries exist in debug mode
1496 const bool debug = Tpetra::Details::Behavior::debug();
1497 if (debug) {
1498 auto rowMap = Aview.origMatrix->getRowMap();
1499 Tpetra::Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node> diags(rowMap);
1500 Aview.origMatrix->getLocalDiagCopy(diags);
1501 size_t diagLength = rowMap->getLocalNumElements();
1502 Teuchos::Array<Scalar> diagonal(diagLength);
1503 diags.get1dCopy(diagonal());
1504
1505 for (size_t i = 0; i < diagLength; ++i) {
1506 TEUCHOS_TEST_FOR_EXCEPTION(diagonal[i] == Teuchos::ScalarTraits<Scalar>::zero(),
1507 std::runtime_error,
1508 "Matrix A has a zero/missing diagonal: " << diagonal[i] << std::endl
1509 << "KokkosKernels Jacobi-fused SpGEMM requires nonzero diagonal entries in A" << std::endl);
1510 }
1511 }
1512
1513 using Teuchos::RCP;
1514 using Teuchos::rcp;
1515 RCP<Tpetra::Details::ProfilingRegion> MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: Jacobi: Newmatrix KokkosKernels"));
1516
1517 // Usings
1518 using device_t = typename Node::device_type;
1520 using graph_t = typename matrix_t::StaticCrsGraphType;
1521 using lno_view_t = typename graph_t::row_map_type::non_const_type;
1522 using int_view_t = Kokkos::View<int*,
1523 typename lno_view_t::array_layout,
1524 typename lno_view_t::memory_space,
1525 typename lno_view_t::memory_traits>;
1526 using lno_nnz_view_t = typename graph_t::entries_type::non_const_type;
1527 using scalar_view_t = typename matrix_t::values_type::non_const_type;
1528
1529 // KokkosKernels handle
1530 using handle_t = typename KokkosKernels::Experimental::KokkosKernelsHandle<
1531 typename lno_view_t::const_value_type, typename lno_nnz_view_t::const_value_type, typename scalar_view_t::const_value_type,
1532 typename device_t::execution_space, typename device_t::memory_space, typename device_t::memory_space>;
1533
1534 using int_handle_t = typename KokkosKernels::Experimental::KokkosKernelsHandle<
1535 typename int_view_t::const_value_type, typename lno_nnz_view_t::const_value_type, typename scalar_view_t::const_value_type,
1536 typename device_t::execution_space, typename device_t::memory_space, typename device_t::memory_space>;
1537
1538 // Merge the B and Bimport matrices
1539 const matrix_t Bmerged = Tpetra::MMdetails::merge_matrices(Aview, Bview, Acol2Brow, Acol2Irow, Bcol2Ccol, Icol2Ccol, C.getColMap()->getLocalNumElements());
1540
1541 // Get the properties and arrays of input matrices
1542 const matrix_t Amat = Aview.origMatrix->getLocalMatrixDevice();
1543 const matrix_t Bmat = Bview.origMatrix->getLocalMatrixDevice();
1544
1545 typename handle_t::nnz_lno_t AnumRows = Amat.numRows();
1546 typename handle_t::nnz_lno_t BnumRows = Bmerged.numRows();
1547 typename handle_t::nnz_lno_t BnumCols = Bmerged.numCols();
1548
1549 // Arrays of the output matrix
1550 lno_view_t row_mapC(Kokkos::ViewAllocateWithoutInitializing("row_mapC"), AnumRows + 1);
1551 lno_nnz_view_t entriesC;
1552 scalar_view_t valuesC;
1553
1554 // Options
1555 int team_work_size = 16;
1556 std::string myalg("SPGEMM_DEFAULT");
1557 if (!params.is_null()) {
1558 const std::string prefixedAlg = backend_type::parameter_prefix() + ": algorithm";
1559 const std::string prefixedTeam = backend_type::parameter_prefix() + ": team work size";
1560 if (params->isParameter(prefixedAlg))
1561 myalg = params->get(prefixedAlg, myalg);
1562 if (params->isParameter(prefixedTeam))
1563 team_work_size = params->get(prefixedTeam, team_work_size);
1564 }
1565
1566 // Get the algorithm mode
1567 const std::string genericAlg = backend_type::algorithm_label() + " algorithm";
1568 if (!params.is_null() && params->isParameter(genericAlg))
1569 myalg = params->get(genericAlg, myalg);
1570 KokkosSparse::SPGEMMAlgorithm alg_enum = KokkosSparse::StringToSPGEMMAlgorithm(myalg);
1571
1572 // Decide whether to use integer-typed row pointers for this spgemm, in order to allow a TPL to be used in the symbolic phase.
1573 // Note: until KokkosKernels 5.3, if spgemm_symbolic uses a TPL but A or B is unsorted, an exception is thrown in
1574 // debug mode. To avoid this exception, we only enable the use of a TPL here when both inputs are known to be sorted.
1575 Tpetra::Details::IntRowPtrHelper<decltype(Bmerged)> irph(Bmerged.nnz(), Bmerged.graph.row_map);
1576#if KOKKOSKERNELS_VERSION >= 50299
1577 const bool useIntRowptrs =
1578 irph.shouldUseIntRowptrs() &&
1579 CrsMatrixApplyHelperAccess::get(*Aview.origMatrix)->shouldUseIntRowptrs();
1580#else
1581 const bool useIntRowptrs =
1582 irph.shouldUseIntRowptrs() &&
1583 CrsMatrixApplyHelperAccess::get(*Aview.origMatrix)->shouldUseIntRowptrs() &&
1584 KokkosSparse::isCrsGraphSorted(Amat.graph.row_map, Amat.graph.entries) &&
1585 KokkosSparse::isCrsGraphSorted(Bmerged.graph.row_map, Bmerged.graph.entries);
1586#endif
1587
1588 const Scalar jacobiOmega = omega * Teuchos::ScalarTraits<Scalar>::one();
1589
1590#if KOKKOSKERNELS_VERSION >= 50299
1591 // Determine whether the local spgemm inputs are known to be sorted.
1592 // For A, we can directly ask its CrsGraph if it's sorted.
1593 // For Bmerged, we don't have a CrsGraph. But if no nontrivial import+merge
1594 // took place (Bview.importMatrix.is_null()), and the unmerged B was sorted, then
1595 // Bmerged must also be sorted.
1596 const bool input_sorted =
1597 Aview.origMatrix->getCrsGraph()->isSorted() &&
1598 (Bview.importMatrix.is_null() && Bview.origMatrix->getCrsGraph()->isSorted());
1599 // We always want C to be sorted.
1600 const bool result_sorted = true;
1601#endif
1602
1603 if (useIntRowptrs) {
1604 int_handle_t kh;
1605#if KOKKOSKERNELS_VERSION >= 50299
1606 kh.create_spgemm_handle(alg_enum, input_sorted, result_sorted);
1607#else
1608 kh.create_spgemm_handle(alg_enum);
1609#endif
1610 kh.set_team_work_size(team_work_size);
1611
1612 int_view_t int_row_mapC(Kokkos::ViewAllocateWithoutInitializing("int_row_mapC"), AnumRows + 1);
1613
1614 auto Aint = CrsMatrixApplyHelperAccess::get(*Aview.origMatrix)->getIntRowptrMatrix(Amat);
1615 auto Bint = irph.getIntRowptrMatrix(Bmerged);
1616
1617 {
1618 Tpetra::Details::ProfilingRegion MM2("TpetraExt: Jacobi: KokkosKernels symbolic int");
1619 KokkosSparse::spgemm_symbolic(&kh, AnumRows, BnumRows, BnumCols,
1620 Aint.graph.row_map, Aint.graph.entries, false,
1621 Bint.graph.row_map, Bint.graph.entries, false,
1622 int_row_mapC);
1623 }
1624
1625 size_t c_nnz_size = kh.get_spgemm_handle()->get_c_nnz();
1626 if (c_nnz_size) {
1627 entriesC = lno_nnz_view_t(Kokkos::ViewAllocateWithoutInitializing("entriesC"), c_nnz_size);
1628 valuesC = scalar_view_t(Kokkos::ViewAllocateWithoutInitializing("valuesC"), c_nnz_size);
1629 }
1630 Tpetra::Details::ProfilingRegion MM2("TpetraExt: Jacobi: KokkosKernels numeric int");
1631
1632 if (c_nnz_size) {
1633 // even though there is no TPL for this, we have to use the same handle that was used in the symbolic phase,
1634 // so need to have a special int-typed call for this as well.
1635 KokkosSparse::Experimental::spgemm_jacobi(&kh, AnumRows, BnumRows, BnumCols,
1636 Aint.graph.row_map, Aint.graph.entries, Amat.values, false,
1637 Bint.graph.row_map, Bint.graph.entries, Bint.values, false,
1638 int_row_mapC, entriesC, valuesC,
1639 jacobiOmega, Dinv.getLocalViewDevice(Access::ReadOnly));
1640 }
1641 // transfer the integer rowptrs back to the correct rowptr type
1642 Kokkos::parallel_for(
1643 Kokkos::RangePolicy<typename device_t::execution_space>(0, int_row_mapC.size()),
1644 KOKKOS_LAMBDA(int i) { row_mapC(i) = int_row_mapC(i); });
1645 kh.destroy_spgemm_handle();
1646 } else {
1647 handle_t kh;
1648#if KOKKOSKERNELS_VERSION >= 50299
1649 kh.create_spgemm_handle(alg_enum, input_sorted, result_sorted);
1650#else
1651 kh.create_spgemm_handle(alg_enum);
1652#endif
1653 kh.set_team_work_size(team_work_size);
1654
1655 {
1656 Tpetra::Details::ProfilingRegion MM2("TpetraExt: Jacobi: KokkosKernels symbolic non-int");
1657 KokkosSparse::spgemm_symbolic(&kh, AnumRows, BnumRows, BnumCols,
1658 Amat.graph.row_map, Amat.graph.entries, false,
1659 Bmerged.graph.row_map, Bmerged.graph.entries, false,
1660 row_mapC);
1661 }
1662
1663 size_t c_nnz_size = kh.get_spgemm_handle()->get_c_nnz();
1664 if (c_nnz_size) {
1665 entriesC = lno_nnz_view_t(Kokkos::ViewAllocateWithoutInitializing("entriesC"), c_nnz_size);
1666 valuesC = scalar_view_t(Kokkos::ViewAllocateWithoutInitializing("valuesC"), c_nnz_size);
1667 }
1668
1669 Tpetra::Details::ProfilingRegion MM2("TpetraExt: Jacobi: KokkosKernels numeric non-int");
1670 if (c_nnz_size) {
1671 KokkosSparse::Experimental::spgemm_jacobi(&kh, AnumRows, BnumRows, BnumCols,
1672 Amat.graph.row_map, Amat.graph.entries, Amat.values, false,
1673 Bmerged.graph.row_map, Bmerged.graph.entries, Bmerged.values, false,
1674 row_mapC, entriesC, valuesC,
1675 jacobiOmega, Dinv.getLocalViewDevice(Access::ReadOnly));
1676 }
1677 kh.destroy_spgemm_handle();
1678 }
1679
1680#if KOKKOSKERNELS_VERSION < 50299
1681 if (debug) {
1682 // In KokkosKernels < 5.3, we did not have the option to pass result_sorted=true to the spgemm handle,
1683 // but the spgemm_jacobi implementation did explicitly sort the output.
1684 // In debug mode, ensure that this worked as expected.
1685 TEUCHOS_TEST_FOR_EXCEPTION(!KokkosSparse::Impl::isCrsGraphSorted(row_mapC, entriesC),
1686 std::runtime_error,
1687 "KokkosKernels spgemm_jacobi did not produce sorted output!");
1688 }
1689#endif
1690
1691 C.setAllValues(row_mapC, entriesC, valuesC);
1692
1693 const std::string esfcLabel = "TpetraExt: Jacobi: Newmatrix " + backend_type::algorithm_label() + "ESFC";
1694 MM = Teuchos::null;
1695 MM = rcp(new Tpetra::Details::ProfilingRegion(esfcLabel));
1696
1697 // Final Fillcomplete
1698 Teuchos::RCP<Teuchos::ParameterList> labelList = rcp(new Teuchos::ParameterList);
1699 labelList->set("Timer Label", label);
1700 if (!params.is_null()) labelList->set("compute global constants", params->get("compute global constants", true));
1701 Teuchos::RCP<const Export<LocalOrdinal, GlobalOrdinal, Node>> dummyExport;
1702 C.expertStaticFillComplete(Bview.origMatrix->getDomainMap(), Aview.origMatrix->getRangeMap(), Cimport, dummyExport, labelList);
1703 TEUCHOS_TEST_FOR_EXCEPTION(!C.getCrsGraph()->isSorted(), std::runtime_error,
1704 "TpetraExt: kokkos_kernels_jacobi_A_B_newmatrix: result C does not have a sorted CrsGraph.");
1705}
1706
1707template <class Scalar,
1708 class LocalOrdinal,
1709 class GlobalOrdinal,
1710 class Node,
1711 class LocalOrdinalViewType>
1712void host_jacobi_A_B_reuse(typename Teuchos::ScalarTraits<Scalar>::magnitudeType omega,
1713 const Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Dinv,
1714 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
1715 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Bview,
1716 const LocalOrdinalViewType& targetMapToOrigRow_dev,
1717 const LocalOrdinalViewType& targetMapToImportRow_dev,
1718 const LocalOrdinalViewType& Bcol2Ccol_dev,
1719 const LocalOrdinalViewType& Icol2Ccol_dev,
1720 CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& C,
1721 Teuchos::RCP<const Import<LocalOrdinal, GlobalOrdinal, Node>> Cimport,
1722 const std::string& label,
1723 const Teuchos::RCP<Teuchos::ParameterList>& params) {
1724 // FIXME: Right now, this is a cut-and-paste of the serial kernel
1725 using Teuchos::RCP;
1726 using Teuchos::rcp;
1727 RCP<Tpetra::Details::ProfilingRegion> MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: Jacobi: Reuse CudaCore"));
1728
1729 // Lots and lots of typedefs
1730 typedef typename Tpetra::CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>::local_matrix_host_type KCRS;
1731 typedef typename KCRS::StaticCrsGraphType graph_t;
1732 typedef typename graph_t::row_map_type::const_type c_lno_view_t;
1733 typedef typename graph_t::entries_type::non_const_type lno_nnz_view_t;
1734 typedef typename KCRS::values_type::non_const_type scalar_view_t;
1735 typedef typename scalar_view_t::memory_space scalar_memory_space;
1736
1737 typedef Scalar SC;
1738 typedef LocalOrdinal LO;
1739 typedef GlobalOrdinal GO;
1740 typedef Node NO;
1741 typedef Map<LO, GO, NO> map_type;
1742 const size_t ST_INVALID = Teuchos::OrdinalTraits<LO>::invalid();
1743 const LO LO_INVALID = Teuchos::OrdinalTraits<LO>::invalid();
1744 const SC SC_ZERO = Teuchos::ScalarTraits<Scalar>::zero();
1745
1746 // Since this is being run on Cuda, we need to fence because the below host code will use UVM
1747 // KDDKDD typename graph_t::execution_space().fence();
1748
1749 // KDDKDD UVM Without UVM, need to copy targetMap arrays to host.
1750 // KDDKDD UVM Ideally, this function would run on device and use
1751 // KDDKDD UVM KokkosKernels instead of this host implementation.
1752 auto targetMapToOrigRow =
1753 Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(),
1754 targetMapToOrigRow_dev);
1755 auto targetMapToImportRow =
1756 Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(),
1757 targetMapToImportRow_dev);
1758 auto Bcol2Ccol =
1759 Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(),
1760 Bcol2Ccol_dev);
1761 auto Icol2Ccol =
1762 Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(),
1763 Icol2Ccol_dev);
1764
1765 // Sizes
1766 RCP<const map_type> Ccolmap = C.getColMap();
1767 size_t m = Aview.origMatrix->getLocalNumRows();
1768 size_t n = Ccolmap->getLocalNumElements();
1769
1770 // Grab the Kokkos::SparseCrsMatrices & inner stuff
1771 const KCRS Amat = Aview.origMatrix->getLocalMatrixHost();
1772 const KCRS Bmat = Bview.origMatrix->getLocalMatrixHost();
1773 const KCRS Cmat = C.getLocalMatrixHost();
1774
1775 c_lno_view_t Arowptr = Amat.graph.row_map, Browptr = Bmat.graph.row_map, Crowptr = Cmat.graph.row_map;
1776 const lno_nnz_view_t Acolind = Amat.graph.entries, Bcolind = Bmat.graph.entries, Ccolind = Cmat.graph.entries;
1777 const scalar_view_t Avals = Amat.values, Bvals = Bmat.values;
1778 scalar_view_t Cvals = Cmat.values;
1779
1780 c_lno_view_t Irowptr;
1781 lno_nnz_view_t Icolind;
1782 scalar_view_t Ivals;
1783 if (!Bview.importMatrix.is_null()) {
1784 auto lclB = Bview.importMatrix->getLocalMatrixHost();
1785 Irowptr = lclB.graph.row_map;
1786 Icolind = lclB.graph.entries;
1787 Ivals = lclB.values;
1788 }
1789
1790 // Jacobi-specific inner stuff
1791 auto Dvals =
1792 Dinv.template getLocalView<scalar_memory_space>(Access::ReadOnly);
1793
1794 // The status array will contain the index into colind where this entry was last deposited.
1795 // c_status[i] < CSR_ip - not in the row yet
1796 // c_status[i] >= CSR_ip - this is the entry where you can find the data
1797 // We start with this filled with INVALID's indicating that there are no entries yet.
1798 // Sadly, this complicates the code due to the fact that size_t's are unsigned.
1799 std::vector<size_t> c_status(n, ST_INVALID);
1800
1801 // For each row of A/C
1802 size_t CSR_ip = 0, OLD_ip = 0;
1803 for (size_t i = 0; i < m; i++) {
1804 // First fill the c_status array w/ locations where we're allowed to
1805 // generate nonzeros for this row
1806 OLD_ip = Crowptr[i];
1807 CSR_ip = Crowptr[i + 1];
1808 for (size_t k = OLD_ip; k < CSR_ip; k++) {
1809 c_status[Ccolind[k]] = k;
1810
1811 // Reset values in the row of C
1812 Cvals[k] = SC_ZERO;
1813 }
1814
1815 SC minusOmegaDval = -omega * Dvals(i, 0);
1816
1817 // Entries of B
1818 for (size_t j = Browptr[i]; j < Browptr[i + 1]; j++) {
1819 Scalar Bval = Bvals[j];
1820 if (Bval == SC_ZERO)
1821 continue;
1822 LO Bij = Bcolind[j];
1823 LO Cij = Bcol2Ccol[Bij];
1824
1825 TEUCHOS_TEST_FOR_EXCEPTION(c_status[Cij] < OLD_ip || c_status[Cij] >= CSR_ip,
1826 std::runtime_error, "Trying to insert a new entry into a static graph");
1827
1828 Cvals[c_status[Cij]] = Bvals[j];
1829 }
1830
1831 // Entries of -omega * Dinv * A * B
1832 for (size_t k = Arowptr[i]; k < Arowptr[i + 1]; k++) {
1833 LO Aik = Acolind[k];
1834 const SC Aval = Avals[k];
1835 if (Aval == SC_ZERO)
1836 continue;
1837
1838 if (targetMapToOrigRow[Aik] != LO_INVALID) {
1839 // Local matrix
1840 size_t Bk = Teuchos::as<size_t>(targetMapToOrigRow[Aik]);
1841
1842 for (size_t j = Browptr[Bk]; j < Browptr[Bk + 1]; ++j) {
1843 LO Bkj = Bcolind[j];
1844 LO Cij = Bcol2Ccol[Bkj];
1845
1846 TEUCHOS_TEST_FOR_EXCEPTION(c_status[Cij] < OLD_ip || c_status[Cij] >= CSR_ip,
1847 std::runtime_error, "Trying to insert a new entry into a static graph");
1848
1849 Cvals[c_status[Cij]] += minusOmegaDval * Aval * Bvals[j];
1850 }
1851
1852 } else {
1853 // Remote matrix
1854 size_t Ik = Teuchos::as<size_t>(targetMapToImportRow[Aik]);
1855 for (size_t j = Irowptr[Ik]; j < Irowptr[Ik + 1]; ++j) {
1856 LO Ikj = Icolind[j];
1857 LO Cij = Icol2Ccol[Ikj];
1858
1859 TEUCHOS_TEST_FOR_EXCEPTION(c_status[Cij] < OLD_ip || c_status[Cij] >= CSR_ip,
1860 std::runtime_error, "Trying to insert a new entry into a static graph");
1861
1862 Cvals[c_status[Cij]] += minusOmegaDval * Aval * Ivals[j];
1863 }
1864 }
1865 }
1866 }
1867
1868 MM = Teuchos::null;
1869 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: Jacobi: Reuse ESFC"));
1870
1871 C.fillComplete(C.getDomainMap(), C.getRangeMap());
1872}
1873
1874/*********************************************************************************************************/
1876// template <class TransferType>
1877// void printMultiplicationStatistics(Teuchos::RCP<TransferType > Transfer, const std::string &label) {
1878// if (Transfer.is_null())
1879// return;
1880//
1881// const Distributor & Distor = Transfer->getDistributor();
1882// Teuchos::RCP<const Teuchos::Comm<int> > Comm = Transfer->getSourceMap()->getComm();
1883//
1884// size_t rows_send = Transfer->getNumExportIDs();
1885// size_t rows_recv = Transfer->getNumRemoteIDs();
1886//
1887// size_t round1_send = Transfer->getNumExportIDs() * sizeof(size_t);
1888// size_t round1_recv = Transfer->getNumRemoteIDs() * sizeof(size_t);
1889// size_t num_send_neighbors = Distor.getNumSends();
1890// size_t num_recv_neighbors = Distor.getNumReceives();
1891// size_t round2_send, round2_recv;
1892// Distor.getLastDoStatistics(round2_send,round2_recv);
1893//
1894// int myPID = Comm->getRank();
1895// int NumProcs = Comm->getSize();
1896//
1897// // Processor by processor statistics
1898// // printf("[%d] %s Statistics: neigh[s/r]=%d/%d rows[s/r]=%d/%d r1bytes[s/r]=%d/%d r2bytes[s/r]=%d/%d\n",
1899// // myPID, label.c_str(),num_send_neighbors,num_recv_neighbors,rows_send,rows_recv,round1_send,round1_recv,round2_send,round2_recv);
1900//
1901// // Global statistics
1902// size_t lstats[8] = {num_send_neighbors,num_recv_neighbors,rows_send,rows_recv,round1_send,round1_recv,round2_send,round2_recv};
1903// size_t gstats_min[8], gstats_max[8];
1904//
1905// double lstats_avg[8], gstats_avg[8];
1906// for(int i=0; i<8; i++)
1907// lstats_avg[i] = ((double)lstats[i])/NumProcs;
1908//
1909// Teuchos::reduceAll(*Comm(),Teuchos::REDUCE_MIN,8,lstats,gstats_min);
1910// Teuchos::reduceAll(*Comm(),Teuchos::REDUCE_MAX,8,lstats,gstats_max);
1911// Teuchos::reduceAll(*Comm(),Teuchos::REDUCE_SUM,8,lstats_avg,gstats_avg);
1912//
1913// if(!myPID) {
1914// printf("%s Send Statistics[min/avg/max]: neigh=%d/%4.1f/%d rows=%d/%4.1f/%d round1=%d/%4.1f/%d round2=%d/%4.1f/%d\n", label.c_str(),
1915// (int)gstats_min[0],gstats_avg[0],(int)gstats_max[0], (int)gstats_min[2],gstats_avg[2],(int)gstats_max[2],
1916// (int)gstats_min[4],gstats_avg[4],(int)gstats_max[4], (int)gstats_min[6],gstats_avg[6],(int)gstats_max[6]);
1917// printf("%s Recv Statistics[min/avg/max]: neigh=%d/%4.1f/%d rows=%d/%4.1f/%d round1=%d/%4.1f/%d round2=%d/%4.1f/%d\n", label.c_str(),
1918// (int)gstats_min[1],gstats_avg[1],(int)gstats_max[1], (int)gstats_min[3],gstats_avg[3],(int)gstats_max[3],
1919// (int)gstats_min[5],gstats_avg[5],(int)gstats_max[5], (int)gstats_min[7],gstats_avg[7],(int)gstats_max[7]);
1920// }
1921// }
1922
1923// Kernel method for computing the local portion of C = A*B
1924template <class Scalar,
1925 class LocalOrdinal,
1926 class GlobalOrdinal,
1927 class Node>
1928void mult_AT_B_newmatrix(
1929 const CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& A,
1930 const CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& B,
1931 CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& C,
1932 const std::string& label,
1933 const Teuchos::RCP<Teuchos::ParameterList>& params) {
1934 using Teuchos::RCP;
1935 using Teuchos::rcp;
1936 typedef Scalar SC;
1937 typedef LocalOrdinal LO;
1938 typedef GlobalOrdinal GO;
1939 typedef Node NO;
1940 typedef CrsMatrixStruct<SC, LO, GO, NO> crs_matrix_struct_type;
1941 typedef RowMatrixTransposer<SC, LO, GO, NO> transposer_type;
1942
1943 RCP<Tpetra::Details::ProfilingRegion> MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: MMM-T: Transpose"));
1944
1945 /*************************************************************/
1946 /* 1) Local Transpose of A */
1947 /*************************************************************/
1948 transposer_type transposer(rcpFromRef(A), label + std::string("XP: "));
1949
1950 using Teuchos::ParameterList;
1951 RCP<ParameterList> transposeParams(new ParameterList);
1952 transposeParams->set("sort", true); // Kokkos Kernels spgemm requires inputs to be sorted
1953 if (!params.is_null()) {
1954 transposeParams->set("compute global constants",
1955 params->get("compute global constants: temporaries",
1956 false));
1957 }
1958 RCP<Tpetra::CrsMatrix<SC, LO, GO, NO>> Atrans =
1959 transposer.createTransposeLocal(transposeParams);
1960
1961 /*************************************************************/
1962 /* 2/3) Call mult_A_B_newmatrix w/ fillComplete */
1963 /*************************************************************/
1964 MM = Teuchos::null;
1965 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: MMM-T: I&X"));
1966
1967 // Get views, asserting that no import is required to speed up computation
1968 crs_matrix_struct_type Aview;
1969 crs_matrix_struct_type Bview;
1970 RCP<const Import<LO, GO, NO>> dummyImporter;
1971
1972 // NOTE: the I&X routine sticks an importer on the paramlist as output, so we have to use a unique guy here
1973 RCP<Teuchos::ParameterList> importParams = Teuchos::rcp(new ParameterList);
1974 importParams->set("compute global constants", false);
1975 if (!params.is_null()) {
1976 importParams->setParameters(*params);
1977 if (params->isParameter("compute global constants: temporaries"))
1978 importParams->set("compute global constants",
1979 params->get<bool>("compute global constants: temporaries"));
1980 }
1981 MMdetails::import_and_extract_views(*Atrans, Atrans->getRowMap(),
1982 Aview, dummyImporter, true,
1983 label, importParams);
1984
1985 if (B.getRowMap()->isSameAs(*Atrans->getColMap())) {
1986 MMdetails::import_and_extract_views(B, B.getRowMap(), Bview, dummyImporter, true, label, importParams);
1987 } else {
1988 MMdetails::import_and_extract_views(B, Atrans->getColMap(), Bview, dummyImporter, false, label, importParams);
1989 }
1990
1991 MM = Teuchos::null;
1992 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: MMM-T: AB-core"));
1993
1994 RCP<Tpetra::CrsMatrix<SC, LO, GO, NO>> Ctemp;
1995
1996 // If Atrans has no Exporter, we can use C instead of having to create a temp matrix
1997 bool needs_final_export = !Atrans->getGraph()->getExporter().is_null();
1998 if (needs_final_export) {
1999 Ctemp = rcp(new Tpetra::CrsMatrix<SC, LO, GO, NO>(Atrans->getRowMap(), 0));
2000 } else {
2001 Ctemp = rcp(&C, false);
2002 }
2003
2004 RCP<Teuchos::ParameterList> multParams = Teuchos::rcp(new ParameterList);
2005 if (!params.is_null()) {
2006 multParams->setParameters(*params);
2007 }
2008 multParams->set("compute global constants", !needs_final_export);
2009 mult_A_B_newmatrix(Aview, Bview, *Ctemp, label, multParams);
2010
2011 /*************************************************************/
2012 /* 4) exportAndFillComplete matrix */
2013 /*************************************************************/
2014 MM = Teuchos::null;
2015 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: MMM-T: exportAndFillComplete"));
2016
2017 RCP<Tpetra::CrsMatrix<SC, LO, GO, NO>> Crcp(&C, false);
2018
2019 if (needs_final_export) {
2020 ParameterList labelList;
2021 labelList.set("Timer Label", label);
2022 if (!params.is_null()) {
2023 labelList.setParameters(*params);
2024 }
2025 ParameterList& labelList_subList = labelList.sublist("matrixmatrix: kernel params", false);
2026 labelList_subList.set("isMatrixMatrix_TransferAndFillComplete", true,
2027 "This parameter should be set to true only for MatrixMatrix operations: the optimization in Epetra that was ported to Tpetra does _not_ take into account the possibility that for any given source PID, a particular GID may not exist on the target PID: i.e. a transfer operation. A fix for this general case is in development.");
2028
2029 Ctemp->exportAndFillComplete(Crcp,
2030 *Ctemp->getGraph()->getExporter(),
2031 B.getDomainMap(),
2032 A.getDomainMap(),
2033 rcp(&labelList, false));
2034 }
2035#ifdef HAVE_TPETRA_MMM_STATISTICS
2036 printMultiplicationStatistics(Ctemp->getGraph()->getExporter(), label + std::string(" AT_B MMM"));
2037#endif
2038}
2039
2040/*********************************************************************************************************/
2041// Kernel method for computing the local portion of C = A*B
2042template <class Scalar,
2043 class LocalOrdinal,
2044 class GlobalOrdinal,
2045 class Node>
2046void mult_A_B(
2047 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
2048 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Bview,
2049 CrsWrapper<Scalar, LocalOrdinal, GlobalOrdinal, Node>& C,
2050 const std::string& /* label */,
2051 const Teuchos::RCP<Teuchos::ParameterList>& params) {
2052 using Teuchos::Array;
2053 using Teuchos::ArrayRCP;
2054 using Teuchos::ArrayView;
2055 using Teuchos::null;
2056 using Teuchos::OrdinalTraits;
2057
2058 bool skipExplicitZero = true;
2059 if (params && params->isParameter("MM Skip Explicit Zeros")) {
2060 skipExplicitZero = params->get<bool>("MM Skip Explicit Zeros");
2061 }
2062
2063 typedef Teuchos::ScalarTraits<Scalar> STS;
2064 // TEUCHOS_FUNC_TIME_MONITOR_DIFF("mult_A_B", mult_A_B);
2065 LocalOrdinal C_firstCol = Bview.colMap->getMinLocalIndex();
2066 LocalOrdinal C_lastCol = Bview.colMap->getMaxLocalIndex();
2067
2068 LocalOrdinal C_firstCol_import = OrdinalTraits<LocalOrdinal>::zero();
2069 LocalOrdinal C_lastCol_import = OrdinalTraits<LocalOrdinal>::invalid();
2070
2071 ArrayView<const GlobalOrdinal> bcols = Bview.colMap->getLocalElementList();
2072 ArrayView<const GlobalOrdinal> bcols_import = null;
2073 if (Bview.importColMap != null) {
2074 C_firstCol_import = Bview.importColMap->getMinLocalIndex();
2075 C_lastCol_import = Bview.importColMap->getMaxLocalIndex();
2076
2077 bcols_import = Bview.importColMap->getLocalElementList();
2078 }
2079
2080 size_t C_numCols = C_lastCol - C_firstCol +
2081 OrdinalTraits<LocalOrdinal>::one();
2082 size_t C_numCols_import = C_lastCol_import - C_firstCol_import +
2083 OrdinalTraits<LocalOrdinal>::one();
2084
2085 if (C_numCols_import > C_numCols)
2086 C_numCols = C_numCols_import;
2087
2088 Array<Scalar> dwork = Array<Scalar>(C_numCols);
2089 Array<GlobalOrdinal> iwork = Array<GlobalOrdinal>(C_numCols);
2090 Array<size_t> iwork2 = Array<size_t>(C_numCols);
2091
2092 Array<Scalar> C_row_i = dwork;
2093 Array<GlobalOrdinal> C_cols = iwork;
2094 Array<size_t> c_index = iwork2;
2095 Array<GlobalOrdinal> combined_index = Array<GlobalOrdinal>(2 * C_numCols);
2096 Array<Scalar> combined_values = Array<Scalar>(2 * C_numCols);
2097
2098 size_t C_row_i_length, j, k, last_index;
2099
2100 // Run through all the hash table lookups once and for all
2101 LocalOrdinal LO_INVALID = OrdinalTraits<LocalOrdinal>::invalid();
2102 Array<LocalOrdinal> Acol2Brow(Aview.colMap->getLocalNumElements(), LO_INVALID);
2103 Array<LocalOrdinal> Acol2Irow(Aview.colMap->getLocalNumElements(), LO_INVALID);
2104 if (Aview.colMap->isSameAs(*Bview.origMatrix->getRowMap())) {
2105 // Maps are the same: Use local IDs as the hash
2106 for (LocalOrdinal i = Aview.colMap->getMinLocalIndex(); i <=
2107 Aview.colMap->getMaxLocalIndex();
2108 i++)
2109 Acol2Brow[i] = i;
2110 } else {
2111 // Maps are not the same: Use the map's hash
2112 for (LocalOrdinal i = Aview.colMap->getMinLocalIndex(); i <=
2113 Aview.colMap->getMaxLocalIndex();
2114 i++) {
2115 GlobalOrdinal GID = Aview.colMap->getGlobalElement(i);
2116 LocalOrdinal BLID = Bview.origMatrix->getRowMap()->getLocalElement(GID);
2117 if (BLID != LO_INVALID)
2118 Acol2Brow[i] = BLID;
2119 else
2120 Acol2Irow[i] = Bview.importMatrix->getRowMap()->getLocalElement(GID);
2121 }
2122 }
2123
2124 // To form C = A*B we're going to execute this expression:
2125 //
2126 // C(i,j) = sum_k( A(i,k)*B(k,j) )
2127 //
2128 // Our goal, of course, is to navigate the data in A and B once, without
2129 // performing searches for column-indices, etc.
2130 auto Arowptr = Aview.origMatrix->getLocalRowPtrsHost();
2131 auto Acolind = Aview.origMatrix->getLocalIndicesHost();
2132 auto Avals = Aview.origMatrix->getLocalValuesHost(Tpetra::Access::ReadOnly);
2133 auto Browptr = Bview.origMatrix->getLocalRowPtrsHost();
2134 auto Bcolind = Bview.origMatrix->getLocalIndicesHost();
2135 auto Bvals = Bview.origMatrix->getLocalValuesHost(Tpetra::Access::ReadOnly);
2136 decltype(Browptr) Irowptr;
2137 decltype(Bcolind) Icolind;
2138 decltype(Bvals) Ivals;
2139 if (!Bview.importMatrix.is_null()) {
2140 Irowptr = Bview.importMatrix->getLocalRowPtrsHost();
2141 Icolind = Bview.importMatrix->getLocalIndicesHost();
2142 Ivals = Bview.importMatrix->getLocalValuesHost(Tpetra::Access::ReadOnly);
2143 }
2144
2145 bool C_filled = C.isFillComplete();
2146
2147 for (size_t i = 0; i < C_numCols; i++)
2148 c_index[i] = OrdinalTraits<size_t>::invalid();
2149
2150 // Loop over the rows of A.
2151 size_t Arows = Aview.rowMap->getLocalNumElements();
2152 for (size_t i = 0; i < Arows; ++i) {
2153 // Only navigate the local portion of Aview... which is, thankfully, all of
2154 // A since this routine doesn't do transpose modes
2155 GlobalOrdinal global_row = Aview.rowMap->getGlobalElement(i);
2156
2157 // Loop across the i-th row of A and for each corresponding row in B, loop
2158 // across columns and accumulate product A(i,k)*B(k,j) into our partial sum
2159 // quantities C_row_i. In other words, as we stride across B(k,:) we're
2160 // calculating updates for row i of the result matrix C.
2161 C_row_i_length = OrdinalTraits<size_t>::zero();
2162
2163 for (k = Arowptr[i]; k < Arowptr[i + 1]; ++k) {
2164 LocalOrdinal Ak = Acol2Brow[Acolind[k]];
2165 const Scalar Aval = Avals[k];
2166 if (Aval == STS::zero() && skipExplicitZero)
2167 continue;
2168
2169 if (Ak == LO_INVALID)
2170 continue;
2171
2172 for (j = Browptr[Ak]; j < Browptr[Ak + 1]; ++j) {
2173 LocalOrdinal col = Bcolind[j];
2174 // assert(col >= 0 && col < C_numCols);
2175
2176 if (c_index[col] == OrdinalTraits<size_t>::invalid()) {
2177 // assert(C_row_i_length >= 0 && C_row_i_length < C_numCols);
2178 // This has to be a += so insertGlobalValue goes out
2179 C_row_i[C_row_i_length] = Aval * Bvals[j];
2180 C_cols[C_row_i_length] = col;
2181 c_index[col] = C_row_i_length;
2182 C_row_i_length++;
2183
2184 } else {
2185 // static cast from impl_scalar_type to Scalar needed for complex
2186 C_row_i[c_index[col]] += Aval * static_cast<Scalar>(Bvals[j]);
2187 }
2188 }
2189 }
2190
2191 for (size_t ii = 0; ii < C_row_i_length; ii++) {
2192 c_index[C_cols[ii]] = OrdinalTraits<size_t>::invalid();
2193 C_cols[ii] = bcols[C_cols[ii]];
2194 combined_index[ii] = C_cols[ii];
2195 combined_values[ii] = C_row_i[ii];
2196 }
2197 last_index = C_row_i_length;
2198
2199 //
2200 // Now put the C_row_i values into C.
2201 //
2202 // We might have to revamp this later.
2203 C_row_i_length = OrdinalTraits<size_t>::zero();
2204
2205 for (k = Arowptr[i]; k < Arowptr[i + 1]; ++k) {
2206 LocalOrdinal Ak = Acol2Brow[Acolind[k]];
2207 const Scalar Aval = Avals[k];
2208 if (Aval == STS::zero() && skipExplicitZero)
2209 continue;
2210
2211 if (Ak != LO_INVALID) continue;
2212
2213 Ak = Acol2Irow[Acolind[k]];
2214 for (j = Irowptr[Ak]; j < Irowptr[Ak + 1]; ++j) {
2215 LocalOrdinal col = Icolind[j];
2216 // assert(col >= 0 && col < C_numCols);
2217
2218 if (c_index[col] == OrdinalTraits<size_t>::invalid()) {
2219 // assert(C_row_i_length >= 0 && C_row_i_length < C_numCols);
2220 // This has to be a += so insertGlobalValue goes out
2221 C_row_i[C_row_i_length] = Aval * Ivals[j];
2222 C_cols[C_row_i_length] = col;
2223 c_index[col] = C_row_i_length;
2224 C_row_i_length++;
2225
2226 } else {
2227 // This has to be a += so insertGlobalValue goes out
2228 // static cast from impl_scalar_type to Scalar needed for complex
2229 C_row_i[c_index[col]] += Aval * static_cast<Scalar>(Ivals[j]);
2230 }
2231 }
2232 }
2233
2234 for (size_t ii = 0; ii < C_row_i_length; ii++) {
2235 c_index[C_cols[ii]] = OrdinalTraits<size_t>::invalid();
2236 C_cols[ii] = bcols_import[C_cols[ii]];
2237 combined_index[last_index] = C_cols[ii];
2238 combined_values[last_index] = C_row_i[ii];
2239 last_index++;
2240 }
2241
2242 // Now put the C_row_i values into C.
2243 // We might have to revamp this later.
2244 C_filled ? C.sumIntoGlobalValues(
2245 global_row,
2246 combined_index.view(OrdinalTraits<size_t>::zero(), last_index),
2247 combined_values.view(OrdinalTraits<size_t>::zero(), last_index))
2248 : C.insertGlobalValues(
2249 global_row,
2250 combined_index.view(OrdinalTraits<size_t>::zero(), last_index),
2251 combined_values.view(OrdinalTraits<size_t>::zero(), last_index));
2252 }
2253}
2254
2255/*********************************************************************************************************/
2256template <class Scalar,
2257 class LocalOrdinal,
2258 class GlobalOrdinal,
2259 class Node>
2260void setMaxNumEntriesPerRow(CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Mview) {
2261 typedef typename Teuchos::Array<Teuchos::ArrayView<const LocalOrdinal>>::size_type local_length_size;
2262 Mview.maxNumRowEntries = Teuchos::OrdinalTraits<local_length_size>::zero();
2263
2264 if (Mview.indices.size() > Teuchos::OrdinalTraits<local_length_size>::zero()) {
2265 Mview.maxNumRowEntries = Mview.indices[0].size();
2266
2267 for (local_length_size i = 1; i < Mview.indices.size(); ++i)
2268 if (Mview.indices[i].size() > Mview.maxNumRowEntries)
2269 Mview.maxNumRowEntries = Mview.indices[i].size();
2270 }
2271}
2272
2273/*********************************************************************************************************/
2274template <class CrsMatrixType>
2275size_t C_estimate_nnz(CrsMatrixType& A, CrsMatrixType& B) {
2276 // Follows the NZ estimate in ML's ml_matmatmult.c
2277 size_t Aest = 100, Best = 100;
2278 if (A.getLocalNumEntries() >= A.getLocalNumRows())
2279 Aest = (A.getLocalNumRows() > 0) ? A.getLocalNumEntries() / A.getLocalNumRows() : 100;
2280 if (B.getLocalNumEntries() >= B.getLocalNumRows())
2281 Best = (B.getLocalNumRows() > 0) ? B.getLocalNumEntries() / B.getLocalNumRows() : 100;
2282
2283 size_t nnzperrow = (size_t)(sqrt((double)Aest) + sqrt((double)Best) - 1);
2284 nnzperrow *= nnzperrow;
2285
2286 return (size_t)(A.getLocalNumRows() * nnzperrow * 0.75 + 100);
2287}
2288
2289/*********************************************************************************************************/
2290// Kernel method for computing the local portion of C = A*B for CrsMatrix
2291//
2292// mfh 27 Sep 2016: Currently, mult_AT_B_newmatrix() also calls this
2293// function, so this is probably the function we want to
2294// thread-parallelize.
2295template <class Scalar,
2296 class LocalOrdinal,
2297 class GlobalOrdinal,
2298 class Node>
2299void mult_A_B_newmatrix(
2300 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
2301 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Bview,
2302 CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& C,
2303 const std::string& label,
2304 const Teuchos::RCP<Teuchos::ParameterList>& params) {
2305 using Teuchos::Array;
2306 using Teuchos::ArrayRCP;
2307 using Teuchos::ArrayView;
2308 using Teuchos::RCP;
2309 using Teuchos::rcp;
2310
2311 // Tpetra typedefs
2312 typedef LocalOrdinal LO;
2313 typedef GlobalOrdinal GO;
2314 typedef Node NO;
2315 typedef Import<LO, GO, NO> import_type;
2316 typedef Map<LO, GO, NO> map_type;
2317
2318 // Kokkos typedefs
2319 typedef typename map_type::local_map_type local_map_type;
2321 typedef typename KCRS::StaticCrsGraphType graph_t;
2322 typedef typename graph_t::row_map_type::non_const_type lno_view_t;
2323 typedef typename NO::execution_space execution_space;
2324 typedef Kokkos::RangePolicy<execution_space, size_t> range_type;
2325 typedef Kokkos::View<LO*, typename lno_view_t::array_layout, typename lno_view_t::device_type> lo_view_t;
2326
2327 Tpetra::Details::ProfilingRegion MM("TpetraExt: MMM: M5 Cmap");
2328
2329 LO LO_INVALID = Teuchos::OrdinalTraits<LO>::invalid();
2330
2331 // Build the final importer / column map, hash table lookups for C
2332 RCP<const import_type> Cimport;
2333 RCP<const map_type> Ccolmap;
2334 RCP<const import_type> Bimport = Bview.origMatrix->getGraph()->getImporter();
2335 RCP<const import_type> Iimport = Bview.importMatrix.is_null() ? Teuchos::null : Bview.importMatrix->getGraph()->getImporter();
2336 local_map_type Acolmap_local = Aview.colMap->getLocalMap();
2337 local_map_type Browmap_local = Bview.origMatrix->getRowMap()->getLocalMap();
2338 local_map_type Irowmap_local;
2339 if (!Bview.importMatrix.is_null()) Irowmap_local = Bview.importMatrix->getRowMap()->getLocalMap();
2340 local_map_type Bcolmap_local = Bview.origMatrix->getColMap()->getLocalMap();
2341 local_map_type Icolmap_local;
2342 if (!Bview.importMatrix.is_null()) Icolmap_local = Bview.importMatrix->getColMap()->getLocalMap();
2343
2344 // mfh 27 Sep 2016: Bcol2Ccol is a table that maps from local column
2345 // indices of B, to local column indices of C. (B and C have the
2346 // same number of columns.) The kernel uses this, instead of
2347 // copying the entire input matrix B and converting its column
2348 // indices to those of C.
2349 lo_view_t Bcol2Ccol(Kokkos::ViewAllocateWithoutInitializing("Bcol2Ccol"), Bview.colMap->getLocalNumElements()), Icol2Ccol;
2350
2351 if (Bview.importMatrix.is_null()) {
2352 // mfh 27 Sep 2016: B has no "remotes," so B and C have the same column Map.
2353 Cimport = Bimport;
2354 Ccolmap = Bview.colMap;
2355 const LO colMapSize = static_cast<LO>(Bview.colMap->getLocalNumElements());
2356 // Bcol2Ccol is trivial
2357 Kokkos::parallel_for(
2358 "Tpetra::mult_A_B_newmatrix::Bcol2Ccol_fill",
2359 Kokkos::RangePolicy<execution_space, LO>(0, colMapSize),
2360 KOKKOS_LAMBDA(const LO i) {
2361 Bcol2Ccol(i) = i;
2362 });
2363 } else {
2364 // mfh 27 Sep 2016: B has "remotes," so we need to build the
2365 // column Map of C, as well as C's Import object (from its domain
2366 // Map to its column Map). C's column Map is the union of the
2367 // column Maps of (the local part of) B, and the "remote" part of
2368 // B. Ditto for the Import. We have optimized this "setUnion"
2369 // operation on Import objects and Maps.
2370
2371 // Choose the right variant of setUnion
2372 if (!Bimport.is_null() && !Iimport.is_null()) {
2373 Cimport = Bimport->setUnion(*Iimport, params);
2374 } else if (!Bimport.is_null() && Iimport.is_null()) {
2375 Cimport = Bimport->setUnion(params);
2376 } else if (Bimport.is_null() && !Iimport.is_null()) {
2377 Cimport = Iimport->setUnion(params);
2378 } else {
2379 throw std::runtime_error("TpetraExt::MMM status of matrix importers is nonsensical");
2380 }
2381 Ccolmap = Cimport->getTargetMap();
2382
2383 // FIXME (mfh 27 Sep 2016) This error check requires an all-reduce
2384 // in general. We should get rid of it in order to reduce
2385 // communication costs of sparse matrix-matrix multiply.
2386 TEUCHOS_TEST_FOR_EXCEPTION(!Cimport->getSourceMap()->isSameAs(*Bview.origMatrix->getDomainMap()),
2387 std::runtime_error, "Tpetra::MMM: Import setUnion messed with the DomainMap in an unfortunate way");
2388
2389 // NOTE: This is not efficient and should be folded into setUnion
2390 //
2391 // mfh 27 Sep 2016: What the above comment means, is that the
2392 // setUnion operation on Import objects could also compute these
2393 // local index - to - local index look-up tables.
2394 Kokkos::resize(Icol2Ccol, Bview.importMatrix->getColMap()->getLocalNumElements());
2395 local_map_type Ccolmap_local = Ccolmap->getLocalMap();
2396 Kokkos::parallel_for(
2397 "Tpetra::mult_A_B_newmatrix::Bcol2Ccol_getGlobalElement", range_type(0, Bview.origMatrix->getColMap()->getLocalNumElements()), KOKKOS_LAMBDA(const LO i) {
2398 Bcol2Ccol(i) = Ccolmap_local.getLocalElement(Bcolmap_local.getGlobalElement(i));
2399 });
2400 Kokkos::parallel_for(
2401 "Tpetra::mult_A_B_newmatrix::Icol2Ccol_getGlobalElement", range_type(0, Bview.importMatrix->getColMap()->getLocalNumElements()), KOKKOS_LAMBDA(const LO i) {
2402 Icol2Ccol(i) = Ccolmap_local.getLocalElement(Icolmap_local.getGlobalElement(i));
2403 });
2404 }
2405
2406 // Replace the column map
2407 //
2408 // mfh 27 Sep 2016: We do this because C was originally created
2409 // without a column Map. Now we have its column Map.
2410 C.replaceColMap(Ccolmap);
2411
2412 // mfh 27 Sep 2016: Construct tables that map from local column
2413 // indices of A, to local row indices of either B_local (the locally
2414 // owned part of B), or B_remote (the "imported" remote part of B).
2415 //
2416 // For column index Aik in row i of A, if the corresponding row of B
2417 // exists in the local part of B ("orig") (which I'll call B_local),
2418 // then targetMapToOrigRow[Aik] is the local index of that row of B.
2419 // Otherwise, targetMapToOrigRow[Aik] is "invalid" (a flag value).
2420 //
2421 // For column index Aik in row i of A, if the corresponding row of B
2422 // exists in the remote part of B ("Import") (which I'll call
2423 // B_remote), then targetMapToImportRow[Aik] is the local index of
2424 // that row of B. Otherwise, targetMapToOrigRow[Aik] is "invalid"
2425 // (a flag value).
2426
2427 // Run through all the hash table lookups once and for all
2428 lo_view_t targetMapToOrigRow(Kokkos::ViewAllocateWithoutInitializing("targetMapToOrigRow"), Aview.colMap->getLocalNumElements());
2429 lo_view_t targetMapToImportRow(Kokkos::ViewAllocateWithoutInitializing("targetMapToImportRow"), Aview.colMap->getLocalNumElements());
2430
2431 Kokkos::parallel_for(
2432 "Tpetra::mult_A_B_newmatrix::construct_tables", range_type(Aview.colMap->getMinLocalIndex(), Aview.colMap->getMaxLocalIndex() + 1), KOKKOS_LAMBDA(const LO i) {
2433 GO aidx = Acolmap_local.getGlobalElement(i);
2434 LO B_LID = Browmap_local.getLocalElement(aidx);
2435 if (B_LID != LO_INVALID) {
2436 targetMapToOrigRow(i) = B_LID;
2437 targetMapToImportRow(i) = LO_INVALID;
2438 } else {
2439 LO I_LID = Irowmap_local.getLocalElement(aidx);
2440 targetMapToOrigRow(i) = LO_INVALID;
2441 targetMapToImportRow(i) = I_LID;
2442 }
2443 });
2444
2445 // Call the actual kernel. We'll rely on partial template specialization to call the correct one ---
2446 // Either the straight-up Tpetra code (SerialNode) or the KokkosKernels one (other NGP node types)
2447 KernelWrappers<Scalar, LocalOrdinal, GlobalOrdinal, Node, lo_view_t>::mult_A_B_newmatrix_kernel_wrapper(Aview, Bview, targetMapToOrigRow, targetMapToImportRow, Bcol2Ccol, Icol2Ccol, C, Cimport, label, params);
2448}
2449
2450/*********************************************************************************************************/
2451// Kernel method for computing the local portion of C = A*B for BlockCrsMatrix
2452template <class Scalar,
2453 class LocalOrdinal,
2454 class GlobalOrdinal,
2455 class Node>
2456void mult_A_B_newmatrix(BlockCrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
2457 BlockCrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Bview,
2458 Teuchos::RCP<BlockCrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>>& C) {
2459 using Teuchos::Array;
2460 using Teuchos::ArrayRCP;
2461 using Teuchos::ArrayView;
2462 using Teuchos::null;
2463 using Teuchos::RCP;
2464 using Teuchos::rcp;
2465
2466 // Tpetra typedefs
2467 typedef LocalOrdinal LO;
2468 typedef GlobalOrdinal GO;
2469 typedef Node NO;
2470 typedef Import<LO, GO, NO> import_type;
2471 typedef Map<LO, GO, NO> map_type;
2472 typedef BlockCrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node> block_crs_matrix_type;
2473 typedef typename block_crs_matrix_type::crs_graph_type graph_t;
2474
2475 // Kokkos typedefs
2476 typedef typename map_type::local_map_type local_map_type;
2477 typedef typename block_crs_matrix_type::local_matrix_device_type KBSR;
2478 typedef typename KBSR::device_type device_t;
2479 typedef typename KBSR::StaticCrsGraphType static_graph_t;
2480 typedef typename static_graph_t::row_map_type::non_const_type lno_view_t;
2481 typedef typename static_graph_t::entries_type::non_const_type lno_nnz_view_t;
2482 typedef typename KBSR::values_type::non_const_type scalar_view_t;
2483 typedef typename NO::execution_space execution_space;
2484 typedef Kokkos::RangePolicy<execution_space, size_t> range_type;
2485 typedef Kokkos::View<LO*, typename lno_view_t::array_layout, typename lno_view_t::device_type> lo_view_t;
2486
2487 LO LO_INVALID = Teuchos::OrdinalTraits<LO>::invalid();
2488
2489 // Build the final importer / column map, hash table lookups for C
2490 RCP<const import_type> Cimport;
2491 RCP<const map_type> Ccolmap;
2492 RCP<const import_type> Bimport = Bview.origMatrix->getGraph()->getImporter();
2493 RCP<const import_type> Iimport = Bview.importMatrix.is_null() ? Teuchos::null : Bview.importMatrix->getGraph()->getImporter();
2494 local_map_type Acolmap_local = Aview.colMap->getLocalMap();
2495 local_map_type Browmap_local = Bview.origMatrix->getRowMap()->getLocalMap();
2496 local_map_type Irowmap_local;
2497 if (!Bview.importMatrix.is_null()) Irowmap_local = Bview.importMatrix->getRowMap()->getLocalMap();
2498 local_map_type Bcolmap_local = Bview.origMatrix->getColMap()->getLocalMap();
2499 local_map_type Icolmap_local;
2500 if (!Bview.importMatrix.is_null()) Icolmap_local = Bview.importMatrix->getColMap()->getLocalMap();
2501
2502 // Bcol2Ccol is a table that maps from local column
2503 // indices of B, to local column indices of C. (B and C have the
2504 // same number of columns.) The kernel uses this, instead of
2505 // copying the entire input matrix B and converting its column
2506 // indices to those of C.
2507 lo_view_t Bcol2Ccol(Kokkos::ViewAllocateWithoutInitializing("Bcol2Ccol"), Bview.colMap->getLocalNumElements()), Icol2Ccol;
2508
2509 if (Bview.importMatrix.is_null()) {
2510 // mfh 27 Sep 2016: B has no "remotes," so B and C have the same column Map.
2511 Cimport = Bimport;
2512 Ccolmap = Bview.colMap;
2513 const LO colMapSize = static_cast<LO>(Bview.colMap->getLocalNumElements());
2514 // Bcol2Ccol is trivial
2515 Kokkos::parallel_for(
2516 "Tpetra::mult_A_B_newmatrix::Bcol2Ccol_fill",
2517 Kokkos::RangePolicy<execution_space, LO>(0, colMapSize),
2518 KOKKOS_LAMBDA(const LO i) {
2519 Bcol2Ccol(i) = i;
2520 });
2521 } else {
2522 // B has "remotes," so we need to build the
2523 // column Map of C, as well as C's Import object (from its domain
2524 // Map to its column Map). C's column Map is the union of the
2525 // column Maps of (the local part of) B, and the "remote" part of
2526 // B. Ditto for the Import. We have optimized this "setUnion"
2527 // operation on Import objects and Maps.
2528
2529 // Choose the right variant of setUnion
2530 if (!Bimport.is_null() && !Iimport.is_null()) {
2531 Cimport = Bimport->setUnion(*Iimport);
2532 } else if (!Bimport.is_null() && Iimport.is_null()) {
2533 Cimport = Bimport->setUnion();
2534 } else if (Bimport.is_null() && !Iimport.is_null()) {
2535 Cimport = Iimport->setUnion();
2536 } else {
2537 throw std::runtime_error("TpetraExt::MMM status of matrix importers is nonsensical");
2538 }
2539 Ccolmap = Cimport->getTargetMap();
2540
2541 // NOTE: This is not efficient and should be folded into setUnion
2542 //
2543 // What the above comment means, is that the
2544 // setUnion operation on Import objects could also compute these
2545 // local index - to - local index look-up tables.
2546 Kokkos::resize(Icol2Ccol, Bview.importMatrix->getColMap()->getLocalNumElements());
2547 local_map_type Ccolmap_local = Ccolmap->getLocalMap();
2548 Kokkos::parallel_for(
2549 "Tpetra::mult_A_B_newmatrix::Bcol2Ccol_getGlobalElement",
2550 range_type(0, Bview.origMatrix->getColMap()->getLocalNumElements()),
2551 KOKKOS_LAMBDA(const LO i) {
2552 Bcol2Ccol(i) = Ccolmap_local.getLocalElement(Bcolmap_local.getGlobalElement(i));
2553 });
2554 Kokkos::parallel_for(
2555 "Tpetra::mult_A_B_newmatrix::Icol2Ccol_getGlobalElement",
2556 range_type(0, Bview.importMatrix->getColMap()->getLocalNumElements()),
2557 KOKKOS_LAMBDA(const LO i) {
2558 Icol2Ccol(i) = Ccolmap_local.getLocalElement(Icolmap_local.getGlobalElement(i));
2559 });
2560 }
2561
2562 // Construct tables that map from local column
2563 // indices of A, to local row indices of either B_local (the locally
2564 // owned part of B), or B_remote (the "imported" remote part of B).
2565 //
2566 // For column index Aik in row i of A, if the corresponding row of B
2567 // exists in the local part of B ("orig") (which I'll call B_local),
2568 // then targetMapToOrigRow[Aik] is the local index of that row of B.
2569 // Otherwise, targetMapToOrigRow[Aik] is "invalid" (a flag value).
2570 //
2571 // For column index Aik in row i of A, if the corresponding row of B
2572 // exists in the remote part of B ("Import") (which I'll call
2573 // B_remote), then targetMapToImportRow[Aik] is the local index of
2574 // that row of B. Otherwise, targetMapToOrigRow[Aik] is "invalid"
2575 // (a flag value).
2576
2577 // Run through all the hash table lookups once and for all
2578 lo_view_t targetMapToOrigRow(Kokkos::ViewAllocateWithoutInitializing("targetMapToOrigRow"),
2579 Aview.colMap->getLocalNumElements());
2580 lo_view_t targetMapToImportRow(Kokkos::ViewAllocateWithoutInitializing("targetMapToImportRow"),
2581 Aview.colMap->getLocalNumElements());
2582
2583 Kokkos::parallel_for(
2584 "Tpetra::mult_A_B_newmatrix::construct_tables",
2585 range_type(Aview.colMap->getMinLocalIndex(), Aview.colMap->getMaxLocalIndex() + 1),
2586 KOKKOS_LAMBDA(const LO i) {
2587 GO aidx = Acolmap_local.getGlobalElement(i);
2588 LO B_LID = Browmap_local.getLocalElement(aidx);
2589 if (B_LID != LO_INVALID) {
2590 targetMapToOrigRow(i) = B_LID;
2591 targetMapToImportRow(i) = LO_INVALID;
2592 } else {
2593 LO I_LID = Irowmap_local.getLocalElement(aidx);
2594 targetMapToOrigRow(i) = LO_INVALID;
2595 targetMapToImportRow(i) = I_LID;
2596 }
2597 });
2598
2599 // Create the KernelHandle
2600 using KernelHandle =
2601 KokkosKernels::Experimental::KokkosKernelsHandle<typename lno_view_t::const_value_type,
2602 typename lno_nnz_view_t::const_value_type,
2603 typename scalar_view_t::const_value_type,
2604 typename device_t::execution_space,
2605 typename device_t::memory_space,
2606 typename device_t::memory_space>;
2607 int team_work_size = 16; // Defaults to 16 as per Deveci 12/7/16 - csiefer
2608 std::string myalg("SPGEMM_DEFAULT");
2609 KokkosSparse::SPGEMMAlgorithm alg_enum = KokkosSparse::StringToSPGEMMAlgorithm(myalg);
2610
2611 KernelHandle kh;
2612 kh.create_spgemm_handle(alg_enum);
2613 kh.set_team_work_size(team_work_size);
2614
2615 // Get KokkosSparse::BsrMatrix for A and Bmerged (B and BImport)
2616 const KBSR Amat = Aview.origMatrix->getLocalMatrixDevice();
2617 const KBSR Bmerged = Tpetra::MMdetails::merge_matrices(Aview, Bview,
2618 targetMapToOrigRow, targetMapToImportRow,
2619 Bcol2Ccol, Icol2Ccol,
2620 Ccolmap.getConst()->getLocalNumElements());
2621
2622 RCP<graph_t> graphC;
2623 typename KBSR::values_type values;
2624 {
2625 // Call KokkosSparse routines to calculate Amat*Bmerged on device.
2626 // NOTE: Need to scope guard this since the BlockCrs constructor will need to copy the host graph
2627 KBSR Cmat;
2628 KokkosSparse::block_spgemm_symbolic(kh, Amat, false, Bmerged, false, Cmat);
2629 KokkosSparse::block_spgemm_numeric(kh, Amat, false, Bmerged, false, Cmat);
2630 kh.destroy_spgemm_handle();
2631
2632 // Build Tpetra::BlockCrsMatrix from KokkosSparse::BsrMatrix
2633 graphC = rcp(new graph_t(Cmat.graph, Aview.origMatrix->getRowMap(), Ccolmap.getConst()));
2634 values = Cmat.values;
2635 }
2636 C = rcp(new block_crs_matrix_type(*graphC, values, Aview.blocksize));
2637}
2638
2639/*********************************************************************************************************/
2640// AB NewMatrix Kernel wrappers (Default non-threaded version for CrsMatrix)
2641template <class Scalar,
2642 class LocalOrdinal,
2643 class GlobalOrdinal,
2644 class Node,
2645 class LocalOrdinalViewType>
2646void KernelWrappers<Scalar, LocalOrdinal, GlobalOrdinal, Node, LocalOrdinalViewType>::mult_A_B_newmatrix_kernel_wrapper(CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
2647 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Bview,
2648 const LocalOrdinalViewType& targetMapToOrigRow,
2649 const LocalOrdinalViewType& targetMapToImportRow,
2650 const LocalOrdinalViewType& Bcol2Ccol,
2651 const LocalOrdinalViewType& Icol2Ccol,
2652 CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& C,
2653 Teuchos::RCP<const Import<LocalOrdinal, GlobalOrdinal, Node>> Cimport,
2654 const std::string& label,
2655 const Teuchos::RCP<Teuchos::ParameterList>& params) {
2656 using Teuchos::Array;
2657 using Teuchos::ArrayRCP;
2658 using Teuchos::ArrayView;
2659 using Teuchos::RCP;
2660 using Teuchos::rcp;
2661
2662 Tpetra::Details::ProfilingRegion MM("TpetraExt: MMM: Newmatrix SerialCore");
2663
2664 // Lots and lots of typedefs
2665 typedef typename Tpetra::CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>::local_matrix_host_type KCRS;
2666 typedef typename KCRS::StaticCrsGraphType graph_t;
2667 typedef typename graph_t::row_map_type::const_type c_lno_view_t;
2668 typedef typename graph_t::row_map_type::non_const_type lno_view_t;
2669 typedef typename graph_t::entries_type::non_const_type lno_nnz_view_t;
2670 typedef typename KCRS::values_type::non_const_type scalar_view_t;
2671
2672 typedef Scalar SC;
2673 typedef LocalOrdinal LO;
2674 typedef GlobalOrdinal GO;
2675 typedef Node NO;
2676 typedef Map<LO, GO, NO> map_type;
2677 const size_t ST_INVALID = Teuchos::OrdinalTraits<LO>::invalid();
2678 const LO LO_INVALID = Teuchos::OrdinalTraits<LO>::invalid();
2679 const SC SC_ZERO = Teuchos::ScalarTraits<Scalar>::zero();
2680
2681 bool skipExplicitZero = true;
2682 if (params && params->isParameter("MM Skip Explicit Zeros")) {
2683 skipExplicitZero = params->get<bool>("MM Skip Explicit Zeros");
2684 }
2685
2686 // Sizes
2687 RCP<const map_type> Ccolmap = C.getColMap();
2688 size_t m = Aview.origMatrix->getLocalNumRows();
2689 size_t n = Ccolmap->getLocalNumElements();
2690 size_t b_max_nnz_per_row = Bview.origMatrix->getLocalMaxNumRowEntries();
2691
2692 // Grab the Kokkos::SparseCrsMatrices & inner stuff
2693 const KCRS Amat = Aview.origMatrix->getLocalMatrixHost();
2694 const KCRS Bmat = Bview.origMatrix->getLocalMatrixHost();
2695
2696 c_lno_view_t Arowptr = Amat.graph.row_map, Browptr = Bmat.graph.row_map;
2697 const lno_nnz_view_t Acolind = Amat.graph.entries, Bcolind = Bmat.graph.entries;
2698 const scalar_view_t Avals = Amat.values, Bvals = Bmat.values;
2699
2700 c_lno_view_t Irowptr;
2701 lno_nnz_view_t Icolind;
2702 scalar_view_t Ivals;
2703 if (!Bview.importMatrix.is_null()) {
2704 auto lclB = Bview.importMatrix->getLocalMatrixHost();
2705 Irowptr = lclB.graph.row_map;
2706 Icolind = lclB.graph.entries;
2707 Ivals = lclB.values;
2708 b_max_nnz_per_row = std::max(b_max_nnz_per_row, Bview.importMatrix->getLocalMaxNumRowEntries());
2709 }
2710
2711 // Classic csr assembly (low memory edition)
2712 //
2713 // mfh 27 Sep 2016: C_estimate_nnz does not promise an upper bound.
2714 // The method loops over rows of A, and may resize after processing
2715 // each row. Chris Siefert says that this reflects experience in
2716 // ML; for the non-threaded case, ML found it faster to spend less
2717 // effort on estimation and risk an occasional reallocation.
2718 size_t CSR_alloc = std::max(C_estimate_nnz(*Aview.origMatrix, *Bview.origMatrix), n);
2719 lno_view_t Crowptr(Kokkos::ViewAllocateWithoutInitializing("Crowptr"), m + 1);
2720 lno_nnz_view_t Ccolind(Kokkos::ViewAllocateWithoutInitializing("Ccolind"), CSR_alloc);
2721 scalar_view_t Cvals(Kokkos::ViewAllocateWithoutInitializing("Cvals"), CSR_alloc);
2722
2723 // mfh 27 Sep 2016: The c_status array is an implementation detail
2724 // of the local sparse matrix-matrix multiply routine.
2725
2726 // The status array will contain the index into colind where this entry was last deposited.
2727 // c_status[i] < CSR_ip - not in the row yet
2728 // c_status[i] >= CSR_ip - this is the entry where you can find the data
2729 // We start with this filled with INVALID's indicating that there are no entries yet.
2730 // Sadly, this complicates the code due to the fact that size_t's are unsigned.
2731 size_t INVALID = Teuchos::OrdinalTraits<size_t>::invalid();
2732 std::vector<size_t> c_status(n, ST_INVALID);
2733
2734 // mfh 27 Sep 2016: Here is the local sparse matrix-matrix multiply
2735 // routine. The routine computes C := A * (B_local + B_remote).
2736 //
2737 // For column index Aik in row i of A, targetMapToOrigRow[Aik] tells
2738 // you whether the corresponding row of B belongs to B_local
2739 // ("orig") or B_remote ("Import").
2740
2741 // For each row of A/C
2742 size_t CSR_ip = 0, OLD_ip = 0;
2743 for (size_t i = 0; i < m; i++) {
2744 // mfh 27 Sep 2016: m is the number of rows in the input matrix A
2745 // on the calling process.
2746 Crowptr[i] = CSR_ip;
2747
2748 // mfh 27 Sep 2016: For each entry of A in the current row of A
2749 for (size_t k = Arowptr[i]; k < Arowptr[i + 1]; k++) {
2750 LO Aik = Acolind[k]; // local column index of current entry of A
2751 const SC Aval = Avals[k]; // value of current entry of A
2752 if (Aval == SC_ZERO && skipExplicitZero)
2753 continue; // skip explicitly stored zero values in A
2754
2755 if (targetMapToOrigRow[Aik] != LO_INVALID) {
2756 // mfh 27 Sep 2016: If the entry of targetMapToOrigRow
2757 // corresponding to the current entry of A is populated, then
2758 // the corresponding row of B is in B_local (i.e., it lives on
2759 // the calling process).
2760
2761 // Local matrix
2762 size_t Bk = static_cast<size_t>(targetMapToOrigRow[Aik]);
2763
2764 // mfh 27 Sep 2016: Go through all entries in that row of B_local.
2765 for (size_t j = Browptr[Bk]; j < Browptr[Bk + 1]; ++j) {
2766 LO Bkj = Bcolind[j];
2767 LO Cij = Bcol2Ccol[Bkj];
2768
2769 if (c_status[Cij] == INVALID || c_status[Cij] < OLD_ip) {
2770 // New entry
2771 c_status[Cij] = CSR_ip;
2772 Ccolind[CSR_ip] = Cij;
2773 Cvals[CSR_ip] = Aval * Bvals[j];
2774 CSR_ip++;
2775
2776 } else {
2777 Cvals[c_status[Cij]] += Aval * Bvals[j];
2778 }
2779 }
2780
2781 } else {
2782 // mfh 27 Sep 2016: If the entry of targetMapToOrigRow
2783 // corresponding to the current entry of A NOT populated (has
2784 // a flag "invalid" value), then the corresponding row of B is
2785 // in B_local (i.e., it lives on the calling process).
2786
2787 // Remote matrix
2788 size_t Ik = static_cast<size_t>(targetMapToImportRow[Aik]);
2789 for (size_t j = Irowptr[Ik]; j < Irowptr[Ik + 1]; ++j) {
2790 LO Ikj = Icolind[j];
2791 LO Cij = Icol2Ccol[Ikj];
2792
2793 if (c_status[Cij] == INVALID || c_status[Cij] < OLD_ip) {
2794 // New entry
2795 c_status[Cij] = CSR_ip;
2796 Ccolind[CSR_ip] = Cij;
2797 Cvals[CSR_ip] = Aval * Ivals[j];
2798 CSR_ip++;
2799 } else {
2800 Cvals[c_status[Cij]] += Aval * Ivals[j];
2801 }
2802 }
2803 }
2804 }
2805
2806 // Resize for next pass if needed
2807 if (i + 1 < m && CSR_ip + std::min(n, (Arowptr[i + 2] - Arowptr[i + 1]) * b_max_nnz_per_row) > CSR_alloc) {
2808 CSR_alloc *= 2;
2809 Kokkos::resize(Ccolind, CSR_alloc);
2810 Kokkos::resize(Cvals, CSR_alloc);
2811 }
2812 OLD_ip = CSR_ip;
2813 }
2814
2815 Crowptr[m] = CSR_ip;
2816
2817 // Downward resize
2818 Kokkos::resize(Ccolind, CSR_ip);
2819 Kokkos::resize(Cvals, CSR_ip);
2820
2821 {
2822 Tpetra::Details::ProfilingRegion MM3("TpetraExt: MMM: Newmatrix Final Sort");
2823
2824 // Final sort & set of CRS arrays
2825 if (params.is_null() || params->get("sort entries", true)) {
2826 // Tpetra's serial SpGEMM results in almost sorted matrices. Use shell sort.
2827 Import_Util::sortCrsEntries(Crowptr, Ccolind, Cvals);
2828 }
2829 C.setAllValues(Crowptr, Ccolind, Cvals);
2830 }
2831
2832 Tpetra::Details::ProfilingRegion MM4("TpetraExt: MMM: Newmatrix ESCC");
2833 {
2834 // Final FillComplete
2835 //
2836 // mfh 27 Sep 2016: So-called "expert static fill complete" bypasses
2837 // Import (from domain Map to column Map) construction (which costs
2838 // lots of communication) by taking the previously constructed
2839 // Import object. We should be able to do this without interfering
2840 // with the implementation of the local part of sparse matrix-matrix
2841 // multply above.
2842 RCP<Teuchos::ParameterList> labelList = rcp(new Teuchos::ParameterList);
2843 labelList->set("Timer Label", label);
2844 if (!params.is_null()) labelList->set("compute global constants", params->get("compute global constants", true));
2845 RCP<const Export<LO, GO, NO>> dummyExport;
2846 C.expertStaticFillComplete(Bview.origMatrix->getDomainMap(), Aview.origMatrix->getRangeMap(), Cimport, dummyExport, labelList);
2847 }
2848}
2849/*********************************************************************************************************/
2850// Kernel method for computing the local portion of C = A*B (reuse)
2851template <class Scalar,
2852 class LocalOrdinal,
2853 class GlobalOrdinal,
2854 class Node>
2855void mult_A_B_reuse(
2856 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
2857 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Bview,
2858 CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& C,
2859 const std::string& label,
2860 const Teuchos::RCP<Teuchos::ParameterList>& params) {
2861 using Teuchos::Array;
2862 using Teuchos::ArrayRCP;
2863 using Teuchos::ArrayView;
2864 using Teuchos::RCP;
2865 using Teuchos::rcp;
2866
2867 // Tpetra typedefs
2868 typedef LocalOrdinal LO;
2869 typedef GlobalOrdinal GO;
2870 typedef Node NO;
2871 typedef Import<LO, GO, NO> import_type;
2872 typedef Map<LO, GO, NO> map_type;
2873
2874 // Kokkos typedefs
2875 typedef typename map_type::local_map_type local_map_type;
2877 typedef typename KCRS::StaticCrsGraphType graph_t;
2878 typedef typename graph_t::row_map_type::non_const_type lno_view_t;
2879 typedef typename NO::execution_space execution_space;
2880 typedef Kokkos::RangePolicy<execution_space, size_t> range_type;
2881 typedef Kokkos::View<LO*, typename lno_view_t::array_layout, typename lno_view_t::device_type> lo_view_t;
2882
2883 Tpetra::Details::ProfilingRegion MM("TpetraExt: MMM: Reuse Cmap");
2884 (void)label;
2885
2886 LO LO_INVALID = Teuchos::OrdinalTraits<LO>::invalid();
2887
2888 // Grab all the maps
2889 RCP<const import_type> Cimport = C.getGraph()->getImporter();
2890 RCP<const map_type> Ccolmap = C.getColMap();
2891 local_map_type Acolmap_local = Aview.colMap->getLocalMap();
2892 local_map_type Browmap_local = Bview.origMatrix->getRowMap()->getLocalMap();
2893 local_map_type Irowmap_local;
2894 if (!Bview.importMatrix.is_null()) Irowmap_local = Bview.importMatrix->getRowMap()->getLocalMap();
2895 local_map_type Bcolmap_local = Bview.origMatrix->getColMap()->getLocalMap();
2896 local_map_type Icolmap_local;
2897 if (!Bview.importMatrix.is_null()) Icolmap_local = Bview.importMatrix->getColMap()->getLocalMap();
2898 local_map_type Ccolmap_local = Ccolmap->getLocalMap();
2899
2900 // Build the final importer / column map, hash table lookups for C
2901 lo_view_t Bcol2Ccol(Kokkos::ViewAllocateWithoutInitializing("Bcol2Ccol"), Bview.colMap->getLocalNumElements()), Icol2Ccol;
2902 {
2903 // Bcol2Col may not be trivial, as Ccolmap is compressed during fillComplete in newmatrix
2904 // So, column map of C may be a strict subset of the column map of B
2905 Kokkos::parallel_for(
2906 range_type(0, Bview.origMatrix->getColMap()->getLocalNumElements()), KOKKOS_LAMBDA(const LO i) {
2907 Bcol2Ccol(i) = Ccolmap_local.getLocalElement(Bcolmap_local.getGlobalElement(i));
2908 });
2909
2910 if (!Bview.importMatrix.is_null()) {
2911 TEUCHOS_TEST_FOR_EXCEPTION(!Cimport->getSourceMap()->isSameAs(*Bview.origMatrix->getDomainMap()),
2912 std::runtime_error, "Tpetra::MMM: Import setUnion messed with the DomainMap in an unfortunate way");
2913
2914 Kokkos::resize(Icol2Ccol, Bview.importMatrix->getColMap()->getLocalNumElements());
2915 Kokkos::parallel_for(
2916 range_type(0, Bview.importMatrix->getColMap()->getLocalNumElements()), KOKKOS_LAMBDA(const LO i) {
2917 Icol2Ccol(i) = Ccolmap_local.getLocalElement(Icolmap_local.getGlobalElement(i));
2918 });
2919 }
2920 }
2921
2922 // Run through all the hash table lookups once and for all
2923 lo_view_t targetMapToOrigRow(Kokkos::ViewAllocateWithoutInitializing("targetMapToOrigRow"), Aview.colMap->getLocalNumElements());
2924 lo_view_t targetMapToImportRow(Kokkos::ViewAllocateWithoutInitializing("targetMapToImportRow"), Aview.colMap->getLocalNumElements());
2925 Kokkos::parallel_for(
2926 range_type(Aview.colMap->getMinLocalIndex(), Aview.colMap->getMaxLocalIndex() + 1), KOKKOS_LAMBDA(const LO i) {
2927 GO aidx = Acolmap_local.getGlobalElement(i);
2928 LO B_LID = Browmap_local.getLocalElement(aidx);
2929 if (B_LID != LO_INVALID) {
2930 targetMapToOrigRow(i) = B_LID;
2931 targetMapToImportRow(i) = LO_INVALID;
2932 } else {
2933 LO I_LID = Irowmap_local.getLocalElement(aidx);
2934 targetMapToOrigRow(i) = LO_INVALID;
2935 targetMapToImportRow(i) = I_LID;
2936 }
2937 });
2938
2939 // Call the actual kernel. We'll rely on partial template specialization to call the correct one ---
2940 // Either the straight-up Tpetra code (SerialNode) or the KokkosKernels one (other NGP node types)
2941 KernelWrappers<Scalar, LocalOrdinal, GlobalOrdinal, Node, lo_view_t>::mult_A_B_reuse_kernel_wrapper(Aview, Bview, targetMapToOrigRow, targetMapToImportRow, Bcol2Ccol, Icol2Ccol, C, Cimport, label, params);
2942}
2943
2944/*********************************************************************************************************/
2945template <class Scalar,
2946 class LocalOrdinal,
2947 class GlobalOrdinal,
2948 class Node,
2949 class LocalOrdinalViewType>
2950void KernelWrappers<Scalar, LocalOrdinal, GlobalOrdinal, Node, LocalOrdinalViewType>::mult_A_B_reuse_kernel_wrapper(CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
2951 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Bview,
2952 const LocalOrdinalViewType& targetMapToOrigRow,
2953 const LocalOrdinalViewType& targetMapToImportRow,
2954 const LocalOrdinalViewType& Bcol2Ccol,
2955 const LocalOrdinalViewType& Icol2Ccol,
2956 CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& C,
2957 Teuchos::RCP<const Import<LocalOrdinal, GlobalOrdinal, Node>> Cimport,
2958 const std::string& label,
2959 const Teuchos::RCP<Teuchos::ParameterList>& params) {
2960 Tpetra::MMdetails::host_mult_A_B_reuse(
2961 Aview, Bview, targetMapToOrigRow, targetMapToImportRow,
2962 Bcol2Ccol, Icol2Ccol, C, Cimport, label, params);
2963}
2964
2965/*********************************************************************************************************/
2966// Kernel method for computing the local portion of C = (I-omega D^{-1} A)*B
2967template <class Scalar,
2968 class LocalOrdinal,
2969 class GlobalOrdinal,
2970 class Node>
2971void jacobi_A_B_newmatrix(
2972 typename Teuchos::ScalarTraits<Scalar>::magnitudeType omega,
2973 const Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Dinv,
2974 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
2975 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Bview,
2976 CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& C,
2977 const std::string& label,
2978 const Teuchos::RCP<Teuchos::ParameterList>& params) {
2979 using Teuchos::Array;
2980 using Teuchos::ArrayRCP;
2981 using Teuchos::ArrayView;
2982 using Teuchos::RCP;
2983 using Teuchos::rcp;
2984 // typedef Scalar SC;
2985 typedef LocalOrdinal LO;
2986 typedef GlobalOrdinal GO;
2987 typedef Node NO;
2988
2989 typedef Import<LO, GO, NO> import_type;
2990 typedef Map<LO, GO, NO> map_type;
2991 typedef typename map_type::local_map_type local_map_type;
2992
2993 // All of the Kokkos typedefs
2995 typedef typename KCRS::StaticCrsGraphType graph_t;
2996 typedef typename graph_t::row_map_type::non_const_type lno_view_t;
2997 typedef typename NO::execution_space execution_space;
2998 typedef Kokkos::RangePolicy<execution_space, size_t> range_type;
2999 typedef Kokkos::View<LO*, typename lno_view_t::array_layout, typename lno_view_t::device_type> lo_view_t;
3000
3001 Tpetra::Details::ProfilingRegion MM3("TpetraExt: Jacobi: M5 Cmap");
3002
3003 LO LO_INVALID = Teuchos::OrdinalTraits<LO>::invalid();
3004
3005 // Build the final importer / column map, hash table lookups for C
3006 RCP<const import_type> Cimport;
3007 RCP<const map_type> Ccolmap;
3008 RCP<const import_type> Bimport = Bview.origMatrix->getGraph()->getImporter();
3009 RCP<const import_type> Iimport = Bview.importMatrix.is_null() ? Teuchos::null : Bview.importMatrix->getGraph()->getImporter();
3010 local_map_type Acolmap_local = Aview.colMap->getLocalMap();
3011 local_map_type Browmap_local = Bview.origMatrix->getRowMap()->getLocalMap();
3012 local_map_type Irowmap_local;
3013 if (!Bview.importMatrix.is_null()) Irowmap_local = Bview.importMatrix->getRowMap()->getLocalMap();
3014 local_map_type Bcolmap_local = Bview.origMatrix->getColMap()->getLocalMap();
3015 local_map_type Icolmap_local;
3016 if (!Bview.importMatrix.is_null()) Icolmap_local = Bview.importMatrix->getColMap()->getLocalMap();
3017
3018 // mfh 27 Sep 2016: Bcol2Ccol is a table that maps from local column
3019 // indices of B, to local column indices of C. (B and C have the
3020 // same number of columns.) The kernel uses this, instead of
3021 // copying the entire input matrix B and converting its column
3022 // indices to those of C.
3023 lo_view_t Bcol2Ccol(Kokkos::ViewAllocateWithoutInitializing("Bcol2Ccol"), Bview.colMap->getLocalNumElements()), Icol2Ccol;
3024
3025 if (Bview.importMatrix.is_null()) {
3026 // mfh 27 Sep 2016: B has no "remotes," so B and C have the same column Map.
3027 Cimport = Bimport;
3028 Ccolmap = Bview.colMap;
3029 // Bcol2Ccol is trivial
3030 // Bcol2Ccol is trivial
3031
3032 Kokkos::RangePolicy<execution_space, LO> range(0, static_cast<LO>(Bview.colMap->getLocalNumElements()));
3033 Kokkos::parallel_for(
3034 range, KOKKOS_LAMBDA(const size_t i) {
3035 Bcol2Ccol(i) = static_cast<LO>(i);
3036 });
3037 } else {
3038 // mfh 27 Sep 2016: B has "remotes," so we need to build the
3039 // column Map of C, as well as C's Import object (from its domain
3040 // Map to its column Map). C's column Map is the union of the
3041 // column Maps of (the local part of) B, and the "remote" part of
3042 // B. Ditto for the Import. We have optimized this "setUnion"
3043 // operation on Import objects and Maps.
3044
3045 // Choose the right variant of setUnion
3046 if (!Bimport.is_null() && !Iimport.is_null()) {
3047 Cimport = Bimport->setUnion(*Iimport, params);
3048 Ccolmap = Cimport->getTargetMap();
3049
3050 } else if (!Bimport.is_null() && Iimport.is_null()) {
3051 Cimport = Bimport->setUnion(params);
3052
3053 } else if (Bimport.is_null() && !Iimport.is_null()) {
3054 Cimport = Iimport->setUnion(params);
3055
3056 } else
3057 throw std::runtime_error("TpetraExt::Jacobi status of matrix importers is nonsensical");
3058
3059 Ccolmap = Cimport->getTargetMap();
3060
3061 TEUCHOS_TEST_FOR_EXCEPTION(!Cimport->getSourceMap()->isSameAs(*Bview.origMatrix->getDomainMap()),
3062 std::runtime_error, "Tpetra:Jacobi Import setUnion messed with the DomainMap in an unfortunate way");
3063
3064 // NOTE: This is not efficient and should be folded into setUnion
3065 //
3066 // mfh 27 Sep 2016: What the above comment means, is that the
3067 // setUnion operation on Import objects could also compute these
3068 // local index - to - local index look-up tables.
3069 Kokkos::resize(Icol2Ccol, Bview.importMatrix->getColMap()->getLocalNumElements());
3070 local_map_type Ccolmap_local = Ccolmap->getLocalMap();
3071 Kokkos::parallel_for(
3072 range_type(0, Bview.origMatrix->getColMap()->getLocalNumElements()), KOKKOS_LAMBDA(const LO i) {
3073 Bcol2Ccol(i) = Ccolmap_local.getLocalElement(Bcolmap_local.getGlobalElement(i));
3074 });
3075 Kokkos::parallel_for(
3076 range_type(0, Bview.importMatrix->getColMap()->getLocalNumElements()), KOKKOS_LAMBDA(const LO i) {
3077 Icol2Ccol(i) = Ccolmap_local.getLocalElement(Icolmap_local.getGlobalElement(i));
3078 });
3079 }
3080
3081 // Replace the column map
3082 //
3083 // mfh 27 Sep 2016: We do this because C was originally created
3084 // without a column Map. Now we have its column Map.
3085 C.replaceColMap(Ccolmap);
3086
3087 // mfh 27 Sep 2016: Construct tables that map from local column
3088 // indices of A, to local row indices of either B_local (the locally
3089 // owned part of B), or B_remote (the "imported" remote part of B).
3090 //
3091 // For column index Aik in row i of A, if the corresponding row of B
3092 // exists in the local part of B ("orig") (which I'll call B_local),
3093 // then targetMapToOrigRow[Aik] is the local index of that row of B.
3094 // Otherwise, targetMapToOrigRow[Aik] is "invalid" (a flag value).
3095 //
3096 // For column index Aik in row i of A, if the corresponding row of B
3097 // exists in the remote part of B ("Import") (which I'll call
3098 // B_remote), then targetMapToImportRow[Aik] is the local index of
3099 // that row of B. Otherwise, targetMapToOrigRow[Aik] is "invalid"
3100 // (a flag value).
3101
3102 // Run through all the hash table lookups once and for all
3103 lo_view_t targetMapToOrigRow(Kokkos::ViewAllocateWithoutInitializing("targetMapToOrigRow"), Aview.colMap->getLocalNumElements());
3104 lo_view_t targetMapToImportRow(Kokkos::ViewAllocateWithoutInitializing("targetMapToImportRow"), Aview.colMap->getLocalNumElements());
3105 Kokkos::parallel_for(
3106 range_type(Aview.colMap->getMinLocalIndex(), Aview.colMap->getMaxLocalIndex() + 1), KOKKOS_LAMBDA(const LO i) {
3107 GO aidx = Acolmap_local.getGlobalElement(i);
3108 LO B_LID = Browmap_local.getLocalElement(aidx);
3109 if (B_LID != LO_INVALID) {
3110 targetMapToOrigRow(i) = B_LID;
3111 targetMapToImportRow(i) = LO_INVALID;
3112 } else {
3113 LO I_LID = Irowmap_local.getLocalElement(aidx);
3114 targetMapToOrigRow(i) = LO_INVALID;
3115 targetMapToImportRow(i) = I_LID;
3116 }
3117 });
3118
3119 // Call the actual kernel. We'll rely on partial template specialization to call the correct one ---
3120 // Either the straight-up Tpetra code (SerialNode) or the KokkosKernels one (other NGP node types)
3121 KernelWrappers2<Scalar, LocalOrdinal, GlobalOrdinal, Node, lo_view_t>::jacobi_A_B_newmatrix_kernel_wrapper(omega, Dinv, Aview, Bview, targetMapToOrigRow, targetMapToImportRow, Bcol2Ccol, Icol2Ccol, C, Cimport, label, params);
3122}
3123
3124/*********************************************************************************************************/
3125// Jacobi AB NewMatrix Kernel wrappers (Default non-threaded version)
3126// Kernel method for computing the local portion of C = (I-omega D^{-1} A)*B
3127
3128template <class Scalar,
3129 class LocalOrdinal,
3130 class GlobalOrdinal,
3131 class Node,
3132 class LocalOrdinalViewType>
3133void KernelWrappers2<Scalar, LocalOrdinal, GlobalOrdinal, Node, LocalOrdinalViewType>::jacobi_A_B_newmatrix_kernel_wrapper(typename Teuchos::ScalarTraits<Scalar>::magnitudeType omega,
3134 const Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Dinv,
3135 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
3136 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Bview,
3137 const LocalOrdinalViewType& targetMapToOrigRow,
3138 const LocalOrdinalViewType& targetMapToImportRow,
3139 const LocalOrdinalViewType& Bcol2Ccol,
3140 const LocalOrdinalViewType& Icol2Ccol,
3141 CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& C,
3142 Teuchos::RCP<const Import<LocalOrdinal, GlobalOrdinal, Node>> Cimport,
3143 const std::string& label,
3144 const Teuchos::RCP<Teuchos::ParameterList>& params) {
3145 Tpetra::Details::ProfilingRegion MM("TpetraExt: Jacobi: Newmatrix SerialCore");
3146
3147 using Teuchos::Array;
3148 using Teuchos::ArrayRCP;
3149 using Teuchos::ArrayView;
3150 using Teuchos::RCP;
3151 using Teuchos::rcp;
3152
3153 // Lots and lots of typedefs
3154 typedef typename Tpetra::CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>::local_matrix_host_type KCRS;
3155 typedef typename KCRS::StaticCrsGraphType graph_t;
3156 typedef typename graph_t::row_map_type::const_type c_lno_view_t;
3157 typedef typename graph_t::row_map_type::non_const_type lno_view_t;
3158 typedef typename graph_t::entries_type::non_const_type lno_nnz_view_t;
3159 typedef typename KCRS::values_type::non_const_type scalar_view_t;
3160
3161 // Jacobi-specific
3162 typedef typename scalar_view_t::memory_space scalar_memory_space;
3163
3164 typedef Scalar SC;
3165 typedef LocalOrdinal LO;
3166 typedef GlobalOrdinal GO;
3167 typedef Node NO;
3168
3169 typedef Map<LO, GO, NO> map_type;
3170 size_t ST_INVALID = Teuchos::OrdinalTraits<LO>::invalid();
3171 LO LO_INVALID = Teuchos::OrdinalTraits<LO>::invalid();
3172
3173 // Sizes
3174 RCP<const map_type> Ccolmap = C.getColMap();
3175 size_t m = Aview.origMatrix->getLocalNumRows();
3176 size_t n = Ccolmap->getLocalNumElements();
3177 size_t b_max_nnz_per_row = Bview.origMatrix->getLocalMaxNumRowEntries();
3178
3179 // Grab the Kokkos::SparseCrsMatrices & inner stuff
3180 const KCRS Amat = Aview.origMatrix->getLocalMatrixHost();
3181 const KCRS Bmat = Bview.origMatrix->getLocalMatrixHost();
3182
3183 c_lno_view_t Arowptr = Amat.graph.row_map, Browptr = Bmat.graph.row_map;
3184 const lno_nnz_view_t Acolind = Amat.graph.entries, Bcolind = Bmat.graph.entries;
3185 const scalar_view_t Avals = Amat.values, Bvals = Bmat.values;
3186
3187 c_lno_view_t Irowptr;
3188 lno_nnz_view_t Icolind;
3189 scalar_view_t Ivals;
3190 if (!Bview.importMatrix.is_null()) {
3191 auto lclB = Bview.importMatrix->getLocalMatrixHost();
3192 Irowptr = lclB.graph.row_map;
3193 Icolind = lclB.graph.entries;
3194 Ivals = lclB.values;
3195 b_max_nnz_per_row = std::max(b_max_nnz_per_row, Bview.importMatrix->getLocalMaxNumRowEntries());
3196 }
3197
3198 // Jacobi-specific inner stuff
3199 auto Dvals =
3200 Dinv.template getLocalView<scalar_memory_space>(Access::ReadOnly);
3201
3202 // Teuchos::ArrayView::operator[].
3203 // The status array will contain the index into colind where this entry was last deposited.
3204 // c_status[i] < CSR_ip - not in the row yet.
3205 // c_status[i] >= CSR_ip, this is the entry where you can find the data
3206 // We start with this filled with INVALID's indicating that there are no entries yet.
3207 // Sadly, this complicates the code due to the fact that size_t's are unsigned.
3208 size_t INVALID = Teuchos::OrdinalTraits<size_t>::invalid();
3209 Array<size_t> c_status(n, ST_INVALID);
3210
3211 // Classic csr assembly (low memory edition)
3212 //
3213 // mfh 27 Sep 2016: C_estimate_nnz does not promise an upper bound.
3214 // The method loops over rows of A, and may resize after processing
3215 // each row. Chris Siefert says that this reflects experience in
3216 // ML; for the non-threaded case, ML found it faster to spend less
3217 // effort on estimation and risk an occasional reallocation.
3218 size_t CSR_alloc = std::max(C_estimate_nnz(*Aview.origMatrix, *Bview.origMatrix), n);
3219 lno_view_t Crowptr(Kokkos::ViewAllocateWithoutInitializing("Crowptr"), m + 1);
3220 lno_nnz_view_t Ccolind(Kokkos::ViewAllocateWithoutInitializing("Ccolind"), CSR_alloc);
3221 scalar_view_t Cvals(Kokkos::ViewAllocateWithoutInitializing("Cvals"), CSR_alloc);
3222 size_t CSR_ip = 0, OLD_ip = 0;
3223
3224 const SC SC_ZERO = Teuchos::ScalarTraits<Scalar>::zero();
3225
3226 // mfh 27 Sep 2016: Here is the local sparse matrix-matrix multiply
3227 // routine. The routine computes
3228 //
3229 // C := (I - omega * D^{-1} * A) * (B_local + B_remote)).
3230 //
3231 // This corresponds to one sweep of (weighted) Jacobi.
3232 //
3233 // For column index Aik in row i of A, targetMapToOrigRow[Aik] tells
3234 // you whether the corresponding row of B belongs to B_local
3235 // ("orig") or B_remote ("Import").
3236
3237 // For each row of A/C
3238 for (size_t i = 0; i < m; i++) {
3239 // mfh 27 Sep 2016: m is the number of rows in the input matrix A
3240 // on the calling process.
3241 Crowptr[i] = CSR_ip;
3242 SC minusOmegaDval = -omega * Dvals(i, 0);
3243
3244 // Entries of B
3245 for (size_t j = Browptr[i]; j < Browptr[i + 1]; j++) {
3246 Scalar Bval = Bvals[j];
3247 if (Bval == SC_ZERO)
3248 continue;
3249 LO Bij = Bcolind[j];
3250 LO Cij = Bcol2Ccol[Bij];
3251
3252 // Assume no repeated entries in B
3253 c_status[Cij] = CSR_ip;
3254 Ccolind[CSR_ip] = Cij;
3255 Cvals[CSR_ip] = Bvals[j];
3256 CSR_ip++;
3257 }
3258
3259 // Entries of -omega * Dinv * A * B
3260 for (size_t k = Arowptr[i]; k < Arowptr[i + 1]; k++) {
3261 LO Aik = Acolind[k];
3262 const SC Aval = Avals[k];
3263 if (Aval == SC_ZERO)
3264 continue;
3265
3266 if (targetMapToOrigRow[Aik] != LO_INVALID) {
3267 // Local matrix
3268 size_t Bk = static_cast<size_t>(targetMapToOrigRow[Aik]);
3269
3270 for (size_t j = Browptr[Bk]; j < Browptr[Bk + 1]; ++j) {
3271 LO Bkj = Bcolind[j];
3272 LO Cij = Bcol2Ccol[Bkj];
3273
3274 if (c_status[Cij] == INVALID || c_status[Cij] < OLD_ip) {
3275 // New entry
3276 c_status[Cij] = CSR_ip;
3277 Ccolind[CSR_ip] = Cij;
3278 Cvals[CSR_ip] = minusOmegaDval * Aval * Bvals[j];
3279 CSR_ip++;
3280
3281 } else {
3282 Cvals[c_status[Cij]] += minusOmegaDval * Aval * Bvals[j];
3283 }
3284 }
3285
3286 } else {
3287 // Remote matrix
3288 size_t Ik = static_cast<size_t>(targetMapToImportRow[Aik]);
3289 for (size_t j = Irowptr[Ik]; j < Irowptr[Ik + 1]; ++j) {
3290 LO Ikj = Icolind[j];
3291 LO Cij = Icol2Ccol[Ikj];
3292
3293 if (c_status[Cij] == INVALID || c_status[Cij] < OLD_ip) {
3294 // New entry
3295 c_status[Cij] = CSR_ip;
3296 Ccolind[CSR_ip] = Cij;
3297 Cvals[CSR_ip] = minusOmegaDval * Aval * Ivals[j];
3298 CSR_ip++;
3299 } else {
3300 Cvals[c_status[Cij]] += minusOmegaDval * Aval * Ivals[j];
3301 }
3302 }
3303 }
3304 }
3305
3306 // Resize for next pass if needed
3307 if (i + 1 < m && CSR_ip + std::min(n, (Arowptr[i + 2] - Arowptr[i + 1] + 1) * b_max_nnz_per_row) > CSR_alloc) {
3308 CSR_alloc *= 2;
3309 Kokkos::resize(Ccolind, CSR_alloc);
3310 Kokkos::resize(Cvals, CSR_alloc);
3311 }
3312 OLD_ip = CSR_ip;
3313 }
3314 Crowptr[m] = CSR_ip;
3315
3316 // Downward resize
3317 Kokkos::resize(Ccolind, CSR_ip);
3318 Kokkos::resize(Cvals, CSR_ip);
3319
3320 {
3321 Tpetra::Details::ProfilingRegion MM2("TpetraExt: Jacobi: Newmatrix Final Sort");
3322
3323 // Replace the column map
3324 //
3325 // mfh 27 Sep 2016: We do this because C was originally created
3326 // without a column Map. Now we have its column Map.
3327 C.replaceColMap(Ccolmap);
3328
3329 // Final sort & set of CRS arrays
3330 //
3331 // TODO (mfh 27 Sep 2016) Will the thread-parallel "local" sparse
3332 // matrix-matrix multiply routine sort the entries for us?
3333 // Final sort & set of CRS arrays
3334 if (params.is_null() || params->get("sort entries", true)) {
3335 // Tpetra's serial SpGEMM results in almost sorted matrices. Use shell sort.
3336 Import_Util::sortCrsEntries(Crowptr, Ccolind, Cvals);
3337 }
3338 C.setAllValues(Crowptr, Ccolind, Cvals);
3339 }
3340 {
3341 Tpetra::Details::ProfilingRegion MM3("TpetraExt: Jacobi: Newmatrix ESFC");
3342
3343 // Final FillComplete
3344 //
3345 // mfh 27 Sep 2016: So-called "expert static fill complete" bypasses
3346 // Import (from domain Map to column Map) construction (which costs
3347 // lots of communication) by taking the previously constructed
3348 // Import object. We should be able to do this without interfering
3349 // with the implementation of the local part of sparse matrix-matrix
3350 // multply above
3351 RCP<Teuchos::ParameterList> labelList = rcp(new Teuchos::ParameterList);
3352 labelList->set("Timer Label", label);
3353 if (!params.is_null()) labelList->set("compute global constants", params->get("compute global constants", true));
3354 RCP<const Export<LO, GO, NO>> dummyExport;
3355 C.expertStaticFillComplete(Bview.origMatrix->getDomainMap(), Aview.origMatrix->getRangeMap(), Cimport, dummyExport, labelList);
3356 }
3357}
3358
3359/*********************************************************************************************************/
3360// Kernel method for computing the local portion of C = (I-omega D^{-1} A)*B
3361template <class Scalar,
3362 class LocalOrdinal,
3363 class GlobalOrdinal,
3364 class Node>
3365void jacobi_A_B_reuse(
3366 typename Teuchos::ScalarTraits<Scalar>::magnitudeType omega,
3367 const Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Dinv,
3368 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
3369 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Bview,
3370 CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& C,
3371 const std::string& label,
3372 const Teuchos::RCP<Teuchos::ParameterList>& params) {
3373 using Teuchos::Array;
3374 using Teuchos::ArrayRCP;
3375 using Teuchos::ArrayView;
3376 using Teuchos::RCP;
3377 using Teuchos::rcp;
3378
3379 typedef LocalOrdinal LO;
3380 typedef GlobalOrdinal GO;
3381 typedef Node NO;
3382
3383 typedef Import<LO, GO, NO> import_type;
3384 typedef Map<LO, GO, NO> map_type;
3385
3386 // Kokkos typedefs
3387 typedef typename map_type::local_map_type local_map_type;
3389 typedef typename KCRS::StaticCrsGraphType graph_t;
3390 typedef typename graph_t::row_map_type::non_const_type lno_view_t;
3391 typedef typename NO::execution_space execution_space;
3392 typedef Kokkos::RangePolicy<execution_space, size_t> range_type;
3393 typedef Kokkos::View<LO*, typename lno_view_t::array_layout, typename lno_view_t::device_type> lo_view_t;
3394
3395 RCP<const import_type> Cimport = C.getGraph()->getImporter();
3396 lo_view_t Bcol2Ccol, Icol2Ccol;
3397 lo_view_t targetMapToOrigRow;
3398 lo_view_t targetMapToImportRow;
3399 {
3400 Tpetra::Details::ProfilingRegion MM("TpetraExt: Jacobi: Reuse Cmap");
3401
3402 LO LO_INVALID = Teuchos::OrdinalTraits<LO>::invalid();
3403
3404 // Grab all the maps
3405 RCP<const map_type> Ccolmap = C.getColMap();
3406 local_map_type Acolmap_local = Aview.colMap->getLocalMap();
3407 local_map_type Browmap_local = Bview.origMatrix->getRowMap()->getLocalMap();
3408 local_map_type Irowmap_local;
3409 if (!Bview.importMatrix.is_null()) Irowmap_local = Bview.importMatrix->getRowMap()->getLocalMap();
3410 local_map_type Bcolmap_local = Bview.origMatrix->getColMap()->getLocalMap();
3411 local_map_type Icolmap_local;
3412 if (!Bview.importMatrix.is_null()) Icolmap_local = Bview.importMatrix->getColMap()->getLocalMap();
3413 local_map_type Ccolmap_local = Ccolmap->getLocalMap();
3414
3415 // Build the final importer / column map, hash table lookups for C
3416 Bcol2Ccol = lo_view_t(Kokkos::ViewAllocateWithoutInitializing("Bcol2Ccol"), Bview.colMap->getLocalNumElements());
3417 {
3418 // Bcol2Col may not be trivial, as Ccolmap is compressed during fillComplete in newmatrix
3419 // So, column map of C may be a strict subset of the column map of B
3420 Kokkos::parallel_for(
3421 range_type(0, Bview.origMatrix->getColMap()->getLocalNumElements()), KOKKOS_LAMBDA(const LO i) {
3422 Bcol2Ccol(i) = Ccolmap_local.getLocalElement(Bcolmap_local.getGlobalElement(i));
3423 });
3424
3425 if (!Bview.importMatrix.is_null()) {
3426 TEUCHOS_TEST_FOR_EXCEPTION(!Cimport->getSourceMap()->isSameAs(*Bview.origMatrix->getDomainMap()),
3427 std::runtime_error, "Tpetra::Jacobi: Import setUnion messed with the DomainMap in an unfortunate way");
3428
3429 Kokkos::resize(Icol2Ccol, Bview.importMatrix->getColMap()->getLocalNumElements());
3430 Kokkos::parallel_for(
3431 range_type(0, Bview.importMatrix->getColMap()->getLocalNumElements()), KOKKOS_LAMBDA(const LO i) {
3432 Icol2Ccol(i) = Ccolmap_local.getLocalElement(Icolmap_local.getGlobalElement(i));
3433 });
3434 }
3435 }
3436
3437 // Run through all the hash table lookups once and for all
3438 targetMapToOrigRow = lo_view_t(Kokkos::ViewAllocateWithoutInitializing("targetMapToOrigRow"), Aview.colMap->getLocalNumElements());
3439 targetMapToImportRow = lo_view_t(Kokkos::ViewAllocateWithoutInitializing("targetMapToImportRow"), Aview.colMap->getLocalNumElements());
3440 Kokkos::parallel_for(
3441 range_type(Aview.colMap->getMinLocalIndex(), Aview.colMap->getMaxLocalIndex() + 1), KOKKOS_LAMBDA(const LO i) {
3442 GO aidx = Acolmap_local.getGlobalElement(i);
3443 LO B_LID = Browmap_local.getLocalElement(aidx);
3444 if (B_LID != LO_INVALID) {
3445 targetMapToOrigRow(i) = B_LID;
3446 targetMapToImportRow(i) = LO_INVALID;
3447 } else {
3448 LO I_LID = Irowmap_local.getLocalElement(aidx);
3449 targetMapToOrigRow(i) = LO_INVALID;
3450 targetMapToImportRow(i) = I_LID;
3451 }
3452 });
3453 }
3454
3455 // Call the actual kernel. We'll rely on partial template specialization to call the correct one ---
3456 // Either the straight-up Tpetra code (SerialNode) or the KokkosKernels one (other NGP node types)
3457 KernelWrappers2<Scalar, LocalOrdinal, GlobalOrdinal, Node, lo_view_t>::jacobi_A_B_reuse_kernel_wrapper(omega, Dinv, Aview, Bview, targetMapToOrigRow, targetMapToImportRow, Bcol2Ccol, Icol2Ccol, C, Cimport, label, params);
3458}
3459
3460/*********************************************************************************************************/
3461template <class Scalar,
3462 class LocalOrdinal,
3463 class GlobalOrdinal,
3464 class Node,
3465 class LocalOrdinalViewType>
3466void KernelWrappers2<Scalar, LocalOrdinal, GlobalOrdinal, Node, LocalOrdinalViewType>::jacobi_A_B_reuse_kernel_wrapper(typename Teuchos::ScalarTraits<Scalar>::magnitudeType omega,
3467 const Vector<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Dinv,
3468 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
3469 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Bview,
3470 const LocalOrdinalViewType& targetMapToOrigRow,
3471 const LocalOrdinalViewType& targetMapToImportRow,
3472 const LocalOrdinalViewType& Bcol2Ccol,
3473 const LocalOrdinalViewType& Icol2Ccol,
3474 CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& C,
3475 Teuchos::RCP<const Import<LocalOrdinal, GlobalOrdinal, Node>> Cimport,
3476 const std::string& label,
3477 const Teuchos::RCP<Teuchos::ParameterList>& params) {
3478 host_jacobi_A_B_reuse(omega, Dinv, Aview, Bview, targetMapToOrigRow, targetMapToImportRow, Bcol2Ccol, Icol2Ccol, C, Cimport, label, params);
3479}
3480
3481/*********************************************************************************************************/
3482template <class Scalar,
3483 class LocalOrdinal,
3484 class GlobalOrdinal,
3485 class Node>
3486void import_and_extract_views(
3487 const CrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& A,
3488 Teuchos::RCP<const Map<LocalOrdinal, GlobalOrdinal, Node>> targetMap,
3489 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
3490 Teuchos::RCP<const Import<LocalOrdinal, GlobalOrdinal, Node>> prototypeImporter,
3491 bool userAssertsThereAreNoRemotes,
3492 const std::string& label,
3493 const Teuchos::RCP<Teuchos::ParameterList>& params) {
3494 using Teuchos::Array;
3495 using Teuchos::ArrayView;
3496 using Teuchos::null;
3497 using Teuchos::RCP;
3498 using Teuchos::rcp;
3499
3500 typedef Scalar SC;
3501 typedef LocalOrdinal LO;
3502 typedef GlobalOrdinal GO;
3503 typedef Node NO;
3504
3505 typedef Map<LO, GO, NO> map_type;
3506 typedef Import<LO, GO, NO> import_type;
3507 typedef CrsMatrix<SC, LO, GO, NO> crs_matrix_type;
3508
3509 RCP<Tpetra::Details::ProfilingRegion> MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: MMM: I&X Alloc"));
3510
3511 // The goal of this method is to populate the 'Aview' struct with views of the
3512 // rows of A, including all rows that correspond to elements in 'targetMap'.
3513 //
3514 // If targetMap includes local elements that correspond to remotely-owned rows
3515 // of A, then those remotely-owned rows will be imported into
3516 // 'Aview.importMatrix', and views of them will be included in 'Aview'.
3517 Aview.deleteContents();
3518
3519 Aview.origMatrix = rcp(&A, false);
3520 // trigger creation of int-typed row pointer array for use in TPLs, but don't actually need it here
3521 Aview.origMatrix->getApplyHelper();
3522 Aview.origRowMap = A.getRowMap();
3523 Aview.rowMap = targetMap;
3524 Aview.colMap = A.getColMap();
3525 Aview.domainMap = A.getDomainMap();
3526 Aview.importColMap = null;
3527 RCP<const map_type> rowMap = A.getRowMap();
3528 const int numProcs = rowMap->getComm()->getSize();
3529
3530 // Short circuit if the user swears there are no remotes (or if we're in serial)
3531 if (userAssertsThereAreNoRemotes || numProcs < 2)
3532 return;
3533
3534 RCP<const import_type> importer;
3535 if (params != null && params->isParameter("importer")) {
3536 importer = params->get<RCP<const import_type>>("importer");
3537
3538 } else {
3539 MM = Teuchos::null;
3540 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: MMM: I&X RemoteMap"));
3541
3542 // Mark each row in targetMap as local or remote, and go ahead and get a view
3543 // for the local rows
3544 RCP<const map_type> remoteRowMap;
3545 size_t numRemote = 0;
3546 int mode = 0;
3547 if (!prototypeImporter.is_null() &&
3548 prototypeImporter->getSourceMap()->isSameAs(*rowMap) &&
3549 prototypeImporter->getTargetMap()->isSameAs(*targetMap)) {
3550 // We have a valid prototype importer --- ask it for the remotes
3551 Tpetra::Details::ProfilingRegion MM2("TpetraExt: MMM: I&X RemoteMap: Mode1");
3552
3553 ArrayView<const LO> remoteLIDs = prototypeImporter->getRemoteLIDs();
3554 numRemote = prototypeImporter->getNumRemoteIDs();
3555
3556 Array<GO> remoteRows(numRemote);
3557 for (size_t i = 0; i < numRemote; i++)
3558 remoteRows[i] = targetMap->getGlobalElement(remoteLIDs[i]);
3559
3560 remoteRowMap = rcp(new map_type(Teuchos::OrdinalTraits<global_size_t>::invalid(), remoteRows(),
3561 rowMap->getIndexBase(), rowMap->getComm(), params));
3562 mode = 1;
3563
3564 } else if (prototypeImporter.is_null()) {
3565 // No prototype importer --- count the remotes the hard way
3566 Tpetra::Details::ProfilingRegion MM2("TpetraExt: MMM: I&X RemoteMap: Mode2");
3567
3568 ArrayView<const GO> rows = targetMap->getLocalElementList();
3569 size_t numRows = targetMap->getLocalNumElements();
3570
3571 Array<GO> remoteRows(numRows);
3572 for (size_t i = 0; i < numRows; ++i) {
3573 const LO mlid = rowMap->getLocalElement(rows[i]);
3574
3575 if (mlid == Teuchos::OrdinalTraits<LO>::invalid())
3576 remoteRows[numRemote++] = rows[i];
3577 }
3578 remoteRows.resize(numRemote);
3579 remoteRowMap = rcp(new map_type(Teuchos::OrdinalTraits<global_size_t>::invalid(), remoteRows(),
3580 rowMap->getIndexBase(), rowMap->getComm(), params));
3581 mode = 2;
3582
3583 } else {
3584 // PrototypeImporter is bad. But if we're in serial that's OK.
3585 mode = 3;
3586 }
3587
3588 if (numProcs < 2) {
3589 TEUCHOS_TEST_FOR_EXCEPTION(numRemote > 0, std::runtime_error,
3590 "MatrixMatrix::import_and_extract_views ERROR, numProcs < 2 but attempting to import remote matrix rows.");
3591 // If only one processor we don't need to import any remote rows, so return.
3592 return;
3593 }
3594
3595 //
3596 // Now we will import the needed remote rows of A, if the remoteRowMap has entries
3597 //
3598 if (!remoteRowMap.is_null() && (remoteRowMap->getGlobalNumElements() > 0)) {
3599 MM = Teuchos::null;
3600 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: MMM: I&X Import-2"));
3601
3602 // Create an importer with target-map remoteRowMap and source-map rowMap.
3603 if (mode == 1)
3604 importer = prototypeImporter->createRemoteOnlyImport(remoteRowMap);
3605 else if (mode == 2)
3606 importer = rcp(new import_type(rowMap, remoteRowMap));
3607 else
3608 throw std::runtime_error("prototypeImporter->SourceMap() does not match A.getRowMap()!");
3609 }
3610
3611 if (params != null)
3612 params->set("importer", importer);
3613 }
3614
3615 if (importer != null) {
3616 MM = Teuchos::null;
3617 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: MMM: I&X Import-3"));
3618
3619 // Now create a new matrix into which we can import the remote rows of A that we need.
3620 Teuchos::ParameterList labelList;
3621 labelList.set("Timer Label", label);
3622 // Minor speedup tweak - avoid computing the global constants
3623 labelList.set("compute global constants", false);
3624 auto& labelList_subList = labelList.sublist("matrixmatrix: kernel params", false);
3625 labelList_subList.set("isMatrixMatrix_TransferAndFillComplete", true);
3626
3627 if (!params.is_null()) {
3628 if (params->isParameter("compute global constants"))
3629 labelList.set("compute global constants", params->get<bool>("compute global constants"));
3630 }
3631
3632 Aview.importMatrix = Tpetra::importAndFillCompleteCrsMatrix<crs_matrix_type>(rcpFromRef(A), *importer,
3633 A.getDomainMap(), importer->getTargetMap(), rcpFromRef(labelList));
3634 // trigger creation of int-typed row pointer array for use in TPLs, but don't actually need it here
3635 Aview.importMatrix->getApplyHelper();
3636
3637#if 0
3638 // Disabled code for dumping input matrices
3639 static int count=0;
3640 char str[80];
3641 sprintf(str,"import_matrix.%d.dat",count);
3643 count++;
3644#endif
3645
3646#ifdef HAVE_TPETRA_MMM_STATISTICS
3647 printMultiplicationStatistics(importer, label + std::string(" I&X MMM"));
3648#endif
3649
3650 MM = Teuchos::null;
3651 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: MMM: I&X Import-4"));
3652
3653 // Save the column map of the imported matrix, so that we can convert indices back to global for arithmetic later
3654 Aview.importColMap = Aview.importMatrix->getColMap();
3655 MM = Teuchos::null;
3656 }
3657}
3658
3659/*********************************************************************************************************/
3660template <class Scalar,
3661 class LocalOrdinal,
3662 class GlobalOrdinal,
3663 class Node>
3664void import_and_extract_views(
3665 const BlockCrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>& M,
3666 Teuchos::RCP<const Map<LocalOrdinal, GlobalOrdinal, Node>> targetMap,
3667 BlockCrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Mview,
3668 Teuchos::RCP<const Import<LocalOrdinal, GlobalOrdinal, Node>> prototypeImporter,
3669 bool userAssertsThereAreNoRemotes) {
3670 using Teuchos::Array;
3671 using Teuchos::ArrayView;
3672 using Teuchos::null;
3673 using Teuchos::RCP;
3674 using Teuchos::rcp;
3675
3676 typedef Scalar SC;
3677 typedef LocalOrdinal LO;
3678 typedef GlobalOrdinal GO;
3679 typedef Node NO;
3680
3681 typedef Map<LO, GO, NO> map_type;
3682 typedef Import<LO, GO, NO> import_type;
3683 typedef BlockCrsMatrix<SC, LO, GO, NO> blockcrs_matrix_type;
3684
3685 // The goal of this method is to populate the 'Mview' struct with views of the
3686 // rows of M, including all rows that correspond to elements in 'targetMap'.
3687 //
3688 // If targetMap includes local elements that correspond to remotely-owned rows
3689 // of M, then those remotely-owned rows will be imported into
3690 // 'Mview.importMatrix', and views of them will be included in 'Mview'.
3691 Mview.deleteContents();
3692
3693 Mview.origMatrix = rcp(&M, false);
3694 // trigger creation of int-typed row pointer array for use in TPLs, but don't actually need it here
3695 Mview.origMatrix->getApplyHelper();
3696 Mview.origRowMap = M.getRowMap();
3697 Mview.rowMap = targetMap;
3698 Mview.colMap = M.getColMap();
3699 Mview.importColMap = null;
3700 RCP<const map_type> rowMap = M.getRowMap();
3701 const int numProcs = rowMap->getComm()->getSize();
3702
3703 // Short circuit if the user swears there are no remotes (or if we're in serial)
3704 if (userAssertsThereAreNoRemotes || numProcs < 2) return;
3705
3706 // Mark each row in targetMap as local or remote, and go ahead and get a view
3707 // for the local rows
3708 RCP<const map_type> remoteRowMap;
3709 size_t numRemote = 0;
3710 int mode = 0;
3711 if (!prototypeImporter.is_null() &&
3712 prototypeImporter->getSourceMap()->isSameAs(*rowMap) &&
3713 prototypeImporter->getTargetMap()->isSameAs(*targetMap)) {
3714 // We have a valid prototype importer --- ask it for the remotes
3715 ArrayView<const LO> remoteLIDs = prototypeImporter->getRemoteLIDs();
3716 numRemote = prototypeImporter->getNumRemoteIDs();
3717
3718 Array<GO> remoteRows(numRemote);
3719 for (size_t i = 0; i < numRemote; i++)
3720 remoteRows[i] = targetMap->getGlobalElement(remoteLIDs[i]);
3721
3722 remoteRowMap = rcp(new map_type(Teuchos::OrdinalTraits<global_size_t>::invalid(), remoteRows(),
3723 rowMap->getIndexBase(), rowMap->getComm()));
3724 mode = 1;
3725
3726 } else if (prototypeImporter.is_null()) {
3727 // No prototype importer --- count the remotes the hard way
3728 ArrayView<const GO> rows = targetMap->getLocalElementList();
3729 size_t numRows = targetMap->getLocalNumElements();
3730
3731 Array<GO> remoteRows(numRows);
3732 for (size_t i = 0; i < numRows; ++i) {
3733 const LO mlid = rowMap->getLocalElement(rows[i]);
3734
3735 if (mlid == Teuchos::OrdinalTraits<LO>::invalid())
3736 remoteRows[numRemote++] = rows[i];
3737 }
3738 remoteRows.resize(numRemote);
3739 remoteRowMap = rcp(new map_type(Teuchos::OrdinalTraits<global_size_t>::invalid(), remoteRows(),
3740 rowMap->getIndexBase(), rowMap->getComm()));
3741 mode = 2;
3742
3743 } else {
3744 // PrototypeImporter is bad. But if we're in serial that's OK.
3745 mode = 3;
3746 }
3747
3748 if (numProcs < 2) {
3749 TEUCHOS_TEST_FOR_EXCEPTION(numRemote > 0, std::runtime_error,
3750 "MatrixMatrix::import_and_extract_views ERROR, numProcs < 2 but attempting to import remote matrix rows.");
3751 // If only one processor we don't need to import any remote rows, so return.
3752 return;
3753 }
3754
3755 // Now we will import the needed remote rows of M, if the remoteRowMap has entries
3756
3757 RCP<const import_type> importer;
3758
3759 if (!remoteRowMap.is_null() && (remoteRowMap->getGlobalNumElements() > 0)) {
3760 // Create an importer with target-map remoteRowMap and source-map rowMap.
3761 if (mode == 1)
3762 importer = prototypeImporter->createRemoteOnlyImport(remoteRowMap);
3763 else if (mode == 2)
3764 importer = rcp(new import_type(rowMap, remoteRowMap));
3765 else
3766 throw std::runtime_error("prototypeImporter->SourceMap() does not match M.getRowMap()!");
3767 }
3768
3769 if (importer != null) {
3770 // Get import matrix
3771 // TODO: create the int-typed row-pointer here
3772 Mview.importMatrix = Tpetra::importAndFillCompleteBlockCrsMatrix<blockcrs_matrix_type>(rcpFromRef(M), *importer);
3773 // trigger creation of int-typed row pointer array for use in TPLs, but don't actually need it here
3774 Mview.importMatrix->getApplyHelper();
3775 // Save the column map of the imported matrix, so that we can convert indices
3776 // back to global for arithmetic later
3777 Mview.importColMap = Mview.importMatrix->getColMap();
3778 }
3779}
3780
3781/*********************************************************************************************************/
3782// This only merges matrices that look like B & Bimport, aka, they have no overlapping rows
3783template <class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node, class LocalOrdinalViewType>
3785merge_matrices(CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
3786 CrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Bview,
3787 const LocalOrdinalViewType& Acol2Brow,
3788 const LocalOrdinalViewType& Acol2Irow,
3789 const LocalOrdinalViewType& Bcol2Ccol,
3790 const LocalOrdinalViewType& Icol2Ccol,
3791 const size_t mergedNodeNumCols) {
3792 using Teuchos::RCP;
3794 typedef typename KCRS::StaticCrsGraphType graph_t;
3795 typedef typename graph_t::row_map_type::non_const_type lno_view_t;
3796 typedef typename graph_t::entries_type::non_const_type lno_nnz_view_t;
3797 typedef typename KCRS::values_type::non_const_type scalar_view_t;
3798 // Grab the Kokkos::SparseCrsMatrices
3799 const KCRS Ak = Aview.origMatrix->getLocalMatrixDevice();
3800 const KCRS Bk = Bview.origMatrix->getLocalMatrixDevice();
3801
3802 // We need to do this dance if either (a) We have Bimport or (b) We don't A's colMap is not the same as B's rowMap
3803 if (!Bview.importMatrix.is_null() || (Bview.importMatrix.is_null() && (&*Aview.origMatrix->getGraph()->getColMap() != &*Bview.origMatrix->getGraph()->getRowMap()))) {
3804 // We do have a Bimport
3805 // NOTE: We're going merge Borig and Bimport into a single matrix and reindex the columns *before* we multiply.
3806 // This option was chosen because we know we don't have any duplicate entries, so we can allocate once.
3807
3808 KCRS Iks;
3809 if (!Bview.importMatrix.is_null()) Iks = Bview.importMatrix->getLocalMatrixDevice();
3810
3811 size_t merge_numrows = Ak.numCols();
3812
3813 // The last entry of this at least, need to be initialized
3814 lno_view_t Mrowptr("Mrowptr", merge_numrows + 1);
3815
3816 const LocalOrdinal LO_INVALID = Teuchos::OrdinalTraits<LocalOrdinal>::invalid();
3817
3818 // Use a Kokkos::parallel_scan to build the rowptr
3819 typedef typename Node::execution_space execution_space;
3820 typedef Kokkos::RangePolicy<execution_space, size_t> range_type;
3821 Kokkos::parallel_scan(
3822 "Tpetra_MatrixMatrix_merge_matrices_buildRowptr", range_type(0, merge_numrows),
3823 KOKKOS_LAMBDA(const size_t i, size_t& update, const bool final) {
3824 if (final) Mrowptr(i) = update;
3825 // Get the row count
3826 size_t ct = 0;
3827 if (Acol2Brow(i) != LO_INVALID)
3828 ct = Bk.graph.row_map(Acol2Brow(i) + 1) - Bk.graph.row_map(Acol2Brow(i));
3829 else
3830 ct = Iks.graph.row_map(Acol2Irow(i) + 1) - Iks.graph.row_map(Acol2Irow(i));
3831 update += ct;
3832
3833 if (final && i + 1 == merge_numrows)
3834 Mrowptr(i + 1) = update;
3835 });
3836
3837 // Allocate nnz
3838 size_t merge_nnz = ::Tpetra::Details::getEntryOnHost(Mrowptr, merge_numrows);
3839 lno_nnz_view_t Mcolind(Kokkos::ViewAllocateWithoutInitializing("Mcolind"), merge_nnz);
3840 scalar_view_t Mvalues(Kokkos::ViewAllocateWithoutInitializing("Mvals"), merge_nnz);
3841
3842 // Use a Kokkos::parallel_for to fill the rowptr/colind arrays
3843 typedef Kokkos::RangePolicy<execution_space, size_t> range_type;
3844 Kokkos::parallel_for(
3845 "Tpetra_MatrixMatrix_merg_matrices_buildColindValues", range_type(0, merge_numrows), KOKKOS_LAMBDA(const size_t i) {
3846 if (Acol2Brow(i) != LO_INVALID) {
3847 size_t row = Acol2Brow(i);
3848 size_t start = Bk.graph.row_map(row);
3849 for (size_t j = Mrowptr(i); j < Mrowptr(i + 1); j++) {
3850 Mvalues(j) = Bk.values(j - Mrowptr(i) + start);
3851 Mcolind(j) = Bcol2Ccol(Bk.graph.entries(j - Mrowptr(i) + start));
3852 }
3853 } else {
3854 size_t row = Acol2Irow(i);
3855 size_t start = Iks.graph.row_map(row);
3856 for (size_t j = Mrowptr(i); j < Mrowptr(i + 1); j++) {
3857 Mvalues(j) = Iks.values(j - Mrowptr(i) + start);
3858 Mcolind(j) = Icol2Ccol(Iks.graph.entries(j - Mrowptr(i) + start));
3859 }
3860 }
3861 });
3862
3863 KCRS newmat("CrsMatrix", merge_numrows, mergedNodeNumCols, merge_nnz, Mvalues, Mrowptr, Mcolind);
3864 return newmat;
3865 } else {
3866 // We don't have a Bimport (the easy case)
3867 return Bk;
3868 }
3869} // end merge_matrices
3870
3871/*********************************************************************************************************/
3872// This only merges matrices that look like B & Bimport, aka, they have no overlapping rows
3873template <class Scalar, class LocalOrdinal, class GlobalOrdinal, class Node, class LocalOrdinalViewType>
3874const typename Tpetra::BlockCrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>::local_matrix_device_type
3875merge_matrices(BlockCrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Aview,
3876 BlockCrsMatrixStruct<Scalar, LocalOrdinal, GlobalOrdinal, Node>& Bview,
3877 const LocalOrdinalViewType& Acol2Brow,
3878 const LocalOrdinalViewType& Acol2Irow,
3879 const LocalOrdinalViewType& Bcol2Ccol,
3880 const LocalOrdinalViewType& Icol2Ccol,
3881 const size_t mergedNodeNumCols) {
3882 using Teuchos::RCP;
3883 typedef typename Tpetra::BlockCrsMatrix<Scalar, LocalOrdinal, GlobalOrdinal, Node>::local_matrix_device_type KBCRS;
3884 typedef typename KBCRS::StaticCrsGraphType graph_t;
3885 typedef typename graph_t::row_map_type::non_const_type lno_view_t;
3886 typedef typename graph_t::entries_type::non_const_type lno_nnz_view_t;
3887 typedef typename KBCRS::values_type::non_const_type scalar_view_t;
3888
3889 // Grab the KokkosSparse::BsrMatrix
3890 const KBCRS Ak = Aview.origMatrix->getLocalMatrixDevice();
3891 const KBCRS Bk = Bview.origMatrix->getLocalMatrixDevice();
3892
3893 // We need to do this dance if either (a) We have Bimport or (b) A's colMap is not the same as B's rowMap
3894 if (!Bview.importMatrix.is_null() ||
3895 (Bview.importMatrix.is_null() &&
3896 (&*Aview.origMatrix->getGraph()->getColMap() != &*Bview.origMatrix->getGraph()->getRowMap()))) {
3897 // We do have a Bimport
3898 // NOTE: We're going merge Borig and Bimport into a single matrix and reindex the columns *before* we multiply.
3899 // This option was chosen because we know we don't have any duplicate entries, so we can allocate once.
3900 KBCRS Iks;
3901 if (!Bview.importMatrix.is_null()) Iks = Bview.importMatrix->getLocalMatrixDevice();
3902 size_t merge_numrows = Ak.numCols();
3903
3904 // The last entry of this at least, need to be initialized
3905 lno_view_t Mrowptr("Mrowptr", merge_numrows + 1);
3906
3907 const LocalOrdinal LO_INVALID = Teuchos::OrdinalTraits<LocalOrdinal>::invalid();
3908
3909 // Use a Kokkos::parallel_scan to build the rowptr
3910 typedef typename Node::execution_space execution_space;
3911 typedef Kokkos::RangePolicy<execution_space, size_t> range_type;
3912 Kokkos::parallel_scan(
3913 "Tpetra_MatrixMatrix_merge_matrices_buildRowptr", range_type(0, merge_numrows),
3914 KOKKOS_LAMBDA(const size_t i, size_t& update, const bool final) {
3915 if (final) Mrowptr(i) = update;
3916 // Get the row count
3917 size_t ct = 0;
3918 if (Acol2Brow(i) != LO_INVALID)
3919 ct = Bk.graph.row_map(Acol2Brow(i) + 1) - Bk.graph.row_map(Acol2Brow(i));
3920 else
3921 ct = Iks.graph.row_map(Acol2Irow(i) + 1) - Iks.graph.row_map(Acol2Irow(i));
3922 update += ct;
3923
3924 if (final && i + 1 == merge_numrows)
3925 Mrowptr(i + 1) = update;
3926 });
3927
3928 // Allocate nnz
3929 size_t merge_nnz = ::Tpetra::Details::getEntryOnHost(Mrowptr, merge_numrows);
3930 const int blocksize = Ak.blockDim();
3931 lno_nnz_view_t Mcolind(Kokkos::ViewAllocateWithoutInitializing("Mcolind"), merge_nnz);
3932 scalar_view_t Mvalues(Kokkos::ViewAllocateWithoutInitializing("Mvals"), merge_nnz * blocksize * blocksize);
3933
3934 // Use a Kokkos::parallel_for to fill the rowptr/colind arrays
3935 typedef Kokkos::RangePolicy<execution_space, size_t> range_type;
3936 Kokkos::parallel_for(
3937 "Tpetra_MatrixMatrix_merg_matrices_buildColindValues", range_type(0, merge_numrows), KOKKOS_LAMBDA(const size_t i) {
3938 if (Acol2Brow(i) != LO_INVALID) {
3939 size_t row = Acol2Brow(i);
3940 size_t start = Bk.graph.row_map(row);
3941 for (size_t j = Mrowptr(i); j < Mrowptr(i + 1); j++) {
3942 Mcolind(j) = Bcol2Ccol(Bk.graph.entries(j - Mrowptr(i) + start));
3943
3944 for (int b = 0; b < blocksize * blocksize; ++b) {
3945 const int val_indx = j * blocksize * blocksize + b;
3946 const int b_val_indx = (j - Mrowptr(i) + start) * blocksize * blocksize + b;
3947 Mvalues(val_indx) = Bk.values(b_val_indx);
3948 }
3949 }
3950 } else {
3951 size_t row = Acol2Irow(i);
3952 size_t start = Iks.graph.row_map(row);
3953 for (size_t j = Mrowptr(i); j < Mrowptr(i + 1); j++) {
3954 Mcolind(j) = Icol2Ccol(Iks.graph.entries(j - Mrowptr(i) + start));
3955
3956 for (int b = 0; b < blocksize * blocksize; ++b) {
3957 const int val_indx = j * blocksize * blocksize + b;
3958 const int b_val_indx = (j - Mrowptr(i) + start) * blocksize * blocksize + b;
3959 Mvalues(val_indx) = Iks.values(b_val_indx);
3960 }
3961 }
3962 }
3963 });
3964
3965 // Build and return merged KokkosSparse matrix
3966 KBCRS newmat("CrsMatrix", merge_numrows, mergedNodeNumCols, merge_nnz, Mvalues, Mrowptr, Mcolind, blocksize);
3967 return newmat;
3968 } else {
3969 // We don't have a Bimport (the easy case)
3970 return Bk;
3971 }
3972} // end merge_matrices
3973
3974/*********************************************************************************************************/
3975template <typename SC, typename LO, typename GO, typename NO>
3976void AddKernels<SC, LO, GO, NO>::
3977 addSorted(
3978 const typename MMdetails::AddKernels<SC, LO, GO, NO>::values_array& Avals,
3979 const typename MMdetails::AddKernels<SC, LO, GO, NO>::row_ptrs_array_const& Arowptrs,
3980 const typename MMdetails::AddKernels<SC, LO, GO, NO>::col_inds_array& Acolinds,
3981 const typename MMdetails::AddKernels<SC, LO, GO, NO>::impl_scalar_type scalarA,
3982 const typename MMdetails::AddKernels<SC, LO, GO, NO>::values_array& Bvals,
3983 const typename MMdetails::AddKernels<SC, LO, GO, NO>::row_ptrs_array_const& Browptrs,
3984 const typename MMdetails::AddKernels<SC, LO, GO, NO>::col_inds_array& Bcolinds,
3985 const typename MMdetails::AddKernels<SC, LO, GO, NO>::impl_scalar_type scalarB,
3986 GO numGlobalCols,
3987 typename MMdetails::AddKernels<SC, LO, GO, NO>::values_array& Cvals,
3988 typename MMdetails::AddKernels<SC, LO, GO, NO>::row_ptrs_array& Crowptrs,
3989 typename MMdetails::AddKernels<SC, LO, GO, NO>::col_inds_array& Ccolinds) {
3990 using Teuchos::rcp;
3991 using Teuchos::TimeMonitor;
3992 using AddKern = MMdetails::AddKernels<SC, LO, GO, NO>;
3993 TEUCHOS_TEST_FOR_EXCEPTION(Arowptrs.extent(0) != Browptrs.extent(0), std::runtime_error, "Can't add matrices with different numbers of rows.");
3994 auto nrows = Arowptrs.extent(0) - 1;
3995 Crowptrs = row_ptrs_array(Kokkos::ViewAllocateWithoutInitializing("C row ptrs"), nrows + 1);
3996 typename AddKern::KKH handle;
3997 handle.create_spadd_handle(true);
3998 auto addHandle = handle.get_spadd_handle();
3999
4000 auto MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: add: sorted symbolic"));
4001
4002 KokkosSparse::spadd_symbolic(&handle,
4003 nrows, numGlobalCols,
4004 Arowptrs, Acolinds, Browptrs, Bcolinds, Crowptrs);
4005 // KokkosKernels requires values to be zeroed
4006 Cvals = values_array("C values", addHandle->get_c_nnz());
4007 Ccolinds = col_inds_array(Kokkos::ViewAllocateWithoutInitializing("C colinds"), addHandle->get_c_nnz());
4008
4009 MM = Teuchos::null;
4010 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: add: sorted numeric"));
4011 KokkosSparse::spadd_numeric(&handle,
4012 nrows, numGlobalCols,
4013 Arowptrs, Acolinds, Avals, scalarA,
4014 Browptrs, Bcolinds, Bvals, scalarB,
4015 Crowptrs, Ccolinds, Cvals);
4016}
4017
4018template <typename SC, typename LO, typename GO, typename NO>
4019void AddKernels<SC, LO, GO, NO>::
4020 addUnsorted(
4021 const typename MMdetails::AddKernels<SC, LO, GO, NO>::values_array& Avals,
4022 const typename MMdetails::AddKernels<SC, LO, GO, NO>::row_ptrs_array_const& Arowptrs,
4023 const typename MMdetails::AddKernels<SC, LO, GO, NO>::col_inds_array& Acolinds,
4024 const typename MMdetails::AddKernels<SC, LO, GO, NO>::impl_scalar_type scalarA,
4025 const typename MMdetails::AddKernels<SC, LO, GO, NO>::values_array& Bvals,
4026 const typename MMdetails::AddKernels<SC, LO, GO, NO>::row_ptrs_array_const& Browptrs,
4027 const typename MMdetails::AddKernels<SC, LO, GO, NO>::col_inds_array& Bcolinds,
4028 const typename MMdetails::AddKernels<SC, LO, GO, NO>::impl_scalar_type scalarB,
4029 GO numGlobalCols,
4030 typename MMdetails::AddKernels<SC, LO, GO, NO>::values_array& Cvals,
4031 typename MMdetails::AddKernels<SC, LO, GO, NO>::row_ptrs_array& Crowptrs,
4032 typename MMdetails::AddKernels<SC, LO, GO, NO>::col_inds_array& Ccolinds) {
4033 using Teuchos::rcp;
4034 using Teuchos::TimeMonitor;
4035 using AddKern = MMdetails::AddKernels<SC, LO, GO, NO>;
4036 TEUCHOS_TEST_FOR_EXCEPTION(Arowptrs.extent(0) != Browptrs.extent(0), std::runtime_error, "Can't add matrices with different numbers of rows.");
4037 auto nrows = Arowptrs.extent(0) - 1;
4038 Crowptrs = row_ptrs_array(Kokkos::ViewAllocateWithoutInitializing("C row ptrs"), nrows + 1);
4039 typedef MMdetails::AddKernels<SC, LO, GO, NO> AddKern;
4040 typename AddKern::KKH handle;
4041 handle.create_spadd_handle(false);
4042 auto addHandle = handle.get_spadd_handle();
4043 auto MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: add: unsorted symbolic"));
4044
4045 KokkosSparse::spadd_symbolic(&handle,
4046 nrows, numGlobalCols,
4047 Arowptrs, Acolinds, Browptrs, Bcolinds, Crowptrs);
4048 // Cvals must be zeroed out
4049 Cvals = values_array("C values", addHandle->get_c_nnz());
4050 Ccolinds = col_inds_array(Kokkos::ViewAllocateWithoutInitializing("C colinds"), addHandle->get_c_nnz());
4051 MM = Teuchos::null;
4052 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: add: unsorted numeric"));
4053 KokkosSparse::spadd_numeric(&handle,
4054 nrows, numGlobalCols,
4055 Arowptrs, Acolinds, Avals, scalarA,
4056 Browptrs, Bcolinds, Bvals, scalarB,
4057 Crowptrs, Ccolinds, Cvals);
4058}
4059
4060template <typename GO,
4061 typename LocalIndicesType,
4062 typename GlobalIndicesType,
4063 typename ColMapType>
4064struct ConvertLocalToGlobalFunctor {
4065 ConvertLocalToGlobalFunctor(
4066 const LocalIndicesType& colindsOrig_,
4067 const GlobalIndicesType& colindsConverted_,
4068 const ColMapType& colmap_)
4069 : colindsOrig(colindsOrig_)
4070 , colindsConverted(colindsConverted_)
4071 , colmap(colmap_) {}
4072 KOKKOS_INLINE_FUNCTION void
4073 operator()(const GO i) const {
4074 colindsConverted(i) = colmap.getGlobalElement(colindsOrig(i));
4075 }
4076 LocalIndicesType colindsOrig;
4077 GlobalIndicesType colindsConverted;
4078 ColMapType colmap;
4079};
4080
4081template <typename SC, typename LO, typename GO, typename NO>
4082void MMdetails::AddKernels<SC, LO, GO, NO>::
4083 convertToGlobalAndAdd(
4084 const typename MMdetails::AddKernels<SC, LO, GO, NO>::KCRS A,
4085 const typename MMdetails::AddKernels<SC, LO, GO, NO>::impl_scalar_type scalarA,
4086 const typename MMdetails::AddKernels<SC, LO, GO, NO>::KCRS B,
4087 const typename MMdetails::AddKernels<SC, LO, GO, NO>::impl_scalar_type scalarB,
4088 const typename MMdetails::AddKernels<SC, LO, GO, NO>::local_map_type& AcolMap,
4089 const typename MMdetails::AddKernels<SC, LO, GO, NO>::local_map_type& BcolMap,
4090 typename MMdetails::AddKernels<SC, LO, GO, NO>::values_array& Cvals,
4091 typename MMdetails::AddKernels<SC, LO, GO, NO>::row_ptrs_array& Crowptrs,
4092 typename MMdetails::AddKernels<SC, LO, GO, NO>::global_col_inds_array& Ccolinds) {
4093 using Teuchos::rcp;
4094 using Teuchos::TimeMonitor;
4095 // Need to use a different KokkosKernelsHandle type than other versions,
4096 // since the ordinals are now GO
4097 using KKH_GO = KokkosKernels::Experimental::KokkosKernelsHandle<size_t, GO, impl_scalar_type,
4098 typename NO::execution_space, typename NO::memory_space, typename NO::memory_space>;
4099
4100 const values_array Avals = A.values;
4101 const values_array Bvals = B.values;
4102 const col_inds_array Acolinds = A.graph.entries;
4103 const col_inds_array Bcolinds = B.graph.entries;
4104 auto Arowptrs = A.graph.row_map;
4105 auto Browptrs = B.graph.row_map;
4106 global_col_inds_array AcolindsConverted(Kokkos::ViewAllocateWithoutInitializing("A colinds (converted)"), Acolinds.extent(0));
4107 global_col_inds_array BcolindsConverted(Kokkos::ViewAllocateWithoutInitializing("B colinds (converted)"), Bcolinds.extent(0));
4108
4109 auto MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: add: diff col map kernel: column map conversion"));
4110
4111 ConvertLocalToGlobalFunctor<GO, col_inds_array, global_col_inds_array, local_map_type> convertA(Acolinds, AcolindsConverted, AcolMap);
4112 Kokkos::parallel_for("Tpetra_MatrixMatrix_convertColIndsA", range_type(0, Acolinds.extent(0)), convertA);
4113 ConvertLocalToGlobalFunctor<GO, col_inds_array, global_col_inds_array, local_map_type> convertB(Bcolinds, BcolindsConverted, BcolMap);
4114 Kokkos::parallel_for("Tpetra_MatrixMatrix_convertColIndsB", range_type(0, Bcolinds.extent(0)), convertB);
4115 KKH_GO handle;
4116 handle.create_spadd_handle(false);
4117 auto addHandle = handle.get_spadd_handle();
4118 MM = Teuchos::null;
4119 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: add: diff col map kernel: unsorted symbolic"));
4120 auto nrows = Arowptrs.extent(0) - 1;
4121 Crowptrs = row_ptrs_array(Kokkos::ViewAllocateWithoutInitializing("C row ptrs"), nrows + 1);
4122 KokkosSparse::spadd_symbolic(&handle,
4123 nrows, A.numCols(),
4124 Arowptrs, AcolindsConverted, Browptrs, BcolindsConverted, Crowptrs);
4125 Cvals = values_array("C values", addHandle->get_c_nnz());
4126 Ccolinds = global_col_inds_array(Kokkos::ViewAllocateWithoutInitializing("C colinds"), addHandle->get_c_nnz());
4127
4128 MM = Teuchos::null;
4129 MM = rcp(new Tpetra::Details::ProfilingRegion("TpetraExt: add: diff col map kernel: unsorted numeric"));
4130 KokkosSparse::spadd_numeric(&handle,
4131 nrows, A.numCols(),
4132 Arowptrs, AcolindsConverted, Avals, scalarA,
4133 Browptrs, BcolindsConverted, Bvals, scalarB,
4134 Crowptrs, Ccolinds, Cvals);
4135}
4136
4137} // namespace MMdetails
4138
4139} // End namespace Tpetra
4140
4141/*********************************************************************************************************/
4142//
4143// Explicit instantiation macro
4144//
4145// Must be expanded from within the Tpetra namespace!
4146//
4147namespace Tpetra {
4148
4149#define TPETRA_MATRIXMATRIX_INSTANT(SCALAR, LO, GO, NODE) \
4150 template void MatrixMatrix::Multiply( \
4151 const CrsMatrix<SCALAR, LO, GO, NODE>& A, \
4152 bool transposeA, \
4153 const CrsMatrix<SCALAR, LO, GO, NODE>& B, \
4154 bool transposeB, \
4155 CrsMatrix<SCALAR, LO, GO, NODE>& C, \
4156 bool call_FillComplete_on_result, \
4157 const std::string& label, \
4158 const Teuchos::RCP<Teuchos::ParameterList>& params); \
4159 \
4160 template void MatrixMatrix::Multiply( \
4161 const Teuchos::RCP<const BlockCrsMatrix<SCALAR, LO, GO, NODE>>& A, \
4162 bool transposeA, \
4163 const Teuchos::RCP<const BlockCrsMatrix<SCALAR, LO, GO, NODE>>& B, \
4164 bool transposeB, \
4165 Teuchos::RCP<BlockCrsMatrix<SCALAR, LO, GO, NODE>>& C, \
4166 const std::string& label); \
4167 \
4168 template void MatrixMatrix::Jacobi( \
4169 typename Teuchos::ScalarTraits<SCALAR>::magnitudeType omega, \
4170 const Vector<SCALAR, LO, GO, NODE>& Dinv, \
4171 const CrsMatrix<SCALAR, LO, GO, NODE>& A, \
4172 const CrsMatrix<SCALAR, LO, GO, NODE>& B, \
4173 CrsMatrix<SCALAR, LO, GO, NODE>& C, \
4174 bool call_FillComplete_on_result, \
4175 const std::string& label, \
4176 const Teuchos::RCP<Teuchos::ParameterList>& params); \
4177 \
4178 template void MatrixMatrix::Add( \
4179 const CrsMatrix<SCALAR, LO, GO, NODE>& A, \
4180 bool transposeA, \
4181 SCALAR scalarA, \
4182 const CrsMatrix<SCALAR, LO, GO, NODE>& B, \
4183 bool transposeB, \
4184 SCALAR scalarB, \
4185 Teuchos::RCP<CrsMatrix<SCALAR, LO, GO, NODE>>& C); \
4186 \
4187 template void MatrixMatrix::Add( \
4188 const CrsMatrix<SCALAR, LO, GO, NODE>& A, \
4189 bool transposeA, \
4190 SCALAR scalarA, \
4191 const CrsMatrix<SCALAR, LO, GO, NODE>& B, \
4192 bool transposeB, \
4193 SCALAR scalarB, \
4194 const Teuchos::RCP<CrsMatrix<SCALAR, LO, GO, NODE>>& C); \
4195 \
4196 template void MatrixMatrix::Add( \
4197 const CrsMatrix<SCALAR, LO, GO, NODE>& A, \
4198 bool transposeA, \
4199 SCALAR scalarA, \
4200 CrsMatrix<SCALAR, LO, GO, NODE>& B, \
4201 SCALAR scalarB); \
4202 \
4203 template Teuchos::RCP<CrsMatrix<SCALAR, LO, GO, NODE>> \
4204 MatrixMatrix::add<SCALAR, LO, GO, NODE>(const SCALAR& alpha, \
4205 const bool transposeA, \
4206 const CrsMatrix<SCALAR, LO, GO, NODE>& A, \
4207 const SCALAR& beta, \
4208 const bool transposeB, \
4209 const CrsMatrix<SCALAR, LO, GO, NODE>& B, \
4210 const Teuchos::RCP<const Map<LO, GO, NODE>>& domainMap, \
4211 const Teuchos::RCP<const Map<LO, GO, NODE>>& rangeMap, \
4212 const Teuchos::RCP<Teuchos::ParameterList>& params); \
4213 \
4214 template void \
4215 MatrixMatrix::add<SCALAR, LO, GO, NODE>(const SCALAR& alpha, \
4216 const bool transposeA, \
4217 const CrsMatrix<SCALAR, LO, GO, NODE>& A, \
4218 const SCALAR& beta, \
4219 const bool transposeB, \
4220 const CrsMatrix<SCALAR, LO, GO, NODE>& B, \
4221 CrsMatrix<SCALAR, LO, GO, NODE>& C, \
4222 const Teuchos::RCP<const Map<LO, GO, NODE>>& domainMap, \
4223 const Teuchos::RCP<const Map<LO, GO, NODE>>& rangeMap, \
4224 const Teuchos::RCP<Teuchos::ParameterList>& params); \
4225 \
4226 template struct MMdetails::AddKernels<SCALAR, LO, GO, NODE>; \
4227 \
4228 template void MMdetails::import_and_extract_views<SCALAR, LO, GO, NODE>(const CrsMatrix<SCALAR, LO, GO, NODE>& M, \
4229 Teuchos::RCP<const Map<LO, GO, NODE>> targetMap, \
4230 CrsMatrixStruct<SCALAR, LO, GO, NODE>& Mview, \
4231 Teuchos::RCP<const Import<LO, GO, NODE>> prototypeImporter, \
4232 bool userAssertsThereAreNoRemotes, \
4233 const std::string& label, \
4234 const Teuchos::RCP<Teuchos::ParameterList>& params); \
4235 \
4236 template void MMdetails::import_and_extract_views<SCALAR, LO, GO, NODE>(const BlockCrsMatrix<SCALAR, LO, GO, NODE>& M, \
4237 Teuchos::RCP<const Map<LO, GO, NODE>> targetMap, \
4238 BlockCrsMatrixStruct<SCALAR, LO, GO, NODE>& Mview, \
4239 Teuchos::RCP<const Import<LO, GO, NODE>> prototypeImporter, \
4240 bool userAssertsThereAreNoRemotes);
4241} // End namespace Tpetra
4242
4243#endif // TPETRA_MATRIXMATRIX_DEF_HPP
Declaration of Tpetra::Details::Behavior, a class that describes Tpetra's behavior.
Declaration of Tpetra::Details::Profiling, a scope guard for Kokkos Profiling.
Declare and define the functions Tpetra::Details::computeOffsetsFromCounts and Tpetra::computeOffsets...
Declaration and definition of Tpetra::Details::getEntryOnHost.
Utility functions for packing and unpacking sparse matrix entries.
void sortCrsEntries(const Teuchos::ArrayView< size_t > &CRS_rowptr, const Teuchos::ArrayView< Ordinal > &CRS_colind, const Teuchos::ArrayView< Scalar > &CRS_vals)
Sort the entries of the (raw CSR) matrix by column index within each row.
Internal functions and macros designed for use with Tpetra::Import and Tpetra::Export objects.
Stand-alone utility functions and macros.
Forward declaration of some Tpetra Matrix Matrix objects.
Sparse matrix whose entries are small dense square blocks, all of the same dimensions.
Struct that holds views of the contents of a BlockCrsMatrix.
A distributed graph accessed by rows (adjacency lists) and stored sparsely.
Sparse matrix that presents a row-oriented interface that lets users read or modify entries.
bool isGloballyIndexed() const override
Whether the matrix is globally indexed on the calling process.
bool haveGlobalConstants() const
Returns true if globalConstants have been computed; false otherwise.
void scale(const Scalar &alpha)
Scale the matrix's values: this := alpha*this.
Teuchos::RCP< const map_type > getDomainMap() const override
The domain Map of this matrix.
void expertStaticFillComplete(const Teuchos::RCP< const map_type > &domainMap, const Teuchos::RCP< const map_type > &rangeMap, const Teuchos::RCP< const import_type > &importer=Teuchos::null, const Teuchos::RCP< const export_type > &exporter=Teuchos::null, const Teuchos::RCP< Teuchos::ParameterList > &params=Teuchos::null)
Perform a fillComplete on a matrix that already has data.
Teuchos::RCP< const Teuchos::Comm< int > > getComm() const override
The communicator over which the matrix is distributed.
bool isFillActive() const
Whether the matrix is not fill complete.
Teuchos::RCP< const map_type > getRangeMap() const override
The range Map of this matrix.
size_t getLocalMaxNumRowEntries() const override
Maximum number of entries in any row of the matrix, on this process.
void replaceColMap(const Teuchos::RCP< const map_type > &newColMap)
Replace the matrix's column Map with the given Map.
global_size_t getGlobalNumRows() const override
Number of global elements in the row map of this matrix.
void insertGlobalValues(const GlobalOrdinal globalRow, const Teuchos::ArrayView< const GlobalOrdinal > &cols, const Teuchos::ArrayView< const Scalar > &vals)
Insert one or more entries into the matrix, using global column indices.
void fillComplete(const Teuchos::RCP< const map_type > &domainMap, const Teuchos::RCP< const map_type > &rangeMap, const Teuchos::RCP< Teuchos::ParameterList > &params=Teuchos::null)
Tell the matrix that you are done changing its structure or values, and that you are ready to do comp...
void setAllValues(const typename local_graph_device_type::row_map_type &ptr, const typename local_graph_device_type::entries_type::non_const_type &ind, const typename local_matrix_device_type::values_type &val)
Set the local matrix using three (compressed sparse row) arrays.
Teuchos::RCP< const RowGraph< LocalOrdinal, GlobalOrdinal, Node > > getGraph() const override
This matrix's graph, as a RowGraph.
bool isStaticGraph() const
Indicates that the graph is static, so that new entries cannot be added to this matrix.
KokkosSparse::CrsMatrix< impl_scalar_type, local_ordinal_type, device_type, void, typename local_graph_device_type::size_type > local_matrix_device_type
The specialization of Kokkos::CrsMatrix that represents the part of the sparse matrix on each MPI pro...
size_t getLocalNumRows() const override
The number of matrix rows owned by the calling process.
bool isFillComplete() const override
Whether the matrix is fill complete.
Teuchos::RCP< const map_type > getRowMap() const override
The Map that describes the row distribution in this matrix.
bool isLocallyIndexed() const override
Whether the matrix is locally indexed on the calling process.
void resumeFill(const Teuchos::RCP< Teuchos::ParameterList > &params=Teuchos::null)
Resume operations that may change the values or structure of the matrix.
Struct that holds views of the contents of a CrsMatrix.
static bool debug()
Whether Tpetra is in debug mode.
virtual Teuchos::RCP< const map_type > getMap() const
The Map describing the parallel distribution of this object.
Communication plan for data redistribution from a (possibly) multiply-owned to a uniquely-owned distr...
Communication plan for data redistribution from a uniquely-owned to a (possibly) multiply-owned distr...
A parallel distribution of indices over processes.
static void writeSparseFile(const std::string &filename, const sparse_matrix_type &matrix, const std::string &matrixName, const std::string &matrixDescription, const bool debug=false)
Construct and (optionally) redistribute the explicitly stored transpose of a CrsMatrix.
Teuchos::RCP< crs_matrix_type > createTranspose(const Teuchos::RCP< Teuchos::ParameterList > &params=Teuchos::null)
Compute and return the transpose of the matrix given to the constructor.
A distributed dense vector.
void start()
Start the deep_copy counter.
int makeColMap(Teuchos::RCP< const Tpetra::Map< LO, GO, NT > > &colMap, Teuchos::Array< int > &remotePIDs, const Teuchos::RCP< const Tpetra::Map< LO, GO, NT > > &domMap, const CrsGraph< LO, GO, NT > &graph, const bool sortEachProcsGids=true, std::ostream *errStrm=NULL)
Make the graph's column Map.
Distributed sparse matrix-matrix multiply and add.
void Jacobi(typename Teuchos::ScalarTraits< Scalar >::magnitudeType omega, const Vector< Scalar, LocalOrdinal, GlobalOrdinal, Node > &Dinv, const CrsMatrix< Scalar, LocalOrdinal, GlobalOrdinal, Node > &A, const CrsMatrix< Scalar, LocalOrdinal, GlobalOrdinal, Node > &B, CrsMatrix< Scalar, LocalOrdinal, GlobalOrdinal, Node > &C, bool call_FillComplete_on_result=true, const std::string &label=std::string(), const Teuchos::RCP< Teuchos::ParameterList > &params=Teuchos::null)
Teuchos::RCP< CrsMatrix< Scalar, LocalOrdinal, GlobalOrdinal, Node > > add(const Scalar &alpha, const bool transposeA, const CrsMatrix< Scalar, LocalOrdinal, GlobalOrdinal, Node > &A, const Scalar &beta, const bool transposeB, const CrsMatrix< Scalar, LocalOrdinal, GlobalOrdinal, Node > &B, const Teuchos::RCP< const Map< LocalOrdinal, GlobalOrdinal, Node > > &domainMap=Teuchos::null, const Teuchos::RCP< const Map< LocalOrdinal, GlobalOrdinal, Node > > &rangeMap=Teuchos::null, const Teuchos::RCP< Teuchos::ParameterList > &params=Teuchos::null)
Compute the sparse matrix sum C = scalarA * Op(A) + scalarB * Op(B), where Op(X) is either X or its t...
void Multiply(const CrsMatrix< Scalar, LocalOrdinal, GlobalOrdinal, Node > &A, bool transposeA, const CrsMatrix< Scalar, LocalOrdinal, GlobalOrdinal, Node > &B, bool transposeB, CrsMatrix< Scalar, LocalOrdinal, GlobalOrdinal, Node > &C, bool call_FillComplete_on_result=true, const std::string &label=std::string(), const Teuchos::RCP< Teuchos::ParameterList > &params=Teuchos::null)
Sparse matrix-matrix multiply.
void Add(const CrsMatrix< Scalar, LocalOrdinal, GlobalOrdinal, Node > &A, bool transposeA, Scalar scalarA, CrsMatrix< Scalar, LocalOrdinal, GlobalOrdinal, Node > &B, Scalar scalarB)
Namespace Tpetra contains the class and methods constituting the Tpetra library.
void removeCrsMatrixZeros(CrsMatrixType &matrix, typename Teuchos::ScalarTraits< typename CrsMatrixType::scalar_type >::magnitudeType const &threshold=Teuchos::ScalarTraits< typename CrsMatrixType::scalar_type >::magnitude(Teuchos::ScalarTraits< typename CrsMatrixType::scalar_type >::zero()))
Remove zero entries from a matrix.
size_t global_size_t
Global size_t object.
Teuchos::RCP< CrsMatrixType > importAndFillCompleteCrsMatrix(const Teuchos::RCP< const CrsMatrixType > &sourceMatrix, const Import< typename CrsMatrixType::local_ordinal_type, typename CrsMatrixType::global_ordinal_type, typename CrsMatrixType::node_type > &importer, const Teuchos::RCP< const Map< typename CrsMatrixType::local_ordinal_type, typename CrsMatrixType::global_ordinal_type, typename CrsMatrixType::node_type > > &domainMap=Teuchos::null, const Teuchos::RCP< const Map< typename CrsMatrixType::local_ordinal_type, typename CrsMatrixType::global_ordinal_type, typename CrsMatrixType::node_type > > &rangeMap=Teuchos::null, const Teuchos::RCP< Teuchos::ParameterList > &params=Teuchos::null)
Nonmember CrsMatrix constructor that fuses Import and fillComplete().