Preconditioned Conjugate Gradient Solver

CMakeLists.txt

@@ -129,7 +129,7 @@ endif()
 # to configure the examples cmake file, so it needs to be a single string as opposed to a list
 # on the other hand, the additional internal (testing) PDEs are just a normal list
 set (_asgard_doc_pdes "continuity_2d inputs_1d continuity diffusion spherical_diffusion elliptic elliptic_nonsep robin_bc sinwav two_stream vplb sod_shock burgers bgk")
-set (_asgard_gpu_pdes "burgers_gpu two_stream_md")
+set (_asgard_gpu_pdes "burgers_gpu two_stream_md landau_damping")
 set (_asgard_int_pdes bound_internal interp_internal interp_bc_internal multichain_internal relaxation_internal varcoeff_internal misc_internal)
 
 set (_asgard_doc_pdes_str "${_asgard_doc_pdes}") # save for examples configuration
@@ -406,6 +406,7 @@ set(asgard_components
     grid_1d
     indexset
     interp
+    iterative_solver
     kronmult
     kronmult_common
     kronmult_cpu
@@ -414,14 +415,14 @@ set(asgard_components
     momentset
     pde
     pde_functions
+    poisson
     program_options
     quadrature
     reconstruct
     refinement
     resource_groups
     small_mats
     solver
-    poisson
     term_build
     term_manager
     term_sources

src/asgard_compute.hpp

@@ -296,6 +296,21 @@ class __signleton_compute_resources {
     P const n = blas_.nrm2(num, x);
     return n;
   }
+  //! dot product between two vectors, result in VRAM
+  template<typename P>
+  void dot_device(int num, P const x[], P const y[], P* result_dev) const {
+    blas_.dot_device(num, x, y, result_dev);
+  }
+  //! dot product of a vector with itself, result in VRAM
+  template<typename P>
+  void dot1_device(int num, P const x[], P* result_dev) const {
+    blas_.dot_device(num, x, x, result_dev);
+  }
+  //! norm-2 of a vector, result in VRAM
+  template<typename P>
+  void nrm2_device(int num, P const x[], P* result_dev) const {
+    blas_.nrm2_device(num, x, result_dev);
+  }
   //! transpose of a matrix, times a vector
   template<typename P>
   void gemtv(int m, int n, no_deduce<P> alpha, P const A[], P const x[],

src/asgard_compute_cuda.hpp

@@ -184,6 +184,38 @@ class blas_engine
     }
   }
 
+  template<typename P>
+  void dot_device(int num, P const x[], P const y[], P* result_dev) const {
+      // 1. Tell CUBLAS to write the result to a pointer on the DEVICE
+      cublasSetPointerMode(cublas, CUBLAS_POINTER_MODE_DEVICE);
+
+      // 2. Perform the dot product
+      if constexpr (std::is_same_v<P, double>) {
+          cublasDdot(cublas, num, x, 1, y, 1, result_dev);
+      } else {
+          cublasSdot(cublas, num, x, 1, y, 1, result_dev);
+      }
+
+      // 3. Reset to HOST mode for standard ASGarD operations
+      cublasSetPointerMode(cublas, CUBLAS_POINTER_MODE_HOST);
+  }
+
+  template<typename P>
+  void nrm2_device(int num, P const x[], P* result_dev) const {
+      // 1. Tell CUBLAS to write the result to a pointer on the DEVICE
+      cublasSetPointerMode(cublas, CUBLAS_POINTER_MODE_DEVICE);
+
+      // 2. Perform the norm calculation
+      if constexpr (std::is_same_v<P, double>) {
+          cublasDnrm2(cublas, num, x, 1, result_dev);
+      } else {
+          cublasSnrm2(cublas, num, x, 1, result_dev);
+      }
+
+      // 3. Reset to HOST mode
+      cublasSetPointerMode(cublas, CUBLAS_POINTER_MODE_HOST);
+  }
+
   template<typename P>
   void gemtv(int m, int n, P alpha, P const A[], P const x[], P beta, P y[]) const {
     if constexpr (is_float<P>) {

src/asgard_compute_rocm.hpp

@@ -169,6 +169,38 @@ class blas_engine
     }
   }
 
+  template<typename P>
+  void dot_device(int num, P const x[], P const y[], P* result_dev) const {
+      // 1. Tell rocBLAS to write the result to a pointer on the DEVICE
+      rocblas_set_pointer_mode(rocblas, rocblas_pointer_mode_device);
+
+      // 2. Perform the dot product
+      if constexpr (std::is_same_v<P, double>) {
+          rocblas_ddot(rocblas, num, x, 1, y, 1, result_dev);
+      } else {
+          rocblas_fdot(rocblas, num, x, 1, y, 1, result_dev);
+      }
+
+      // 3. Reset to HOST mode for standard ASGarD operations
+      rocblas_set_pointer_mode(rocblas, rocblas_pointer_mode_host);
+  }
+
+  template<typename P>
+  void nrm2_device(int num, P const x[], P* result_dev) const {
+      // 1. Tell rocBLAS to write the result to a pointer on the DEVICE
+      rocblas_set_pointer_mode(rocblas, rocblas_pointer_mode_device);
+
+      // 2. Perform the norm calculation
+      if constexpr (std::is_same_v<P, double>) {
+          cublas_dnrm2(rocblas, num, x, 1, result_dev);
+      } else {
+          cublas_snrm2(rocblas, num, x, 1, result_dev);
+      }
+
+      // 3. Reset to HOST mode
+      rocblas_set_pointer_mode(rocblas, rocblas_pointer_mode_host);
+  }
+
   template<typename P>
   void gemtv(int m, int n, P alpha, P const A[], P const x[], P beta, P y[]) const {
     if constexpr (is_float<P>) {

src/asgard_discretization.cpp

@@ -1,4 +1,5 @@
 #include "asgard_discretization.hpp"
+#include "asgard_small_mats.hpp"
 
 #ifdef ASGARD_USE_GPU
 #include "asgard_gpu_algorithms.hpp"
@@ -219,19 +220,17 @@ void discretization_manager<precision>::start_moments() {
     auto build_func = [this](term_entry<precision> &tentry, int const dim, int const level) -> void {
       this->terms.rebuild_term1d(tentry, dim, level);
     };
-    #ifdef ASGARD_USE_GPU
-    auto iter_solve_func = [this](solvers::operation_apply_lhs<precision> apply_lhs,
-                                  gpu::vector<precision> const &rhs, gpu::vector<precision> &x) -> int {
-      return this->poisson_iter.solve(apply_lhs, rhs, x);
-    };
-    #else
-    auto iter_solve_func = [this](solvers::operation_apply_lhs<precision> apply_lhs,
-                                  std::vector<precision> const &rhs, std::vector<precision> &x) -> int {
-      return this->poisson_iter.solve(apply_lhs, rhs, x);
-    };
-    #endif
-    terms.moms.set_poisson(terms.max_level, terms.grid, terms.xleft, terms.xright, terms.conn, terms.hier, build_func, iter_solve_func);
+    terms.moms.set_poisson(terms.max_level, terms.grid, terms.xleft, terms.xright, terms.conn, terms.hier, build_func);
   }
+  terms.moms.update_position_grid(terms.grid);
+  std::visit([&](auto &p)
+    {
+      if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_1d<precision>>) {
+        p.update_level(terms.grid.current_level(0));
+      } else if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_md<precision>>) {
+        p.update_preconditioner(terms.moms.get_position_grid(), terms.conn, poisson_bc::periodic);
+      }
+    }, get_poisson());
   compute_moments_(group_id::all(), state);
 }
 
