Teko Version of the Day
Loading...
Searching...
No Matches
Teko_InterlacedTpetra.cpp
1// @HEADER
2// *****************************************************************************
3// Teko: A package for block and physics based preconditioning
4//
5// Copyright 2010 NTESS and the Teko contributors.
6// SPDX-License-Identifier: BSD-3-Clause
7// *****************************************************************************
8// @HEADER
9
10#include "Teko_InterlacedTpetra.hpp"
11#include "Tpetra_Import.hpp"
12
13#include <vector>
14
15using Teuchos::RCP;
16using Teuchos::rcp;
17
18namespace Teko {
19namespace TpetraHelpers {
20namespace Strided {
21
22// this assumes that there are numGlobals with numVars each interlaced
23// i.e. for numVars = 2 (u,v) then the vector is
24// [u_0,v_0,u_1,v_1,u_2,v_2, ..., u_(numGlobals-1),v_(numGlobals-1)]
25void buildSubMaps(GO numGlobals, int numVars, const Teuchos::Comm<int>& comm,
26 std::vector<std::pair<int, RCP<Tpetra::Map<LO, GO, NT> > > >& subMaps) {
27 std::vector<int> vars;
28
29 // build vector describing the sub maps
30 for (int i = 0; i < numVars; i++) vars.push_back(1);
31
32 // build all the submaps
33 buildSubMaps(numGlobals, vars, comm, subMaps);
34}
35
36// build maps to make other conversions
37void buildSubMaps(const Tpetra::Map<LO, GO, NT>& globalMap, const std::vector<int>& vars,
38 const Teuchos::Comm<int>& comm,
39 std::vector<std::pair<int, Teuchos::RCP<Tpetra::Map<LO, GO, NT> > > >& subMaps) {
40 buildSubMaps(globalMap.getGlobalNumElements(), globalMap.getLocalNumElements(),
41 globalMap.getMinGlobalIndex(), vars, comm, subMaps);
42}
43
44// build maps to make other conversions
45void buildSubMaps(GO numGlobals, const std::vector<int>& vars, const Teuchos::Comm<int>& comm,
46 std::vector<std::pair<int, Teuchos::RCP<Tpetra::Map<LO, GO, NT> > > >& subMaps) {
47 std::vector<int>::const_iterator varItr;
48
49 // compute total number of variables
50 int numGlobalVars = 0;
51 for (varItr = vars.begin(); varItr != vars.end(); ++varItr) numGlobalVars += *varItr;
52
53 // must be an even number of globals
54 TEUCHOS_ASSERT((numGlobals % numGlobalVars) == 0);
55
56 Tpetra::Map<LO, GO, NT> sampleMap(numGlobals / numGlobalVars, 0, rcpFromRef(comm));
57
58 buildSubMaps(numGlobals, numGlobalVars * sampleMap.getLocalNumElements(),
59 numGlobalVars * sampleMap.getMinGlobalIndex(), vars, comm, subMaps);
60}
61
62// build maps to make other conversions
63void buildSubMaps(GO numGlobals, LO numMyElements, GO minMyGID, const std::vector<int>& vars,
64 const Teuchos::Comm<int>& comm,
65 std::vector<std::pair<int, Teuchos::RCP<Tpetra::Map<LO, GO, NT> > > >& subMaps) {
66 std::vector<int>::const_iterator varItr;
67
68 // compute total number of variables
69 int numGlobalVars = 0;
70 for (varItr = vars.begin(); varItr != vars.end(); ++varItr) numGlobalVars += *varItr;
71
72 // must be an even number of globals
73 TEUCHOS_ASSERT((numGlobals % numGlobalVars) == 0);
74 TEUCHOS_ASSERT((numMyElements % numGlobalVars) == 0);
75 TEUCHOS_ASSERT((minMyGID % numGlobalVars) == 0);
76
77 LO numBlocks = numMyElements / numGlobalVars;
78 GO minBlockID = minMyGID / numGlobalVars;
79
80 subMaps.clear();
81
82 // index into local block in strided map
83 GO blockOffset = 0;
84 for (varItr = vars.begin(); varItr != vars.end(); ++varItr) {
85 LO numLocalVars = *varItr;
86 GO numAllElmts = numLocalVars * numGlobals / numGlobalVars;
87#ifndef NDEBUG
88 LO numMyElmts = numLocalVars * numBlocks;
89#endif
90
91 // create global arrays describing the as of yet uncreated maps
92 std::vector<GO> subGlobals;
93 std::vector<GO> contigGlobals; // the contiguous globals
94
95 // loop over each block of variables
96 LO count = 0;
97 for (LO blockNum = 0; blockNum < numBlocks; blockNum++) {
98 // loop over each local variable in the block
99 for (LO local = 0; local < numLocalVars; ++local) {
100 // global block number = minGID+blockNum
101 // block begin global id = numGlobalVars*(minGID+blockNum)
102 // global id block offset = blockOffset+local
103 subGlobals.push_back((minBlockID + blockNum) * numGlobalVars + blockOffset + local);
104
105 // also build the contiguous IDs
106 contigGlobals.push_back(numLocalVars * minBlockID + count);
107 count++;
108 }
109 }
110
111 // sanity check
112 assert((size_t)numMyElmts == subGlobals.size());
113
114 // create the map with contiguous elements and the map with global elements
115 RCP<Tpetra::Map<LO, GO, NT> > subMap = rcp(new Tpetra::Map<LO, GO, NT>(
116 numAllElmts, Teuchos::ArrayView<GO>(subGlobals), 0, rcpFromRef(comm)));
117 RCP<Tpetra::Map<LO, GO, NT> > contigMap = rcp(new Tpetra::Map<LO, GO, NT>(
118 numAllElmts, Teuchos::ArrayView<GO>(contigGlobals), 0, rcpFromRef(comm)));
119
120 Teuchos::set_extra_data(contigMap, "contigMap", Teuchos::inOutArg(subMap));
121 subMaps.push_back(std::make_pair(numLocalVars, subMap));
122
123 // update the block offset
124 blockOffset += numLocalVars;
125 }
126}
127
128void buildExportImport(const Tpetra::Map<LO, GO, NT>& baseMap,
129 const std::vector<std::pair<int, RCP<Tpetra::Map<LO, GO, NT> > > >& subMaps,
130 std::vector<RCP<Tpetra::Export<LO, GO, NT> > >& subExport,
131 std::vector<RCP<Tpetra::Import<LO, GO, NT> > >& subImport) {
132 std::vector<std::pair<int, RCP<Tpetra::Map<LO, GO, NT> > > >::const_iterator mapItr;
133
134 // build importers and exporters
135 for (mapItr = subMaps.begin(); mapItr != subMaps.end(); ++mapItr) {
136 // exctract basic map
137 const Tpetra::Map<LO, GO, NT>& map = *(mapItr->second);
138
139 // add new elements to vectors
140 subImport.push_back(rcp(new Tpetra::Import<LO, GO, NT>(rcpFromRef(baseMap), rcpFromRef(map))));
141 subExport.push_back(rcp(new Tpetra::Export<LO, GO, NT>(rcpFromRef(map), rcpFromRef(baseMap))));
142 }
143}
144
145void buildSubVectors(const std::vector<std::pair<int, RCP<Tpetra::Map<LO, GO, NT> > > >& subMaps,
146 std::vector<RCP<Tpetra::MultiVector<ST, LO, GO, NT> > >& subVectors,
147 int count) {
148 std::vector<std::pair<int, RCP<Tpetra::Map<LO, GO, NT> > > >::const_iterator mapItr;
149
150 // build vectors
151 for (mapItr = subMaps.begin(); mapItr != subMaps.end(); ++mapItr) {
152 // exctract basic map
153 const Tpetra::Map<LO, GO, NT>& map =
154 *(Teuchos::get_extra_data<RCP<Tpetra::Map<LO, GO, NT> > >(mapItr->second, "contigMap"));
155
156 // add new elements to vectors
157 RCP<Tpetra::MultiVector<ST, LO, GO, NT> > mv =
158 rcp(new Tpetra::MultiVector<ST, LO, GO, NT>(rcpFromRef(map), count));
159 Teuchos::set_extra_data(mapItr->second, "globalMap", Teuchos::inOutArg(mv));
160 subVectors.push_back(mv);
161 }
162}
163
164void associateSubVectors(
165 const std::vector<std::pair<int, RCP<Tpetra::Map<LO, GO, NT> > > >& subMaps,
166 std::vector<RCP<const Tpetra::MultiVector<ST, LO, GO, NT> > >& subVectors) {
167 std::vector<std::pair<int, RCP<Tpetra::Map<LO, GO, NT> > > >::const_iterator mapItr;
168 std::vector<RCP<const Tpetra::MultiVector<ST, LO, GO, NT> > >::iterator vecItr;
169
170 TEUCHOS_ASSERT(subMaps.size() == subVectors.size());
171
172 // associate the sub vectors with the subMaps
173 for (mapItr = subMaps.begin(), vecItr = subVectors.begin(); mapItr != subMaps.end();
174 ++mapItr, ++vecItr)
175 Teuchos::set_extra_data(mapItr->second, "globalMap", Teuchos::inOutArg(*vecItr),
176 Teuchos::POST_DESTROY, false);
177}
178
179// build a single subblock Epetra_CrsMatrix
180RCP<Tpetra::CrsMatrix<ST, LO, GO, NT> > buildSubBlock(
181 int i, int j, const RCP<const Tpetra::CrsMatrix<ST, LO, GO, NT> >& A,
182 const std::vector<std::pair<int, RCP<Tpetra::Map<LO, GO, NT> > > >& subMaps) {
183 // get the number of variables families
184 int numVarFamily = subMaps.size();
185
186 TEUCHOS_ASSERT(i >= 0 && i < numVarFamily);
187 TEUCHOS_ASSERT(j >= 0 && j < numVarFamily);
188
189 const Tpetra::Map<LO, GO, NT>& gRowMap = *subMaps[i].second;
190 const Tpetra::Map<LO, GO, NT>& rowMap =
191 *Teuchos::get_extra_data<RCP<Tpetra::Map<LO, GO, NT> > >(subMaps[i].second, "contigMap");
192 const Tpetra::Map<LO, GO, NT>& colMap =
193 *Teuchos::get_extra_data<RCP<Tpetra::Map<LO, GO, NT> > >(subMaps[j].second, "contigMap");
194 int colFamilyCnt = subMaps[j].first;
195
196 // compute the number of global variables
197 // and the row and column block offset
198 GO numGlobalVars = 0;
199 GO rowBlockOffset = 0;
200 GO colBlockOffset = 0;
201 for (int k = 0; k < numVarFamily; k++) {
202 numGlobalVars += subMaps[k].first;
203
204 // compute block offsets
205 if (k < i) rowBlockOffset += subMaps[k].first;
206 if (k < j) colBlockOffset += subMaps[k].first;
207 }
208
209 // copy all global rows to here
210 Tpetra::Import<LO, GO, NT> import(A->getRowMap(), rcpFromRef(gRowMap));
211 Tpetra::CrsMatrix<ST, LO, GO, NT> localA(rcpFromRef(gRowMap), 0);
212 localA.doImport(*A, import, Tpetra::INSERT);
213
214 // get entry information
215 LO numMyRows = rowMap.getLocalNumElements();
216 LO maxNumEntries = A->getGlobalMaxNumRowEntries();
217
218 // for extraction
219 auto indices = typename Tpetra::CrsMatrix<ST, LO, GO, NT>::nonconst_global_inds_host_view_type(
220 Kokkos::ViewAllocateWithoutInitializing("rowIndices"), maxNumEntries);
221 auto values = typename Tpetra::CrsMatrix<ST, LO, GO, NT>::nonconst_values_host_view_type(
222 Kokkos::ViewAllocateWithoutInitializing("rowIndices"), maxNumEntries);
223
224 // for counting row sizes
225 std::vector<size_t> numEntriesPerRow(numMyRows, 0);
226
227 const size_t invalid = Teuchos::OrdinalTraits<size_t>::invalid();
228
229 // Count the sizes of each row, using same logic as insertion below
230 for (LO localRow = 0; localRow < numMyRows; localRow++) {
231 size_t numEntries = invalid;
232 GO globalRow = gRowMap.getGlobalElement(localRow);
233 GO contigRow = rowMap.getGlobalElement(localRow);
234
235 TEUCHOS_ASSERT(globalRow >= 0);
236 TEUCHOS_ASSERT(contigRow >= 0);
237
238 // extract a global row copy
239 localA.getGlobalRowCopy(globalRow, indices, values, numEntries);
240 LO numOwnedCols = 0;
241 for (size_t localCol = 0; localCol < numEntries; localCol++) {
242 GO globalCol = indices(localCol);
243
244 // determinate which block this column ID is in
245 int block = globalCol / numGlobalVars;
246
247 bool inFamily = true;
248
249 // test the beginning of the block
250 inFamily &= (block * numGlobalVars + colBlockOffset <= globalCol);
251 inFamily &= ((block * numGlobalVars + colBlockOffset + colFamilyCnt) > globalCol);
252
253 // is this column in the variable family
254 if (inFamily) {
255 numOwnedCols++;
256 }
257 }
258 numEntriesPerRow[localRow] += numOwnedCols;
259 }
260
261 RCP<Tpetra::CrsMatrix<ST, LO, GO, NT> > mat = rcp(new Tpetra::CrsMatrix<ST, LO, GO, NT>(
262 rcpFromRef(rowMap), Teuchos::ArrayView<const size_t>(numEntriesPerRow)));
263
264 // for insertion
265 std::vector<GO> colIndices(maxNumEntries);
266 std::vector<ST> colValues(maxNumEntries);
267
268 // insert each row into subblock
269 // let FillComplete handle column distribution
270 for (LO localRow = 0; localRow < numMyRows; localRow++) {
271 size_t numEntries = invalid;
272 GO globalRow = gRowMap.getGlobalElement(localRow);
273 GO contigRow = rowMap.getGlobalElement(localRow);
274
275 TEUCHOS_ASSERT(globalRow >= 0);
276 TEUCHOS_ASSERT(contigRow >= 0);
277
278 // extract a global row copy
279 localA.getGlobalRowCopy(globalRow, indices, values, numEntries);
280 LO numOwnedCols = 0;
281 for (size_t localCol = 0; localCol < numEntries; localCol++) {
282 GO globalCol = indices(localCol);
283
284 // determinate which block this column ID is in
285 int block = globalCol / numGlobalVars;
286
287 bool inFamily = true;
288
289 // test the beginning of the block
290 inFamily &= (block * numGlobalVars + colBlockOffset <= globalCol);
291 inFamily &= ((block * numGlobalVars + colBlockOffset + colFamilyCnt) > globalCol);
292
293 // is this column in the variable family
294 if (inFamily) {
295 GO familyOffset = globalCol - (block * numGlobalVars + colBlockOffset);
296
297 colIndices[numOwnedCols] = block * colFamilyCnt + familyOffset;
298 colValues[numOwnedCols] = values(localCol);
299
300 numOwnedCols++;
301 }
302 }
303
304 // insert it into the new matrix
305 colIndices.resize(numOwnedCols);
306 colValues.resize(numOwnedCols);
307 mat->insertGlobalValues(contigRow, Teuchos::ArrayView<GO>(colIndices),
308 Teuchos::ArrayView<ST>(colValues));
309 colIndices.resize(maxNumEntries);
310 colValues.resize(maxNumEntries);
311 }
312
313 // fill it and automagically optimize the storage
314 mat->fillComplete(rcpFromRef(colMap), rcpFromRef(rowMap));
315
316 return mat;
317}
318
319// rebuild a single subblock Epetra_CrsMatrix
320void rebuildSubBlock(int i, int j, const RCP<const Tpetra::CrsMatrix<ST, LO, GO, NT> >& A,
321 const std::vector<std::pair<int, RCP<Tpetra::Map<LO, GO, NT> > > >& subMaps,
322 Tpetra::CrsMatrix<ST, LO, GO, NT>& mat) {
323 // get the number of variables families
324 int numVarFamily = subMaps.size();
325
326 TEUCHOS_ASSERT(i >= 0 && i < numVarFamily);
327 TEUCHOS_ASSERT(j >= 0 && j < numVarFamily);
328 TEUCHOS_ASSERT(mat.isFillComplete());
329
330 const Tpetra::Map<LO, GO, NT>& gRowMap = *subMaps[i].second;
331 const Tpetra::Map<LO, GO, NT>& rowMap =
332 *Teuchos::get_extra_data<RCP<Tpetra::Map<LO, GO, NT> > >(subMaps[i].second, "contigMap");
333 const Tpetra::Map<LO, GO, NT>& colMap =
334 *Teuchos::get_extra_data<RCP<Tpetra::Map<LO, GO, NT> > >(subMaps[j].second, "contigMap");
335 int colFamilyCnt = subMaps[j].first;
336
337 // compute the number of global variables
338 // and the row and column block offset
339 GO numGlobalVars = 0;
340 GO rowBlockOffset = 0;
341 GO colBlockOffset = 0;
342 for (int k = 0; k < numVarFamily; k++) {
343 numGlobalVars += subMaps[k].first;
344
345 // compute block offsets
346 if (k < i) rowBlockOffset += subMaps[k].first;
347 if (k < j) colBlockOffset += subMaps[k].first;
348 }
349
350 // copy all global rows to here
351 Tpetra::Import<LO, GO, NT> import(A->getRowMap(), rcpFromRef(gRowMap));
352 Tpetra::CrsMatrix<ST, LO, GO, NT> localA(rcpFromRef(gRowMap), 0);
353 localA.doImport(*A, import, Tpetra::INSERT);
354
355 // clear out the old matrix
356 mat.resumeFill();
357 mat.setAllToScalar(0.0);
358
359 // get entry information
360 LO numMyRows = rowMap.getLocalNumElements();
361 GO maxNumEntries = A->getGlobalMaxNumRowEntries();
362
363 // for extraction
364 auto indices = typename Tpetra::CrsMatrix<ST, LO, GO, NT>::nonconst_global_inds_host_view_type(
365 Kokkos::ViewAllocateWithoutInitializing("rowIndices"), maxNumEntries);
366 auto values = typename Tpetra::CrsMatrix<ST, LO, GO, NT>::nonconst_values_host_view_type(
367 Kokkos::ViewAllocateWithoutInitializing("rowIndices"), maxNumEntries);
368
369 // for insertion
370 std::vector<GO> colIndices(maxNumEntries);
371 std::vector<ST> colValues(maxNumEntries);
372
373 const size_t invalid = Teuchos::OrdinalTraits<size_t>::invalid();
374
375 // insert each row into subblock
376 // let FillComplete handle column distribution
377 for (LO localRow = 0; localRow < numMyRows; localRow++) {
378 size_t numEntries = invalid;
379 GO globalRow = gRowMap.getGlobalElement(localRow);
380 GO contigRow = rowMap.getGlobalElement(localRow);
381
382 TEUCHOS_ASSERT(globalRow >= 0);
383 TEUCHOS_ASSERT(contigRow >= 0);
384
385 // extract a global row copy
386 localA.getGlobalRowCopy(globalRow, indices, values, numEntries);
387
388 LO numOwnedCols = 0;
389 for (size_t localCol = 0; localCol < numEntries; localCol++) {
390 GO globalCol = indices(localCol);
391
392 // determinate which block this column ID is in
393 int block = globalCol / numGlobalVars;
394
395 bool inFamily = true;
396
397 // test the beginning of the block
398 inFamily &= (block * numGlobalVars + colBlockOffset <= globalCol);
399 inFamily &= ((block * numGlobalVars + colBlockOffset + colFamilyCnt) > globalCol);
400
401 // is this column in the variable family
402 if (inFamily) {
403 GO familyOffset = globalCol - (block * numGlobalVars + colBlockOffset);
404
405 colIndices[numOwnedCols] = block * colFamilyCnt + familyOffset;
406 colValues[numOwnedCols] = values(localCol);
407
408 numOwnedCols++;
409 }
410 }
411
412 // insert it into the new matrix
413 colIndices.resize(numOwnedCols);
414 colValues.resize(numOwnedCols);
415 mat.sumIntoGlobalValues(contigRow, Teuchos::ArrayView<GO>(colIndices),
416 Teuchos::ArrayView<ST>(colValues));
417 colIndices.resize(maxNumEntries);
418 colValues.resize(maxNumEntries);
419 }
420 mat.fillComplete(rcpFromRef(colMap), rcpFromRef(rowMap));
421}
422
423// collect subvectors into a single global vector
424void many2one(Tpetra::MultiVector<ST, LO, GO, NT>& one,
425 const std::vector<RCP<const Tpetra::MultiVector<ST, LO, GO, NT> > >& many,
426 const std::vector<RCP<Tpetra::Export<LO, GO, NT> > >& subExport) {
427 // std::vector<RCP<const Epetra_Vector> >::const_iterator vecItr;
428 std::vector<RCP<const Tpetra::MultiVector<ST, LO, GO, NT> > >::const_iterator vecItr;
429 std::vector<RCP<Tpetra::Export<LO, GO, NT> > >::const_iterator expItr;
430
431 // using Exporters fill the empty vector from the sub-vectors
432 for (vecItr = many.begin(), expItr = subExport.begin(); vecItr != many.end();
433 ++vecItr, ++expItr) {
434 // for ease of access to the source
435 RCP<const Tpetra::MultiVector<ST, LO, GO, NT> > srcVec = *vecItr;
436
437 // extract the map with global indicies from the current vector
438 const Tpetra::Map<LO, GO, NT>& globalMap =
439 *(Teuchos::get_extra_data<RCP<Tpetra::Map<LO, GO, NT> > >(srcVec, "globalMap"));
440
441 // build the export vector as a view of the destination
442 GO lda = srcVec->getStride();
443 GO srcSize = srcVec->getGlobalLength() * srcVec->getNumVectors();
444 std::vector<ST> srcArray(srcSize);
445 Teuchos::ArrayView<ST> srcVals(srcArray);
446 srcVec->get1dCopy(srcVals, lda);
447 Tpetra::MultiVector<ST, LO, GO, NT> exportVector(rcpFromRef(globalMap), srcVals, lda,
448 srcVec->getNumVectors());
449
450 // perform the export
451 one.doExport(exportVector, **expItr, Tpetra::INSERT);
452 }
453}
454
455// distribute one global vector into a many subvectors
456void one2many(std::vector<RCP<Tpetra::MultiVector<ST, LO, GO, NT> > >& many,
457 const Tpetra::MultiVector<ST, LO, GO, NT>& single,
458 const std::vector<RCP<Tpetra::Import<LO, GO, NT> > >& subImport) {
459 // std::vector<RCP<Epetra_Vector> >::const_iterator vecItr;
460 std::vector<RCP<Tpetra::MultiVector<ST, LO, GO, NT> > >::const_iterator vecItr;
461 std::vector<RCP<Tpetra::Import<LO, GO, NT> > >::const_iterator impItr;
462
463 // using Importers fill the sub vectors from the mama vector
464 for (vecItr = many.begin(), impItr = subImport.begin(); vecItr != many.end();
465 ++vecItr, ++impItr) {
466 // for ease of access to the destination
467 RCP<Tpetra::MultiVector<ST, LO, GO, NT> > destVec = *vecItr;
468
469 // extract the map with global indicies from the current vector
470 const Tpetra::Map<LO, GO, NT>& globalMap =
471 *(Teuchos::get_extra_data<RCP<Tpetra::Map<LO, GO, NT> > >(destVec, "globalMap"));
472
473 // build the import vector as a view on the destination
474 GO destLDA = destVec->getStride();
475 GO destSize = destVec->getGlobalLength() * destVec->getNumVectors();
476 std::vector<ST> destArray(destSize);
477 Teuchos::ArrayView<ST> destVals(destArray);
478 destVec->get1dCopy(destVals, destLDA);
479 Tpetra::MultiVector<ST, LO, GO, NT> importVector(rcpFromRef(globalMap), destVals, destLDA,
480 destVec->getNumVectors());
481
482 // perform the import
483 importVector.doImport(single, **impItr, Tpetra::INSERT);
484
485 Tpetra::Import<LO, GO, NT> importer(destVec->getMap(), destVec->getMap());
486 importVector.replaceMap(destVec->getMap());
487 destVec->doImport(importVector, importer, Tpetra::INSERT);
488 }
489}
490
491} // namespace Strided
492} // namespace TpetraHelpers
493} // end namespace Teko
Generated on for Teko by doxygen 1.15.0