Fixes dense solver for Kokkos interface, clean up CG solvers

trilinos · Feb 28, 2024 · 0e89333 · 0e89333
1 parent b7a7c60
commit 0e89333
Show file tree

Hide file tree

Showing 8 changed files with 13 additions and 27 deletions.
diff --git a/packages/belos/src/BelosBlockCGIter.hpp b/packages/belos/src/BelosBlockCGIter.hpp
@@ -57,7 +57,6 @@
 #include "BelosOperatorTraits.hpp"
 #include "BelosMultiVecTraits.hpp"
 #include "BelosDenseMatTraits.hpp"
-#include "BelosTeuchosDenseAdapter.hpp"
 
 #include "Teuchos_ScalarTraits.hpp"
 #include "Teuchos_ParameterList.hpp"
@@ -73,8 +72,7 @@ namespace Belos {
 
 /// \brief Stub implementation of BlockCGIter, for ScalarType types
 ///   for which Teuchos::LAPACK does NOT have a valid implementation.
-template<class ScalarType, class MV, class OP, 
-         class DM = Teuchos::SerialDenseMatrix<int, ScalarType>,
+template<class ScalarType, class MV, class OP, class DM,
          const bool lapackSupportsScalarType =
          Belos::Details::LapackSupportsScalar<ScalarType>::value>
 class BlockCGIter : virtual public CGIteration<ScalarType, MV, OP, DM> {

diff --git a/packages/belos/src/BelosBlockCGSolMgr.hpp b/packages/belos/src/BelosBlockCGSolMgr.hpp
@@ -130,7 +130,6 @@ namespace Belos {
   class BlockCGSolMgr<ScalarType, MV, OP, DM, true> :
     public Details::SolverManagerRequiresLapack<ScalarType, MV, OP, DM, true>
   {
-
   private:
     typedef MultiVecTraits<ScalarType,MV,DM> MVT;
     typedef OperatorTraits<ScalarType,MV,OP> OPT;

diff --git a/packages/belos/src/BelosCGIter.hpp b/packages/belos/src/BelosCGIter.hpp
@@ -57,7 +57,6 @@
 #include "BelosOperatorTraits.hpp"
 #include "BelosMultiVecTraits.hpp"
 #include "BelosDenseMatTraits.hpp"
-#include "BelosTeuchosDenseAdapter.hpp"
 
 #include "Teuchos_ScalarTraits.hpp"
 #include "Teuchos_ParameterList.hpp"
@@ -75,7 +74,7 @@
 
 namespace Belos {
 
-template<class ScalarType, class MV, class OP, class DM = Teuchos::SerialDenseMatrix<int, ScalarType>>  
+template<class ScalarType, class MV, class OP, class DM>  
 class CGIter : virtual public CGIteration<ScalarType,MV,OP,DM> {
 
   public:

diff --git a/packages/belos/src/BelosCGSingleRedIter.hpp b/packages/belos/src/BelosCGSingleRedIter.hpp
@@ -57,7 +57,6 @@
 #include "BelosOperatorTraits.hpp"
 #include "BelosMultiVecTraits.hpp"
 #include "BelosDenseMatTraits.hpp"
-#include "BelosTeuchosDenseAdapter.hpp"
 
 #include "Teuchos_ScalarTraits.hpp"
 #include "Teuchos_ParameterList.hpp"
@@ -75,7 +74,7 @@
 
 namespace Belos {
 
-template<class ScalarType, class MV, class OP, class DM = Teuchos::SerialDenseMatrix<int, ScalarType>>  
+template<class ScalarType, class MV, class OP, class DM>  
 class CGSingleRedIter : virtual public CGIteration<ScalarType,MV,OP,DM> {
 
   public:

diff --git a/packages/belos/src/BelosKokkosDenseAdapter.hpp b/packages/belos/src/BelosKokkosDenseAdapter.hpp
@@ -129,7 +129,6 @@ namespace Belos {
     */
     int factor()
     {
-      //std::cout << "Calling DenseSolver<Kokkos> factor!" << std::endl;
       int INFO = 0;
 
       // Only factor matrix if it is new, otherwise factors have been computed
@@ -160,16 +159,6 @@ namespace Belos {
           else
             lapack.GEEQU (M, N, Aptr, LDA, &R_[0], &C_[0], &ROWCND, &COLCND, &AMAX, &INFO); 
 
-          //for (int i=0; i<M; ++i)
-            //std::cout << "R[ " << i << " ] = " << R_[i] << std::endl;
-          //std::cout << std::endl;
-          if (!spd_)
-          {
-          //for (int i=0; i<N; ++i)
-            //std::cout << "C[ " << i << " ] = " << C_[i] << std::endl;
-          //std::cout << std::endl;
-          }
-
           if (INFO)
             return INFO;
 
@@ -181,7 +170,6 @@ namespace Belos {
               Scalar s1 = R_[j];
               for (int i=0; i<=j; i++) {
                 *ptr = *ptr*s1*R_[i];
-                //std::cout << "A[ " << i << ", " << j << " ] = " << *ptr << std::endl;
                 ptr++;
               }
             }
@@ -219,15 +207,15 @@ namespace Belos {
     */
     int solve()
     {
-      //std::cout << "Calling DenseMatrix<Kokkos> solve!" << std::endl;
       bool transpose = (TRANS_ != Teuchos::NO_TRANS) ? true : false;
 
+      DMT::SyncDeviceToHost(*X_);
+
       // LAPACK overwrites RHS vector with solution vector, so copy if necessary
-      if (DMT::GetRawHostPtr(*B_) != DMT::GetRawHostPtr(*X_))
+      if (B_ != X_)
         DMT::Assign(*X_, *B_); // Copy B to X if needed
 
       // Since B_ = X_, perform operations on X_.
-      DMT::SyncHostToDevice(*X_);
 
       int M = DMT::GetNumRows(*X_); 
       int NRHS = DMT::GetNumCols(*X_);
@@ -245,7 +233,6 @@ namespace Belos {
           ptr = X + j*LDX;
           for (int i=0; i<M; i++) {
             *ptr = *ptr*R_tmp[i];
-            //std::cout << "B[ " << i << " ] = " << *ptr << std::endl; 
             ptr++;
           }
         } 
@@ -272,7 +259,6 @@ namespace Belos {
           ptr = X + j*LDX; 
             for (int i=0; i<M; i++) {
             *ptr = *ptr*R_[i];
-            //std::cout << "X[ " << i << " ] = " << *ptr << std::endl;
             ptr++;
           }
         }

diff --git a/packages/belos/src/BelosPseudoBlockCGIter.hpp b/packages/belos/src/BelosPseudoBlockCGIter.hpp
@@ -57,7 +57,6 @@
 #include "BelosOperatorTraits.hpp"
 #include "BelosMultiVecTraits.hpp"
 #include "BelosDenseMatTraits.hpp"
-#include "BelosTeuchosDenseAdapter.hpp"
 
 #include "Teuchos_ScalarTraits.hpp"
 #include "Teuchos_ParameterList.hpp"
@@ -76,7 +75,7 @@
 
 namespace Belos {
 
-  template<class ScalarType, class MV, class OP, class DM = Teuchos::SerialDenseMatrix<int, ScalarType>>
+  template<class ScalarType, class MV, class OP, class DM>
   class PseudoBlockCGIter : virtual public CGIteration<ScalarType,MV,OP,DM> {
 
   public:

diff --git a/packages/belos/tpetra/test/BlockCG/CMakeLists.txt b/packages/belos/tpetra/test/BlockCG/CMakeLists.txt
@@ -22,6 +22,7 @@ TRIBITS_ADD_EXECUTABLE_AND_TEST(
   Tpetra_KokkosDense_BlockCG_hb_test
   SOURCES test_kokkos_bl_cg_hb.cpp
   ARGS "--verbose"
+       "--verbose --use-single-red"
        "--verbose --num-rhs=2 --block-size=2"
   COMM serial mpi
   ARGS

diff --git a/packages/belos/tpetra/test/BlockCG/test_kokkos_bl_cg_hb.cpp b/packages/belos/tpetra/test/BlockCG/test_kokkos_bl_cg_hb.cpp
@@ -98,6 +98,7 @@ int main(int argc, char *argv[]) {
     //
     bool proc_verbose = false;
     bool debug = false;
+    bool use_single_red = false;
     int frequency = -1;  // how often residuals are printed by solver
     int numrhs = 1;      // total number of right-hand sides to solve for
     int blocksize = 1;   // blocksize used by solver
@@ -114,6 +115,7 @@ int main(int argc, char *argv[]) {
     cmdp.setOption("num-rhs",&numrhs,"Number of right-hand sides to be solved for.");
     cmdp.setOption("max-iters",&maxiters,"Maximum number of iterations per linear system (-1 := adapted to problem/block size).");
     cmdp.setOption("block-size",&blocksize,"Block size to be used by the CG solver.");
+    cmdp.setOption("use-single-red","use-standard-red",&use_single_red,"Use single-reduction CG iteration.");
     if (cmdp.parse(argc,argv) != CommandLineProcessor::PARSE_SUCCESSFUL) {
       return -1;
     }
@@ -158,6 +160,9 @@ int main(int argc, char *argv[]) {
     belosList.set( "Block Size", blocksize );              // Blocksize to be used by iterative solver
     belosList.set( "Maximum Iterations", maxiters );       // Maximum number of iterations allowed
     belosList.set( "Convergence Tolerance", tol );         // Relative convergence tolerance requested
+    if ((blocksize==1) && use_single_red)
+      belosList.set( "Use Single Reduction", use_single_red ); // Use single reduction CG iteration
+
     int verbLevel = Belos::Errors + Belos::Warnings;
     if (debug) {
       verbLevel += Belos::Debug;