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

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_discretization.cpp

@@ -1,4 +1,5 @@
 #include "asgard_discretization.hpp"
+#include "asgard_small_mats.hpp"
 
 #ifdef ASGARD_USE_GPU
 #include "asgard_gpu_algorithms.hpp"
@@ -222,16 +223,39 @@ void discretization_manager<precision>::start_moments() {
     #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);
+      if (this->poisson_preconditioner.method() == precon_method::jacobi) {
+        return this->poisson_iter.solve([&](precision y[]) -> void
+          {
+            tools::time_event timing_("poisson_md jacobi preconditioner");
+            gpu::jacobi_apply(this->poisson_preconditioner.gpu_jacobi(), y);
+          }, apply_lhs, rhs, x);
+      }
+      return this->poisson_iter.solve(nullptr, 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);
+      if (this->poisson_preconditioner.method() == precon_method::jacobi) {
+        return this->poisson_iter.solve([&](precision y[]) -> void
+          {
+            tools::time_event timing_("poisson_md jacobi preconditioner");
+            fm::jacobi_apply(x.size(), this->poisson_preconditioner.jacobi(), y);
+          }, apply_lhs, rhs, x);
+      }
+      return this->poisson_iter.solve(nullptr, 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.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_preconditioner);
+      }
+    }, get_poisson());
   compute_moments_(group_id::all(), state);
 }
 
@@ -407,7 +431,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() : poisson_iter(1e-6, 10000), poisson_preconditioner(precon_method::jacobi)
   {
     #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_preconditioner);
+      }
+    }, 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_preconditioner);
+      }
+    }, get_poisson());
   }
   //! refines the sparse grid using the given strategy and
   void refine(sparse_grid::strategy mode, gpu::vector<precision> const &f)
@@ -843,7 +847,8 @@ class discretization_manager
   std::vector<aux_field_entry<precision>> aux_fields;
 
   //! Just testing for now
-  solvers::bicgstab<precision> poisson_iter;
+  solvers::cg<precision> poisson_iter;
+  mutable preconditioner_data<precision> poisson_preconditioner;
 };
 
 } // namespace asgard

src/asgard_gpu_algorithms.hpp

@@ -57,4 +57,28 @@ 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()); }
 
+//! alpha = rho / p dot q
+template<typename P>
+void cg_calc_alpha(P const *rho, P const *p_dot_q, P *alpha);
+
+//! beta = rho_new / rho
+template<typename P>
+void cg_calc_beta(P const *rho_new, P const *rho, P *beta);
+
+//! x = x + alpha * p, r = r - alpha * q
+template<typename P>
+void cg_update_x_r(int64_t n, P const *alpha, 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 *beta, P const r[], P p[]);
+
+//! rho = rho_new
+template<typename P>
+void cg_update_rho(P const *rho_new, P *rho);
+
+//! x = 0
+template<typename P>
+void set_zero(P *x);
+
 }

src/asgard_gpu_algorithms_cuda.cpp

@@ -235,6 +235,73 @@ int num_non_finite(int64_t num, P const x[]) {
   return cpu_res;
 }
 
+template<typename P>
+__global__ void cg_calc_alpha_kernel(P const *rho, P const *p_dot_q, P *alpha) {
+  *alpha = *rho / *p_dot_q;
+}
+
+template<typename P>
+__global__ void cg_calc_beta_kernel(P const *rho_new, P const *rho, P *beta) {
+  *beta = *rho_new / *rho;
+}
+
+template<typename P>
+__global__ void cg_update_x_r_kernel(int64_t n, P const *alpha, P const p[], P const q[], P x[], P r[]) {
+  int64_t i = blockIdx.x * blockDim.x + threadIdx.x;
+  if (i < n) {
+    P a = *alpha;
+    x[i] += a * p[i];
+    r[i] -= a * q[i];
+  }
+}
+
+template<typename P>
+__global__ void cg_update_p_kernel(int64_t n, P const *beta, P const r[], P p[]) {
+  int64_t i = blockIdx.x * blockDim.x + threadIdx.x;
+  if (i < n) {
+    p[i] = r[i] + (*beta) * p[i];
+  }
+}
+
+template<typename P>
+__global__ void cg_update_rho_kernel(P const *rho_new, P *rho) {
+  *rho = *rho_new;
+}
+
+template<typename P>
+__global__ void zero_kernel(P *x) {
+  *x = 0;
+}
+
+template<typename P>
+void cg_calc_alpha(P const *rho, P const *p_dot_q, P *alpha) {
+  cg_calc_alpha_kernel<<<1, 1>>>(rho, p_dot_q, alpha);
+}
+template<typename P>
+void cg_calc_beta(P const *rho_new, P const *rho, P *beta) {
+  cg_calc_beta_kernel<<<1, 1>>>(rho_new, rho, beta);
+}
+template<typename P>
+void cg_update_x_r(int64_t n, P const *alpha, P const p[], P const q[], P x[], P r[]) {
+  int const threads = 1024;
+  int const blocks = (n + threads - 1) / threads;
+  cg_update_x_r_kernel<<<blocks, threads>>>(n, alpha, p, q, x, r);
+}
+template<typename P>
+void cg_update_p(int64_t n, P const *beta, P const r[], P p[]) {
+  int const threads = 1024;
+  int const blocks = (n + threads - 1) / threads;
+  cg_update_p_kernel<<<blocks, threads>>>(n, beta, r, p);
+}
+template<typename P>
+void cg_update_rho(P const *rho_new, P *rho) {
+  cg_update_rho_kernel<<<1, 1>>>(rho_new, rho);
+}
+template<typename P>
+void set_zero(P *x) {
+  zero_kernel<<<1, 1>>>(x);
+}
+
 #ifdef ASGARD_ENABLE_DOUBLE
 template void jacobi_apply<double>(gpu::vector<double> const &, double[]);
 
