42 #ifndef __TPETRA_DISTRIBUTORLOWERTRIANGULARBLOCK_HPP
43 #define __TPETRA_DISTRIBUTORLOWERTRIANGULARBLOCK_HPP
46 #include "Tpetra_Distributor.hpp"
50 #include "Tpetra_Map.hpp"
51 #include "Tpetra_Operator.hpp"
52 #include "Tpetra_Vector.hpp"
53 #include "Tpetra_CrsMatrix.hpp"
59 template <
typename gno_t,
typename scalar_t>
60 class DistributionLowerTriangularBlock :
public Distribution<gno_t,scalar_t> {
142 using Distribution<gno_t,scalar_t>::me;
143 using Distribution<gno_t,scalar_t>::np;
144 using Distribution<gno_t,scalar_t>::comm;
145 using Distribution<gno_t,scalar_t>::nrows;
146 using typename Distribution<gno_t,scalar_t>::NZindex_t;
147 using typename Distribution<gno_t,scalar_t>::LocalNZmap_t;
152 DistributionLowerTriangularBlock(
size_t nrows_,
153 const Teuchos::RCP<
const Teuchos::Comm<int> > &comm_,
154 const Teuchos::ParameterList ¶ms) :
155 Distribution<gno_t,scalar_t>(nrows_, comm_, params),
156 initialDist(nrows_, comm_, params),
157 sortByDegree(false), permMatrix(Teuchos::null),
158 redistributed(false), chunksComputed(false), nChunks(0)
161 nChunks = int(std::sqrt(
float(npnp)));
162 while (nChunks * (nChunks + 1) < npnp) nChunks++;
164 TEUCHOS_TEST_FOR_EXCEPTION(nChunks * (nChunks+1) != npnp, std::logic_error,
165 "Number of processors np = " << np <<
166 " must satisfy np = q(q+1)/2 for some q" <<
167 " for LowerTriangularBlock distribution");
168 nChunksPerRow = double(nChunks) / double(nrows);
170 const Teuchos::ParameterEntry *pe = params.getEntryPtr(
"sortByDegree");
171 if (pe != NULL) sortByDegree = pe->getValue<
bool>(&sortByDegree);
173 pe = params.getEntryPtr(
"readPerProcess");
174 if (pe != NULL) redistributed = pe->getValue<
bool>(&redistributed);
176 if (me == 0) std::cout <<
"\n LowerTriangularBlock Distribution: "
178 <<
"\n nChunks = " << nChunks
182 enum DistributionType DistType() {
return LowerTriangularBlock; }
184 bool areChunksComputed() {
return chunksComputed; }
186 Teuchos::Array<gno_t> getChunkCuts() {
190 throw std::runtime_error(
"Error: Requested chunk cuts have not been computed yet.");
197 inline bool VecMine(gno_t i) {
return initialDist.VecMine(i); }
202 bool Mine(gno_t i, gno_t j) {
204 if (j > i)
return false;
205 else return (procFromChunks(i,j) == me);
208 return initialDist.Mine(i,j);
211 inline bool Mine(gno_t i, gno_t j,
int p) {
return Mine(i,j);}
214 void Redistribute(LocalNZmap_t &localNZ)
220 gno_t myFirstRow = initialDist.getFirstRow(me);
221 gno_t nMyRows = initialDist.getNumRow(me);
222 Teuchos::Array<gno_t> nnzPerRow(nMyRows, 0);
224 Teuchos::Array<int> rcvcnt(np);
225 Teuchos::Array<int> disp(np);
226 for (
int sum = 0, p = 0; p < np; p++) {
227 int prows = initialDist.getNumRow(p);
238 Teuchos::Array<gno_t> permuteIndex;
239 Teuchos::Array<gno_t> sortedOrder;
241 Teuchos::Array<gno_t> globalRowBuf;
245 globalRowBuf.resize(nrows, 0);
250 for (
auto it = localNZ.begin(); it != localNZ.end(); it++) {
251 gno_t I = it->first.first;
252 nnzPerRow[I-myFirstRow]++;
255 Teuchos::gatherv<int,gno_t>(nnzPerRow.getRawPtr(), nMyRows,
256 globalRowBuf.getRawPtr(),
257 rcvcnt.getRawPtr(), disp.getRawPtr(),
260 permuteIndex.resize(nrows);
261 sortedOrder.resize(nrows);
265 for (
size_t i = 0 ; i != nrows; i++) sortedOrder[i] = i;
267 std::sort(sortedOrder.begin(), sortedOrder.end(),
268 [&](
const size_t& a,
const size_t& b) {
269 return (globalRowBuf[a] > globalRowBuf[b]);
274 for (
size_t i = 0; i < nrows; i++) {
275 permuteIndex[sortedOrder[i]] = i;
279 Teuchos::broadcast<int,gno_t>(*comm, 0, permuteIndex(0,nrows));
288 Teuchos::Array<gno_t> myRows;
289 for (
size_t i = 0; i < nrows; i++) {
290 if (VecMine(i)) myRows.push_back(i);
294 Teuchos::OrdinalTraits<Tpetra::global_size_t>::invalid();
295 Teuchos::RCP<const map_t> permMap =
296 rcp(
new map_t(dummy, myRows(), 0, comm));
298 permMatrix = rcp(
new matrix_t(permMap, 1));
300 Teuchos::Array<gno_t> cols(1);
301 Teuchos::Array<scalar_t> vals(1); vals[0] = 1.;
303 for (
size_t i = 0; i < permMap->getLocalNumElements(); i++) {
304 gno_t gid = permMap->getGlobalElement(i);
305 cols[0] = permuteIndex[gid];
306 permMatrix->insertGlobalValues(gid, cols(), vals());
309 permMatrix->fillComplete(permMap, permMap);
315 nnzPerRow.assign(nMyRows, 0);
317 for (
auto it = localNZ.begin(); it != localNZ.end(); it++) {
318 gno_t I = (sortByDegree ? permuteIndex[it->first.first]
320 gno_t J = (sortByDegree ? permuteIndex[it->first.second]
323 nnzPerRow[it->first.first - myFirstRow]++;
331 Teuchos::gatherv<int,gno_t>(nnzPerRow.getRawPtr(), nMyRows,
332 globalRowBuf.getRawPtr(),
333 rcvcnt.getRawPtr(), disp.getRawPtr(),
336 Teuchos::Array<int>().swap(rcvcnt);
337 Teuchos::Array<int>().swap(disp);
340 Teuchos::reduceAll(*comm, Teuchos::REDUCE_SUM, 1, &nnz, &gNnz);
342 chunkCuts.resize(nChunks+1, 0);
349 size_t target = gNnz / np;
350 size_t targetRunningTotal = 0;
351 size_t currentRunningTotal = 0;
353 for (
int chunkCnt = 0; chunkCnt < nChunks; chunkCnt++) {
354 targetRunningTotal = (target * (chunkCnt+1));
355 currentRunningTotal = 0;
356 while (I < static_cast<gno_t>(nrows)) {
357 size_t nextNnz = (sortByDegree ? globalRowBuf[sortedOrder[I]]
359 if (currentRunningTotal + nextNnz <= targetRunningTotal) {
360 currentRunningTotal += nextNnz;
366 chunkCuts[chunkCnt+1] = I;
368 chunkCuts[nChunks] =
static_cast<gno_t
>(nrows);
372 Teuchos::Array<gno_t>().swap(globalRowBuf);
374 Teuchos::broadcast<int,gno_t>(*comm, 0, chunkCuts(0,nChunks+1));
375 chunksComputed =
true;
378 Kokkos::View<gno_t*, Kokkos::HostSpace> iOut(
"iOut", localNZ.size());
379 Kokkos::View<gno_t*, Kokkos::HostSpace> jOut(
"jOut", localNZ.size());
380 Kokkos::View<scalar_t*, Kokkos::HostSpace> vOut(
"vOut", localNZ.size());
381 Teuchos::Array<int> pOut(localNZ.size());
384 for (
auto it = localNZ.begin(); it != localNZ.end(); it++) {
385 iOut[sendCnt] = (sortByDegree ? permuteIndex[it->first.first]
387 jOut[sendCnt] = (sortByDegree ? permuteIndex[it->first.second]
389 if (jOut[sendCnt] <= iOut[sendCnt]) {
390 vOut[sendCnt] = it->second;
391 pOut[sendCnt] = procFromChunks(iOut[sendCnt], jOut[sendCnt]);
398 LocalNZmap_t().swap(localNZ);
399 if (sortByDegree) Teuchos::Array<gno_t>().swap(permuteIndex);
403 size_t nrecvs = plan.createFromSends(pOut(0,sendCnt));
404 Kokkos::View<gno_t*, Kokkos::HostSpace> iIn(
"iIn", nrecvs);
405 Kokkos::View<gno_t*, Kokkos::HostSpace> jIn(
"jIn", nrecvs);
406 Kokkos::View<scalar_t*, Kokkos::HostSpace> vIn(
"vIn", nrecvs);
409 auto sendIndices = std::make_pair(static_cast<size_t>(0), sendCnt);
410 plan.doPostsAndWaits(Kokkos::subview(iOut, sendIndices), 1, iIn);
411 plan.doPostsAndWaits(Kokkos::subview(jOut, sendIndices), 1, jIn);
412 plan.doPostsAndWaits(Kokkos::subview(vOut, sendIndices), 1, vIn);
415 for (
size_t n = 0; n < nrecvs; n++) {
416 NZindex_t nz(iIn[n], jIn[n]);
417 localNZ[nz] = vIn[n];
420 redistributed =
true;
423 Teuchos::RCP<matrix_t> getPermutationMatrix()
const {
return permMatrix; }
427 Distribution1DLinear<gno_t,scalar_t> initialDist;
434 Teuchos::RCP<matrix_t> permMatrix;
448 double nChunksPerRow;
449 Teuchos::Array<gno_t> chunkCuts;
451 int findIdxInChunks(gno_t I) {
452 int m = I * nChunksPerRow;
453 while (I < chunkCuts[m]) m--;
454 while (I >= chunkCuts[m+1]) m++;
458 int procFromChunks(gno_t I, gno_t J) {
459 int m = findIdxInChunks(I);
460 int n = findIdxInChunks(J);
461 int p = m*(m+1)/2 + n;
470 template <
typename scalar_t,
471 class Node = ::Tpetra::Details::DefaultTypes::node_type>
472 class LowerTriangularBlockOperator :
474 Tpetra::Map<>::global_ordinal_type,
486 using dist_t = Tpetra::DistributionLowerTriangularBlock<gno_t, scalar_t>;
488 LowerTriangularBlockOperator(
489 const Teuchos::RCP<const matrix_t> &lowerTriangularMatrix_,
491 : lowerTriangularMatrix(lowerTriangularMatrix_),
492 permMatrix(dist.getPermutationMatrix())
496 TEUCHOS_TEST_FOR_EXCEPTION(
497 !lowerTriangularMatrix->getRangeMap()->isSameAs(
498 *lowerTriangularMatrix->getDomainMap()),
500 "The Domain and Range maps of the LowerTriangularBlock matrix "
505 vector_t diagByRowMap(lowerTriangularMatrix->getRowMap());
506 lowerTriangularMatrix->getLocalDiagCopy(diagByRowMap);
507 diag = Teuchos::rcp(
new vector_t(lowerTriangularMatrix->getRangeMap()));
509 lowerTriangularMatrix->getRangeMap());
510 diag->doExport(diagByRowMap, exporter,
Tpetra::ADD);
513 void apply(
const mvector_t &x, mvector_t &y, Teuchos::ETransp mode,
514 scalar_t alpha, scalar_t beta)
const
516 scalar_t
ZERO = Teuchos::ScalarTraits<scalar_t>::zero();
517 scalar_t ONE = Teuchos::ScalarTraits<scalar_t>::one();
519 if (beta == ZERO) y.putScalar(ZERO);
524 if (permMatrix == Teuchos::null) {
527 lowerTriangularMatrix->apply(x, y, Teuchos::NO_TRANS, alpha, beta);
530 lowerTriangularMatrix->apply(x, y, Teuchos::TRANS, alpha, ONE);
533 y.elementWiseMultiply(-alpha, *diag, x, ONE);
542 vector_t xtmp(x.getMap(), x.getNumVectors());
543 vector_t ytmp(y.getMap(), y.getNumVectors());
545 permMatrix->apply(x, xtmp, Teuchos::TRANS);
546 if (beta != ZERO) permMatrix->apply(y, ytmp, Teuchos::TRANS);
549 lowerTriangularMatrix->apply(xtmp, ytmp, Teuchos::NO_TRANS, alpha, beta);
552 lowerTriangularMatrix->apply(xtmp, ytmp, Teuchos::TRANS, alpha, ONE);
555 ytmp.elementWiseMultiply(-alpha, *diag, xtmp, ONE);
557 permMatrix->apply(ytmp, y, Teuchos::NO_TRANS);
561 Teuchos::RCP<const map_t> getDomainMap()
const {
562 return lowerTriangularMatrix->getDomainMap();
565 Teuchos::RCP<const map_t> getRangeMap()
const {
566 return lowerTriangularMatrix->getRangeMap();
569 bool hasTransposeApply()
const {
return true;}
572 const Teuchos::RCP<const matrix_t > lowerTriangularMatrix;
573 const Teuchos::RCP<const matrix_t > permMatrix;
574 Teuchos::RCP<vector_t> diag;
Communication plan for data redistribution from a uniquely-owned to a (possibly) multiply-owned distr...
LocalOrdinal local_ordinal_type
The type of local indices.
Sparse matrix that presents a row-oriented interface that lets users read or modify entries...
GlobalOrdinal global_ordinal_type
The type of global indices.
One or more distributed dense vectors.
virtual void apply(const MultiVector< scalar_t, Tpetra::Map<>::local_ordinal_type, Tpetra::Map<>::global_ordinal_type, Node > &X, MultiVector< scalar_t, Tpetra::Map<>::local_ordinal_type, Tpetra::Map<>::global_ordinal_type, Node > &Y, Teuchos::ETransp mode=Teuchos::NO_TRANS, scalar_talpha=Teuchos::ScalarTraits< scalar_t >::one(), scalar_tbeta=Teuchos::ScalarTraits< scalar_t >::zero()) const =0
Computes the operator-multivector application.
void sort(View &view, const size_t &size)
Convenience wrapper for std::sort for host-accessible views.
size_t global_size_t
Global size_t object.
Abstract interface for operators (e.g., matrices and preconditioners).
Communication plan for data redistribution from a (possibly) multiply-owned to a uniquely-owned distr...
Functions for initializing and finalizing Tpetra.
Sets up and executes a communication plan for a Tpetra DistObject.
Replace old values with zero.
A parallel distribution of indices over processes.
A distributed dense vector.