@@ -407,7 +406,7 @@ std::vector<precision> discretization_manager<precision>::get_moment(moment_id i
       else
         throw std::runtime_error("an electric moment was requested but no poisson solver is set");
     }, get_poisson());
-    terms.moms.solve_poisson(terms.grid, terms.conn, terms.hier, terms.interp, terms.kwork);
+    terms.moms.solve_poisson(terms.grid, terms.conn, terms.hier, terms.interp, terms.kwork, true);
     return terms.moms.get_cached_raw(id);
   }
   std::vector<precision> result;

src/asgard_discretization.hpp

@@ -33,7 +33,7 @@ class discretization_manager
 {
 public:
   //! allows the creation of a null manager, has to be reinitialized later
-  discretization_manager() : poisson_iter(1e-6, 10000)
+  discretization_manager()
   {
     #ifdef ASGARD_ENABLE_DOUBLE
     #ifdef ASGARD_ENABLE_FLOAT
@@ -748,13 +748,15 @@ class discretization_manager
     if (is_leader())
       terms.grid.remap(terms.block_size(), next);
     terms.prapare_kron_workspace();
-    if (has_poisson())
-      std::visit([&](auto &p)
-      {
-        if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_1d<precision>>) {
-          p.update_level(terms.grid.current_level(0));
-        }
-      }, get_poisson());
+    std::visit([&](auto &p)
+    {
+      if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_1d<precision>>) {
+        p.update_level(terms.grid.current_level(0));
+      } else if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_md<precision>>) {
+        terms.moms.update_position_grid(terms.grid);
+        p.update_preconditioner(terms.moms.get_position_grid(), terms.conn, poisson_bc::periodic);
+      }
+    }, get_poisson());
   }
   //! refines the sparse grid using the given strategy and
   void refine(sparse_grid::strategy mode, std::vector<precision> const &f)
