Teko  Version of the Day
 All Classes Files Functions Variables Pages
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 
15 using Teuchos::RCP;
16 using Teuchos::rcp;
17 
18 namespace Teko {
19 namespace TpetraHelpers {
20 namespace 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)]
25 void 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
37 void 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
45 void 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
63 void 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 
128 void 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 
145 void 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 
164 void 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
180 RCP<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
320 void 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
424 void 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
456 void 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 Fri Nov 22 2024 09:19:41 for Teko by   doxygen 1.8.5