cherry pick Kokkos-kernels PR trilinos#921

Two-stage GS: add damping factors trilinos#921
iyamazaki · Apr 1, 2021 · c60c502 · c60c502
1 parent 2dd64d2
commit c60c502
Show file tree

Hide file tree

Showing 4 changed files with 576 additions and 237 deletions.
diff --git a/packages/kokkos-kernels/src/common/KokkosKernels_Handle.hpp b/packages/kokkos-kernels/src/common/KokkosKernels_Handle.hpp
@@ -580,12 +580,24 @@ class KokkosKernelsHandle
     return gs2;
   }
   // ---------------------------------------- //
-  // Specify to use either Two-stage or Classical (i.e., inner Jacobi-Richardson or SpTrsv)
+  // Specify numer of outer sweeps for two-stage Gauss-Seidel
+  void set_gs_set_num_outer_sweeps (int num_outer_sweeps) {
+    auto gs2 = get_twostage_gs_handle();
+    gs2->setNumOuterSweeps (num_outer_sweeps);
+  }
+  // ---------------------------------------- //
+  // Specify numer of inner sweeps for two-stage Gauss-Seidel
   void set_gs_set_num_inner_sweeps (int num_inner_sweeps) {
     auto gs2 = get_twostage_gs_handle();
     gs2->setNumInnerSweeps (num_inner_sweeps);
   }
   // ---------------------------------------- //
+  // Specify damping factor of inner sweeps for two-stage Gauss-Seidel
+  void set_gs_set_inner_damp_factor (nnz_scalar_t damp_factor) {
+    auto gs2 = get_twostage_gs_handle();
+    gs2->setInnerDampFactor (damp_factor);
+  }
+  // ---------------------------------------- //
   // Specify to use either Two-stage or Classical (i.e., inner Jacobi-Richardson or SpTrsv)
   void set_gs_twostage (bool two_stage, size_type nrows) {
     auto gs2 = get_twostage_gs_handle();
@@ -608,6 +620,13 @@ class KokkosKernelsHandle
       }
     }
   }
+  // ---------------------------------------- //
+  // Specify to use either Compact or Classical form of recurrence
+  void set_gs_twostage_compact_form (bool compact_form) {
+    auto gs2 = get_twostage_gs_handle();
+    gs2->setCompactForm (compact_form);
+  }
+
 
   void create_gs_handle(KokkosSparse::ClusteringAlgorithm clusterAlgo, nnz_lno_t hint_verts_per_cluster) {
     this->destroy_gs_handle();

diff --git a/packages/kokkos-kernels/src/sparse/KokkosSparse_gauss_seidel_handle.hpp b/packages/kokkos-kernels/src/sparse/KokkosSparse_gauss_seidel_handle.hpp
@@ -573,8 +573,13 @@ namespace KokkosSparse{
     nrhs (1),
     direction (GS_SYMMETRIC),
     two_stage (true),
-    num_inner_sweeps (1)
-    {}
+    compact_form (false),
+    num_inner_sweeps (1),
+    num_outer_sweeps (1)
+    {
+      const scalar_t one (1.0);
+      inner_omega = one;
+    }
 
     // Sweep direction
     void setSweepDirection (GSDirection direction_) {
@@ -592,6 +597,22 @@ namespace KokkosSparse{
       return this->two_stage;
     }
 
+    // specify whether to use compact form of recurrence
+    void setCompactForm (bool compact_form_) {
+      this->compact_form = compact_form_;
+    }
+    bool isCompactForm () {
+      return this->compact_form;
+    }
+
+    // Number of outer sweeps
+    void setNumOuterSweeps (int num_outer_sweeps_) {
+      this->num_outer_sweeps = num_outer_sweeps_;
+    }
+    int getNumOuterSweeps () {
+      return this->num_outer_sweeps;
+    }
+
     // Number of inner sweeps
     void setNumInnerSweeps (int num_inner_sweeps_) {
       this->num_inner_sweeps = num_inner_sweeps_;
@@ -600,27 +621,57 @@ namespace KokkosSparse{
       return this->num_inner_sweeps;
     }
 
-    // workspaces
+    // Inner damping factor
+    void setInnerDampFactor (scalar_t inner_omega_) {
+      this->inner_omega = inner_omega_;
+    }
+    scalar_t getInnerDampFactor () {
+      return this->inner_omega;
+    }
+
+    // Workspaces
+    // > diagonal (inverse)
     void setD (values_view_t D_) {
       this->D = D_;
     }
     values_view_t getD () {
       return this->D;
     }
-
+    // > Lower part of diagonal block
     void setL (crsmat_t L) {
       this->crsmatL = L;
     }
     crsmat_t getL () {
       return this->crsmatL;
     }
-
+    // > Upper part of diagonal block
     void setU (crsmat_t U) {
       this->crsmatU = U;
     }
     crsmat_t getU () {
       return this->crsmatU;
     }
+    // > Complement of U
+    void setLa (crsmat_t La) {
+      this->crsmatLa = La;
+    }
+    crsmat_t getLa () {
+      return this->crsmatLa;
+    }
+    // > Complement of L
+    void setUa (crsmat_t Ua) {
+      this->crsmatUa = Ua;
+    }
+    crsmat_t getUa () {
+      return this->crsmatUa;
+    }
+    // > diagonal (not-inverse)
+    void setDa (values_view_t Da_) {
+      this->Da = Da_;
+    }
+    values_view_t getDa () {
+      return this->Da;
+    }
 
     void initVectors (int nrows_, int nrhs_) {
       if (this->nrows != nrows_ || this->nrhs != nrhs_) {
@@ -650,6 +701,11 @@ namespace KokkosSparse{
     values_view_t D;
     crsmat_t crsmatL;
     crsmat_t crsmatU;
+    // > complements for compact form of recurrence
+    //   where La = A - U and Ua = A - L
+    values_view_t Da;
+    crsmat_t crsmatLa;
+    crsmat_t crsmatUa;
 
     // > residual vector for outer GS, Rk = B-A*Xk
     vector_view_t localR;
@@ -661,7 +717,10 @@ namespace KokkosSparse{
     // solver parameters
     GSDirection direction;
     bool two_stage;
+    bool compact_form;
     int num_inner_sweeps;
+    int num_outer_sweeps;
+    scalar_t inner_omega;
   };
   // -------------------------------------
 }