@@ -770,13 +772,15 @@ class discretization_manager
     if (is_leader())
       terms.grid.remap(terms.block_size(), next);
     terms.prapare_kron_workspace();
-    if (has_poisson())
-      std::visit([&](auto &p)
-      {
-        if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_1d<precision>>) {
-          p.update_level(terms.grid.current_level(0));
-        }
-      }, get_poisson());
+    std::visit([&](auto &p)
+    {
+      if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_1d<precision>>) {
+        p.update_level(terms.grid.current_level(0));
+      } else if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_md<precision>>) {
+        terms.moms.update_position_grid(terms.grid);
+        p.update_preconditioner(terms.moms.get_position_grid(), terms.conn, poisson_bc::periodic);
+      }
+    }, get_poisson());
   }
   //! refines the sparse grid using the given strategy and
   void refine(sparse_grid::strategy mode, gpu::vector<precision> const &f)
@@ -841,9 +845,6 @@ class discretization_manager
 
   //! fields to store and save for plotting
   std::vector<aux_field_entry<precision>> aux_fields;
-
-  //! Just testing for now
-  solvers::bicgstab<precision> poisson_iter;
 };
 
 } // namespace asgard

src/asgard_gpu_algorithms.hpp

@@ -57,4 +57,12 @@ int num_non_finite(int64_t num, P const x[]);
 template<typename P>
 int num_non_finite(gpu::vector<P> const &data) { return num_non_finite(data.size(), data.data()); }
 
+//! x = x + alpha * p, r = r - alpha * q
+template<typename P>
+void cg_update_x_r(int64_t n, P const *rho, P const *p_dot_q, P const p[], P const q[], P x[], P r[]);
+
+//! p = r + beta * p
+template<typename P>
+void cg_update_p(int64_t n, P const *rho_new, P const *rho, P const r[], P p[]);
+
 }

src/asgard_gpu_algorithms_cuda.cpp

@@ -235,6 +235,40 @@ int num_non_finite(int64_t num, P const x[]) {
   return cpu_res;
 }
 
+template<typename P, int num_threads>
+__global__ void cg_update_x_r_kernel(int64_t num, P const *rho, P const *p_dot_q, P const p[], P const q[], P x[], P r[]) {
+  int i = threadIdx.x + blockIdx.x * num_threads;
+  P alpha = *rho / *p_dot_q;
+  while (i < num) {
+    x[i] += alpha * p[i];
+    r[i] -= alpha * q[i];
+    i += num_threads * gridDim.x;
+  }
+}
+
+template<typename P, int num_threads>
+__global__ void cg_update_p_kernel(int64_t num, P const *rho_new, P const *rho, P const r[], P p[]) {
+  int i = threadIdx.x + blockIdx.x * num_threads;
+  P beta = *rho_new / *rho;
+  while (i < num) {
+    p[i] = r[i] + beta * p[i];
+    i += num_threads * gridDim.x;
+  }
+}
+
+template<typename P>
+void cg_update_x_r(int64_t num, P const *rho, P const *p_dot_q, P const p[], P const q[], P x[], P r[]) {
+  constexpr int max_threads = 1024;
+  int const num_blocks = round_up(num, max_threads);
+  cg_update_x_r_kernel<P, max_threads><<<num_blocks, max_threads>>>(num, rho, p_dot_q, p, q, x, r);
+}
+template<typename P>
+void cg_update_p(int64_t num, P const *rho_new, P const *rho, P const r[], P p[]) {
+  constexpr int max_threads = 1024;
+  int const num_blocks = round_up(num, max_threads);
+  cg_update_p_kernel<P, max_threads><<<num_blocks, max_threads>>>(num, rho_new, rho, r, p);
+}
+
 #ifdef ASGARD_ENABLE_DOUBLE
 template void jacobi_apply<double>(gpu::vector<double> const &, double[]);
 
