[WIP] Implement HYPRE solver on top of diffusion example

cmake/FindHYPRE.cmake

@@ -0,0 +1,70 @@
+#[=======================================================================[.rst:
+FindHYPRE
+-------
+
+Finds the HYPRE library.
+
+Imported Targets
+^^^^^^^^^^^^^^^^
+
+This module provides the following imported targets, if found:
+
+``HYPRE::HYPRE``
+  The HYPRE library
+
+We will try looking in the HYPRE_DIR user provided path in site.cmake
+
+Result Variables
+^^^^^^^^^^^^^^^^
+
+This will define the following variables:
+
+``HYPRE_FOUND``
+  True if the system has the HYPRE library.
+``HYPRE_VERSION``
+  The version of the HYPRE library which was found.
+``HYPRE_INCLUDE_DIRS``
+  Include directories needed to use HYPRE.
+``HYPRE_LIBRARIES``
+  Libraries needed to link to HYPRE.
+
+Cache Variables
+^^^^^^^^^^^^^^^
+
+The following cache variables may also be set:
+
+``HYPRE_INCLUDE_DIR``
+  The directory containing ``foo.h``.
+``HYPRE_LIBRARY``
+  The path to the HYPRE library.
+
+#]=======================================================================]
+
+
+find_path(HYPRE_INCLUDE_DIR NAMES HYPRE.h HINTS ${HYPRE_DIR}/include)
+find_library(HYPRE_LIBRARY NAMES HYPRE HINTS ${HYPRE_DIR}/lib)
+
+include(FindPackageHandleStandardArgs)
+find_package_handle_standard_args(HYPRE
+  FOUND_VAR HYPRE_FOUND
+  REQUIRED_VARS
+    HYPRE_LIBRARY
+    HYPRE_INCLUDE_DIR
+)
+
+  # VERSION_VAR HYPRE_VERSION
+if(HYPRE_FOUND AND NOT TARGET HYPRE::HYPRE)
+  add_library(HYPRE::HYPRE UNKNOWN IMPORTED)
+  set_target_properties(HYPRE::HYPRE PROPERTIES
+    IMPORTED_LOCATION "${HYPRE_LIBRARY}"
+    INTERFACE_COMPILE_OPTIONS "${PC_HYPRE_CFLAGS_OTHER}"
+    INTERFACE_INCLUDE_DIRECTORIES "${HYPRE_INCLUDE_DIR}"
+  )
+endif()
+
+mark_as_advanced(
+  HYPRE_INCLUDE_DIR
+  HYPRE_LIBRARY
+)
+
+

example/diffusion/CMakeLists.txt

@@ -34,6 +34,13 @@ if( "diffusion-example" IN_LIST DRIVER_LIST OR NOT PARTHENON_DISABLE_EXAMPLES)
 
   lint_target(diffusion-example)
 