@@ -255,6 +322,13 @@ 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_calc_alpha<double>(double const*, double const*, double*);
+template void cg_calc_beta<double>(double const*, double const*, double*);
+template void cg_update_x_r<double>(int64_t, double const*, double const*, double const*, double*, double*);
+template void cg_update_p<double>(int64_t, double const*, double const*, double*);
+template void cg_update_rho<double>(double const*, double*);
+template void set_zero<double>(double*);
 #endif
 
 #ifdef ASGARD_ENABLE_FLOAT
@@ -277,5 +351,12 @@ 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_calc_alpha<float>(float const*, float const*, float*);
+template void cg_calc_beta<float>(float const*, float const*, float*);
+template void cg_update_x_r<float>(int64_t, float const*, float const*, float const*, float*, float*);
+template void cg_update_p<float>(int64_t, float const*, float const*, float*);
+template void cg_update_rho<float>(float const*, float*);
+template void set_zero<float>(float*);
 #endif
 }

src/asgard_moment_manager.cpp

@@ -510,8 +510,7 @@ void moment_manager<P>::mcompute(sparse_grid const &grid, moment_id id,
 }
 
 template<typename P>
-void moment_manager<P>::mcompute(sparse_grid const &grid, moment_id id,
-                                 std::vector<P> const &state, std::vector<P> &vals) const
+void moment_manager<P>::update_position_grid(sparse_grid const &grid) const
 {
   if (pos_grid.generation() != grid.generation()) { // grid changed, must rebuild
     switch (pos_grid.num_dims()) {
@@ -529,6 +528,13 @@ void moment_manager<P>::mcompute(sparse_grid const &grid, moment_id id,
     };
     pos_grid.generation_ = grid.generation();
   }
+}
+
+template<typename P>
+void moment_manager<P>::mcompute(sparse_grid const &grid, moment_id id,
+                                 std::vector<P> const &state, std::vector<P> &vals) const
+{
+  update_position_grid(grid);
 
   switch (num_vel_) {
   case 1:
@@ -634,7 +640,7 @@ void moment_manager<P>::complete_level(hierarchy_manipulator<P> const &hier,
 template<typename P>
 void moment_manager<P>::solve_poisson(sparse_grid const &grid, connection_patterns const &conn,
                                       hierarchy_manipulator<P> const &hier, interpolation_manager<P> const &interp,
-                                      kronmult::workspace<P> &work) const
+                                      kronmult::workspace<P> &work, bool raw_on_cpu) const
 {
   std::visit([&](auto &p) {
       if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_1d<P>>) {
@@ -651,7 +657,7 @@ void moment_manager<P>::solve_poisson(sparse_grid const &grid, connection_patter
           gpu::vector<P> &res = gpu_interps[0][mid];
           res.resize(full_block * grid.num_indexes());
         }
-        solve_poisson_gpu(conn, hier, interp, work);
+        solve_poisson_gpu(conn, hier, interp, work, raw_on_cpu);
         auto find_mom = [this](moment_id mid) -> const mom_on_gpu &
         {
           for (auto const &mom : this->gpu_moments[0]) {
@@ -670,6 +676,7 @@ void moment_manager<P>::solve_poisson(sparse_grid const &grid, connection_patter
         #else
         std::ignore = interp;
         std::ignore = grid;
+        std::ignore = raw_on_cpu;
         update_position_grid_dsort();
         p.solve_periodic(raw_vals.get(p.moment0()), raw_vals, pos_grid, conn, work);
         #endif
@@ -680,7 +687,7 @@ void moment_manager<P>::solve_poisson(sparse_grid const &grid, connection_patter
 #ifdef ASGARD_USE_GPU
 template<typename P>
 void moment_manager<P>::solve_poisson_gpu(connection_patterns const &conn, hierarchy_manipulator<P> const &hier,
-                                          interpolation_manager<P> const &interp, kronmult::workspace<P> &work) const
+                                          interpolation_manager<P> const &interp, kronmult::workspace<P> &work, bool raw_on_cpu) const
 {
   std::visit([&](auto &p) {
       if constexpr (std::is_same_v<std::decay_t<decltype(p)>, poisson_1d<P>>) {
@@ -706,6 +713,10 @@ void moment_manager<P>::solve_poisson_gpu(connection_patterns const &conn, hiera
         };
         auto interp_func = [&](gpu::vector<P> const &efield, moment_id mid) -> void
         {
+          if (raw_on_cpu) {
+            efield.copy_to_host(this->raw_vals[mid]);
+            this->full_level.get(mid).resize(0);
+          }
           interp.pos2nodal(gpu::device{0}, this->pos_grid, efield.data(), this->wav_scale, w2.vec.data(), work);
           gpu::vector<P> &res = this->gpu_interps[0][mid];
           moment_expand(pdof, this->pos_grid.num_dims(), num_vel_, reduce_ij[0], w2.vec, res);