Merge branch 'amesos2-gather' of https://github.com/iyamazaki/Trilinos …

…into develop
trilinos · Jan 22, 2025 · 2b41dd2 · 2b41dd2
2 parents 90df967 + 0da82d5
commit 2b41dd2
Show file tree

Hide file tree

Showing 2 changed files with 73 additions and 4 deletions.
diff --git a/packages/amesos2/src/Amesos2_TpetraCrsMatrix_MatrixAdapter_def.hpp b/packages/amesos2/src/Amesos2_TpetraCrsMatrix_MatrixAdapter_def.hpp
@@ -186,6 +186,7 @@ namespace Amesos2 {
         auto nRanks = comm->getSize();
         auto myRank = comm->getRank();
 
+        global_ordinal_t indexBase = rowMap->getIndexBase();
         global_ordinal_t nRows = this->mat_->getGlobalNumRows();
         auto lclMatrix = this->mat_->getLocalMatrixDevice();
 
@@ -231,7 +232,7 @@ namespace Amesos2 {
               }
               recvDispls(nRanks) = 0;
             }
-            // gether g2l perm
+            // gether g2l perm (convert to 0-base)
             {
               host_ordinal_type_array lclMap;
               Kokkos::resize(lclMap, myNRows);
@@ -248,6 +249,7 @@ namespace Amesos2 {
                                                    0, *comm);
               if (myRank == 0) {
                 for (int i=0; i < nRows; i++) {
+                  perm_g2l(i) -= indexBase;
                   if (i != perm_g2l(i)) need_to_perm = true;
                 }
               }
@@ -299,9 +301,9 @@ namespace Amesos2 {
                 recvDispls(p+1) = recvDispls(p) + recvCounts(p);
               }
             }
-            // -- convert to global colids
+            // -- convert to global colids & convert to 0-base
             KV_GO lclColind_ ("localColind_", lclColind.extent(0));
-            for (int i = 0; i < int(lclColind.extent(0)); i++) lclColind_(i) = colMap->getGlobalElement((lclColind(i)));
+            for (int i = 0; i < int(lclColind.extent(0)); i++) lclColind_(i) = (colMap->getGlobalElement((lclColind(i))) - indexBase);
             if (column_major || need_to_perm) {
               Kokkos::resize(indices_t, indices.extent(0));
               Teuchos::gatherv<int, LocalOrdinal> (lclColind_.data(), lclColind_.extent(0), indices_t.data(), 

diff --git a/packages/amesos2/test/solvers/KLU2_UnitTests.cpp b/packages/amesos2/test/solvers/KLU2_UnitTests.cpp
@@ -451,6 +451,72 @@ namespace {
     TEST_COMPARE_FLOATING_ARRAYS( xhatnorms, xnorms, 0.005 );
   }
 
+  //! @test Test for one-base
+  TEUCHOS_UNIT_TEST_TEMPLATE_3_DECL( KLU2, BaseOne, SCALAR, LO, GO )
+  {
+    typedef CrsMatrix<SCALAR,LO,GO,Node> MAT;
+    typedef ScalarTraits<SCALAR> ST;
+    typedef MultiVector<SCALAR,LO,GO,Node> MV;
+    typedef typename ST::magnitudeType Mag;
+
+    using Tpetra::global_size_t;
+    using Teuchos::tuple;
+    using Teuchos::RCP;
+    using Teuchos::rcp;
+
+    RCP<const Comm<int> > comm = Tpetra::getDefaultComm();
+
+    const global_size_t numVectors = 1;
+    const global_size_t nrows = 6;
+    const GO numGlobalEntries = nrows;
+
+    // Create one-base Map
+    const GO indexBase = 1;
+    typedef Tpetra::Map<LO,GO>  map_type;
+    RCP< const map_type > map = rcp( new map_type(numGlobalEntries, indexBase, comm) );
+    const LO numLocalEntries = LO(map->getLocalNumElements());
+
+    // Create a diagobal matrix = diag(1:nrows)
+    RCP<MAT> A = rcp( new MAT(map,1) );
+
+    Teuchos::Array<GO> gblColIndsBuf(1);
+    Teuchos::Array<SCALAR> valsBuf(1);
+    for (LO lclRow = 0; lclRow < numLocalEntries; ++lclRow) {
+      const GO gblRow = map->getGlobalElement(lclRow);
+      const GO gblCol = gblRow;
+      const LO numEnt = 1;
+      valsBuf[0] = SCALAR(gblRow);
+      gblColIndsBuf[0] = gblCol;
+
+      Teuchos::ArrayView<GO> gblColInds = gblColIndsBuf.view(0, numEnt);
+      Teuchos::ArrayView<SCALAR> vals = valsBuf.view(0, numEnt);
+      A->insertGlobalValues(gblRow, gblColInds, vals);
+    }
+    A->fillComplete();
+
+    // Create Xhat = ones(nrows, 1), X, and B
+    RCP<MV> Xhat = rcp(new MV(map,numVectors));
+    RCP<MV> X = rcp(new MV(map,numVectors));
+    RCP<MV> B = rcp(new MV(map,numVectors));
+    Xhat->putScalar(SCALAR(1.0));
+    A->apply(*Xhat, *B);
+
+    // Create solver interface with Amesos2 factory method
+    RCP<Amesos2::Solver<MAT,MV> > solver = Amesos2::create<MAT,MV>("KLU2", A, X, B);
+    solver->symbolicFactorization().numericFactorization().solve();
+
+    A->describe(out, Teuchos::VERB_EXTREME);
+    B->describe(out, Teuchos::VERB_EXTREME);
+    Xhat->describe(out, Teuchos::VERB_EXTREME);
+    X->describe(out, Teuchos::VERB_EXTREME);
+
+    // Check result of solve
+    Array<Mag> xhatnorms(numVectors), xnorms(numVectors);
+    Xhat->norm2(xhatnorms());
+    X->norm2(xnorms());
+    TEST_COMPARE_FLOATING_ARRAYS( xhatnorms, xnorms, 0.005 );
+  }
+
   //! @test Test for non-contiguous GIDs.
   TEUCHOS_UNIT_TEST_TEMPLATE_3_DECL( KLU2, NonContigGID, SCALAR, LO, GO )
   {
@@ -970,7 +1036,8 @@ namespace {
   TEUCHOS_UNIT_TEST_TEMPLATE_3_INSTANT( KLU2, Solve, SCALAR, LO, GO ) \
   TEUCHOS_UNIT_TEST_TEMPLATE_3_INSTANT( KLU2, SolveIR, SCALAR, LO, GO ) \
   TEUCHOS_UNIT_TEST_TEMPLATE_3_INSTANT( KLU2, SolveTrans, SCALAR, LO, GO ) \
-  TEUCHOS_UNIT_TEST_TEMPLATE_3_INSTANT( KLU2, NonContigGID, SCALAR, LO, GO )
+  TEUCHOS_UNIT_TEST_TEMPLATE_3_INSTANT( KLU2, NonContigGID, SCALAR, LO, GO ) \
+  TEUCHOS_UNIT_TEST_TEMPLATE_3_INSTANT( KLU2, BaseOne, SCALAR, LO, GO )
 
 #ifdef HAVE_AMESOS2_EPETRA
 #define UNIT_TEST_GROUP_EPETRA( LO, GO, SCALAR)  \