+  if (DIFFUSION_WITH_HYPRE)
+    find_package(HYPRE REQUIRED)
+    target_link_libraries(diffusion-example PRIVATE HYPRE::HYPRE)
+    target_compile_definitions(diffusion-example PRIVATE DIFFUSION_WITH_HYPRE)
+    target_sources(diffusion-example PRIVATE diffusion_hypre.cpp)
+  endif()
+
   add_custom_command(
     OUTPUT ${DOC_GEN_PATH}/diffusion-parth-table.csv
     COMMAND ${CMAKE_COMMAND} -E env

example/diffusion/diffusion_driver.cpp

@@ -19,9 +19,11 @@
 
 // Local Includes
 #include "amr_criteria/refinement_package.hpp"
+#include "basic_types.hpp"
 #include "bvals/comms/bvals_in_one.hpp"
 #include "diffusion_driver.hpp"
 #include "diffusion_equation.hpp"
+#include "diffusion_hypre.hpp"
 #include "diffusion_package.hpp"
 #include "interface/metadata.hpp"
 #include "interface/update.hpp"
@@ -32,6 +34,7 @@
 #include "solvers/cg_solver.hpp"
 #include "solvers/mg_solver.hpp"
 #include "solvers/solver_utils.hpp"
+#include "utils/error_checking.hpp"
 
 using namespace parthenon::driver::prelude;
 
@@ -55,10 +58,26 @@ TaskCollection DiffusionDriver::MakeTaskCollection() {
   TaskCollection tc;
   TaskID none(0);
 
+  auto pkg = pmesh->packages.Get("diffusion_package");
+  const auto use_hypre = pkg->Param<bool>("use_hypre");
+  if (use_hypre) {
+    return MakeTaskCollectionHypre();
+  } else {
+    return MakeTaskCollectionNative();
+  }
+}
+
+TaskCollection DiffusionDriver::MakeTaskCollectionNative() {
+  using namespace parthenon;
+  using namespace diffusion_package;
+  TaskCollection tc;
+  TaskID none(0);
+
   auto pkg = pmesh->packages.Get("diffusion_package");
   auto psolver =
       pkg->Param<std::shared_ptr<parthenon::solvers::SolverBase>>("solver_pointer");
   const auto alpha = pkg->Param<Real>("diagonal_alpha");
+  const auto rel_res = pkg->Param<Real>("rel_res");
   auto peqs = pkg->Param<std::shared_ptr<diffusion_package::DiffusionEquation<u, D>>>(
       "diffusion_equation");
 
@@ -86,9 +105,22 @@ TaskCollection DiffusionDriver::MakeTaskCollection() {
         tl.AddTask(Au, solvers::utils::AddFieldsAndStore<parthenon::TypeList<u>>, md,
                    md_rhs, md_rhs, alpha, -1.0);
 
+    // Get the RHS scale for correct comparison to Hypre solver
+    set_rhs = solvers::utils::DotProduct<parthenon::TypeList<u>>(set_rhs, tl, &u2, md, md,
+                                                                 true);
+
+    set_rhs = tl.AddTask(
+        set_rhs,
+        [alpha](parthenon::AllReduce<Real> *u2,
+                std::shared_ptr<parthenon::solvers::SolverBase> psolver, Real rel_res) {
+          *(psolver->absolute_residual_tolerance) = alpha * rel_res * sqrt(u2->val);
+          return parthenon::TaskStatus::complete;
+        },
+        &u2, psolver, rel_res);
+
     // Set initial solution guess to zero
     auto zero_u = tl.AddTask(set_rhs, TF(solvers::utils::SetToZero<u>), md_deltau);
-    psolver->initial_guess_is_zero = true;
+    psolver->initial_guess_is_zero = false;
     auto setup = psolver->AddSetupTasks(tl, zero_u, i, pmesh);
     auto solve = psolver->AddTasks(tl, setup, i, pmesh);
 
@@ -102,5 +134,77 @@ TaskCollection DiffusionDriver::MakeTaskCollection() {
   }
   return tc;
 }
+TaskCollection DiffusionDriver::MakeTaskCollectionHypre() {
+  using namespace parthenon;
+  using namespace diffusion_package;
+  TaskCollection tc;
+  TaskID none(0);
+
+#ifdef DIFFUSION_WITH_HYPRE
+
+  auto pkg = pmesh->packages.Get("diffusion_package");
+  auto hypre_solver = pkg->Param<std::shared_ptr<HypreSolver>>("hypre_solver");
+
+  TaskRegion &grid_region = tc.AddRegion(1);
+  grid_region[0].AddTask(
+      none,
+      [](HypreSolver *solver, parthenon::Mesh *pmesh) {
+        if (pmesh->modified || solver->needs_grid_setup || !solver->grid_is_setup) {
+          solver->DestroyGrid();
+          solver->SetupGrid(pmesh);
+        }
+        return TaskStatus::complete;
+      },
+      hypre_solver.get(), pmesh);
+
+  auto partitions = pmesh->GetDefaultBlockPartitions();
+  const int num_partitions = partitions.size();
+  TaskRegion &region = tc.AddRegion(num_partitions);
+  for (int i = 0; i < num_partitions; ++i) {
+    TaskList &tl = region[i];
+    auto &md = pmesh->mesh_data.Add("base", partitions[i]);
+
+    auto start_fluxcor = tl.AddTask(none, parthenon::StartReceiveFluxCorrections, md);
+
+    // SetDiffusionCoefficient
+    auto set_d = tl.AddTask(none, TF(SetDiffusionCoefficientHypre), md, tm.dt);
+
+    auto set_fluxcor = parthenon::AddFluxCorrectionTasks(set_d | start_fluxcor, tl, md,
+                                                         pmesh->multilevel);
+  }
+
+  auto &blocks = pmesh->block_list;
+  TaskRegion &build_matrix_region = tc.AddRegion(blocks.size());
+  for (int i = 0; i < blocks.size(); i++) {
+    auto &tl = build_matrix_region[i];
+    auto &pmb = blocks[i];
+    auto build_block = tl.AddTask(none, TF(HypreSolver::BuildMatrixVector),
+                                  hypre_solver.get(), i, pmb.get(), integrator.dt);
+    // probably have a task for setting RHS and initial guess
+  }
+
+  TaskRegion &solve_region = tc.AddRegion(1);
+  auto solve = solve_region[0].AddTask(none, TF(HypreSolver::Solve), hypre_solver.get());
+
+  TaskRegion &update_region = tc.AddRegion(blocks.size());
+  for (int i = 0; i < blocks.size(); ++i) {
+    TaskList &tl = update_region[i];
+    auto &pmb = blocks[i];
+    auto update_block = tl.AddTask(none, TF(HypreSolver::UpdateSolution),
+                                   hypre_solver.get(), i, pmb.get());
+  }
+
+  TaskRegion &dt_region = tc.AddRegion(num_partitions);
+  for (int i = 0; i < num_partitions; ++i) {
+    TaskList &tl = dt_region[i];
+    auto &md = pmesh->mesh_data.Add("base", partitions[i]);
+
+    // Update the timestep
+    tl.AddTask(none, parthenon::Update::EstimateTimestep<MeshData<Real>>, md.get());
+  }
+
+#endif // DIFFUSION_WITH_HYPRE
+  return tc;
+}
 
 } // namespace diffusion_example

example/diffusion/diffusion_driver.hpp

@@ -17,6 +17,7 @@
 #include <memory>
 #include <vector>
 
+#include "diffusion_hypre.hpp"
 #include <kokkos_abstraction.hpp>
 #include <parthenon/driver.hpp>
 #include <parthenon/package.hpp>
@@ -30,23 +31,41 @@ class DiffusionDriver : public EvolutionDriver {
  public:
   DiffusionDriver(ParameterInput *pin, ApplicationInput *app_in, Mesh *pm)
       : EvolutionDriver(pin, app_in, pm), integrator(pin) {
+    u2.val = 1e200;
     //    InitializeOutputs();
   }
   // This next function essentially defines the driver.
   TaskCollection MakeTaskCollection();
+  TaskCollection MakeTaskCollectionHypre();
+  TaskCollection MakeTaskCollectionNative();
   TaskListStatus Step() override;
 
   // DriverStatus Execute() override;
   void OutputDownstreamCycleDiagnostics() override {
     auto pkg = pmesh->packages.Get("diffusion_package");
-    auto solver_type = pkg->Param<std::string>("solver");
-    auto psolver =
-        pkg->Param<std::shared_ptr<parthenon::solvers::SolverBase>>("solver_pointer");
-    int v_cycles = psolver->GetFinalIterations();
-    if (solver_type == "BiCGSTAB") v_cycles *= 2;
-    std::cout << " v-cycles=" << v_cycles;
+    bool print{true};
+#ifdef DIFFUSION_WITH_HYPRE
+    if (pkg->Param<bool>("use_hypre")) {
+      auto hypre_solver =
+          pkg->Param<std::shared_ptr<diffusion_package::HypreSolver>>("hypre_solver");
+      std::cout << " v-cycles=" << hypre_solver->niter * 2
+                << " rel_resid=" << hypre_solver->rnorm;
+      print = false;
+    }
+#endif
+    if (print) {
+      auto solver_type = pkg->Param<std::string>("solver");
+      auto psolver =
+          pkg->Param<std::shared_ptr<parthenon::solvers::SolverBase>>("solver_pointer");
+      const auto alpha = pkg->Param<Real>("diagonal_alpha");
+      int v_cycles = psolver->GetFinalIterations();
+      auto res = psolver->GetFinalResidual();
+      if (solver_type == "BiCGSTAB") v_cycles *= 2;
+      std::cout << " v-cycles=" << v_cycles
+                << " rel_resid=" << res / (alpha * sqrt(u2.val));
+    }
   }
-  
+
   void PostExecute(DriverStatus status) override {
     EvolutionDriver::PostExecute(status);
     if (parthenon::Globals::my_rank == 0) {
@@ -58,12 +77,12 @@ class DiffusionDriver : public EvolutionDriver {
         std::cout << "Solver breakdown: \n" << psolver->solver_timings;
         psolver->solver_timings.clear();
       }
-
     }
   }
 
  private:
   LowStorageIntegrator integrator;
+  parthenon::AllReduce<Real> u2;
 };
 
 void ProblemGenerator(Mesh *pm, parthenon::ParameterInput *pin, MeshData<Real> *md);