teko: fixes to support staticProfile matrix construction in Tpetra #5602

trilinos · Aug 2, 2019 · afa0c9f · afa0c9f
1 parent 0cc8a93
commit afa0c9f
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 2 deletions.
diff --git a/packages/teko/src/Tpetra/Teko_BlockingTpetra.cpp b/packages/teko/src/Tpetra/Teko_BlockingTpetra.cpp
@@ -236,7 +236,6 @@ RCP<Tpetra::CrsMatrix<ST,LO,GO,NT> > buildSubBlock(int i,int j,const RCP<const T
    const RCP<Tpetra::Vector<GO,LO,GO,NT> > plocal2ContigGIDs = getSubBlockColumnGIDs(*A,subMaps[j]);
    Tpetra::Vector<GO,LO,GO,NT> & local2ContigGIDs = *plocal2ContigGIDs;
 
-   RCP<Tpetra::CrsMatrix<ST,LO,GO,NT> > mat = rcp(new Tpetra::CrsMatrix<ST,LO,GO,NT>(rowMap,0));
 
    // get entry information
    LO numMyRows = rowMap->getNodeNumElements();
@@ -250,6 +249,34 @@ RCP<Tpetra::CrsMatrix<ST,LO,GO,NT> > buildSubBlock(int i,int j,const RCP<const T
    std::vector<GO> colIndices(maxNumEntries);
    std::vector<ST> colValues(maxNumEntries);
 
+   // for counting
+   std::vector<size_t> nEntriesPerRow(numMyRows);
+
+   // insert each row into subblock
+   // Count the number of entries per row in the new matrix
+   for(LO localRow=0;localRow<numMyRows;localRow++) {
+      size_t numEntries = -1;
+      GO globalRow = gRowMap->getGlobalElement(localRow);
+      LO lid = A->getRowMap()->getLocalElement(globalRow);
+      TEUCHOS_ASSERT(lid>-1);
+
+      A->getLocalRowCopy(lid, Teuchos::ArrayView<LO>(indices), Teuchos::ArrayView<ST>(values),numEntries);
+
+      LO numOwnedCols = 0;
+      for(size_t localCol=0;localCol<numEntries;localCol++) {
+         TEUCHOS_ASSERT(indices[localCol]>-1);
+
+         // if global id is not owned by this column
+
+         GO gid = local2ContigGIDs.getData()[indices[localCol]];
+         if(gid==-1) continue; // in contiguous row
+         numOwnedCols++;
+      }
+      nEntriesPerRow[localRow] = numOwnedCols;
+   }
+
+   RCP<Tpetra::CrsMatrix<ST,LO,GO,NT> > mat = rcp(new Tpetra::CrsMatrix<ST,LO,GO,NT>(rowMap,Teuchos::ArrayView<size_t>(nEntriesPerRow)));
+
    // insert each row into subblock
    // let FillComplete handle column distribution
    for(LO localRow=0;localRow<numMyRows;localRow++) {

diff --git a/packages/teko/tests/src/tGraphLaplacian_tpetra.cpp b/packages/teko/tests/src/tGraphLaplacian_tpetra.cpp
@@ -77,7 +77,13 @@ void tGraphLaplacian_tpetra::initializeTest()
 Teuchos::RCP<Tpetra::CrsMatrix<ST,LO,GO,NT> > stencil(const Teuchos::Comm<int> & comm)
 {
    Teuchos::RCP<Tpetra::Map<LO,GO,NT> > map = rcp(new Tpetra::Map<LO,GO,NT>(5,0,rcpFromRef(comm)));
-   Teuchos::RCP<Tpetra::CrsMatrix<ST,LO,GO,NT> > mat = rcp(new Tpetra::CrsMatrix<ST,LO,GO,NT>(map,0));
+   Teuchos::RCP<Tpetra::CrsMatrix<ST,LO,GO,NT> > mat = rcp(new Tpetra::CrsMatrix<ST,LO,GO,NT>(map,6));
+   // KDD 8/2019:  The value 6 above is needed because this test inserts 
+   // three indices twice into row 4 of the matrix.
+   // Tpetra is robust enough to correctly handle the double insertion. 
+   // But if this double insertion is, in fact, an error, the 6 can probably
+   // be changed to 3.  Note that this test doesn't insert any nonzeros into
+   // row 0 of the matrix.
 
    // arrays for indicies and values
    GO indicies[3];