@@ -255,6 +289,9 @@ template void sum5(gpu::vector<double> const &, double, gpu::vector<double> cons
                    double, gpu::vector<double> const &, double, gpu::vector<double> const &, gpu::vector<double> &);
 
 template int num_non_finite(int64_t num, double const x[]);
+
+template void cg_update_x_r<double>(int64_t, double const*, double const*, double const[], double const[], double[], double[]);
+template void cg_update_p<double>(int64_t, double const*, double const*, double const[], double[]);
 #endif
 
 #ifdef ASGARD_ENABLE_FLOAT
@@ -277,5 +314,8 @@ template void sum5(gpu::vector<float> const &, float, gpu::vector<float> const &
                    float, gpu::vector<float> const &, float, gpu::vector<float> const &, gpu::vector<float> &);
 
 template int num_non_finite(int64_t num, float const x[]);
+
+template void cg_update_x_r<float>(int64_t, float const*, float const*, float const[], float const[], float[], float[]);
+template void cg_update_p<float>(int64_t, float const*, float const*, float const[], float[]);
 #endif
 }

src/asgard_gpu_algorithms_rocm.cpp

@@ -415,6 +415,40 @@ int num_non_finite(int64_t num, P const x[]) {
   return cpu_res;
 }
 
+template<typename P, int num_threads>
+__global__ void cg_update_x_r_kernel(int64_t num, P const *rho, P const *p_dot_q, P const p[], P const q[], P x[], P r[]) {
+  int i = threadIdx.x + blockIdx.x * num_threads;
+  P alpha = *rho / *p_dot_q;
+  while (i < num) {
+    x[i] += alpha * p[i];
+    r[i] -= alpha * q[i];
+    i += num_threads * gridDim.x;
+  }
+}
+
+template<typename P, int num_threads>
+__global__ void cg_update_p_kernel(int64_t num, P const *rho_new, P const *rho, P const r[], P p[]) {
+  int i = threadIdx.x + blockIdx.x * num_threads;
+  P beta = *rho_new / *rho;
+  while (i < num) {
+    p[i] = r[i] + beta * p[i];
+    i += num_threads * gridDim.x;
+  }
+}
+
+template<typename P>
+void cg_update_x_r(int64_t num, P const *rho, P const *p_dot_q, P const p[], P const q[], P x[], P r[]) {
+  constexpr int max_threads = 1024;
+  int const num_blocks = round_up(num, max_threads);
+  cg_update_x_r_kernel<P, max_threads><<<num_blocks, max_threads>>>(num, rho, p_dot_q, p, q, x, r);
+}
+template<typename P>
+void cg_update_p(int64_t num, P const *rho_new, P const *rho, P const r[], P p[]) {
+  constexpr int max_threads = 1024;
+  int const num_blocks = round_up(num, max_threads);
+  cg_update_p_kernel<P, max_threads><<<num_blocks, max_threads>>>(num, rho_new, rho, r, p);
+}
+
 #ifdef ASGARD_ENABLE_DOUBLE
 template void jacobi_apply<double>(gpu::vector<double> const &, double[]);
 template void compute_last_bicgstab<double>(double, double, gpu::vector<double> const &,
@@ -436,6 +470,9 @@ template void tensor_by_index(int, int, int, int const[], double const[], double
                               double const[], double const[], double const[], double[]);
 
 template int num_non_finite(int64_t num, double const x[]);
+
+template void cg_update_x_r<double>(int64_t, double const*, double const*, double const[], double const[], double[], double[]);
+template void cg_update_p<double>(int64_t, double const*, double const*, double const[], double[]);
 #endif
 
 #ifdef ASGARD_ENABLE_FLOAT
@@ -459,5 +496,8 @@ template void tensor_by_index(int, int, int, int const[], float const[], float c
                               float const[], float const[], float const[], float[]);
 
 template int num_non_finite(int64_t num, float const x[]);
+
+template void cg_update_x_r<float>(int64_t, float const*, float const*, float const[], float const[], float[], float[]);
+template void cg_update_p<float>(int64_t, float const*, float const*, float const[], float[]);
 #endif
 }