diff --git a/cmake/FindHYPRE.cmake b/cmake/FindHYPRE.cmake new file mode 100644 index 0000000000000..0dda43dd53d7b --- /dev/null +++ b/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 +) + + diff --git a/example/diffusion/CMakeLists.txt b/example/diffusion/CMakeLists.txt index cf8bf49ecd67b..f623a815194d7 100644 --- a/example/diffusion/CMakeLists.txt +++ b/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 diff --git a/example/diffusion/diffusion_driver.cpp b/example/diffusion/diffusion_driver.cpp index e9b29a632a154..14153016cbdad 100644 --- a/example/diffusion/diffusion_driver.cpp +++ b/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("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>("solver_pointer"); const auto alpha = pkg->Param("diagonal_alpha"); + const auto rel_res = pkg->Param("rel_res"); auto peqs = pkg->Param>>( "diffusion_equation"); @@ -86,9 +105,22 @@ TaskCollection DiffusionDriver::MakeTaskCollection() { tl.AddTask(Au, solvers::utils::AddFieldsAndStore>, md, md_rhs, md_rhs, alpha, -1.0); + // Get the RHS scale for correct comparison to Hypre solver + set_rhs = solvers::utils::DotProduct>(set_rhs, tl, &u2, md, md, + true); + + set_rhs = tl.AddTask( + set_rhs, + [alpha](parthenon::AllReduce *u2, + std::shared_ptr 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), 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>("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 ®ion = 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>, md.get()); + } + +#endif // DIFFUSION_WITH_HYPRE + return tc; +} } // namespace diffusion_example diff --git a/example/diffusion/diffusion_driver.hpp b/example/diffusion/diffusion_driver.hpp index f0af76a530d80..90e4e0e08d0c8 100644 --- a/example/diffusion/diffusion_driver.hpp +++ b/example/diffusion/diffusion_driver.hpp @@ -17,6 +17,7 @@ #include #include +#include "diffusion_hypre.hpp" #include #include #include @@ -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("solver"); - auto psolver = - pkg->Param>("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("use_hypre")) { + auto hypre_solver = + pkg->Param>("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("solver"); + auto psolver = + pkg->Param>("solver_pointer"); + const auto alpha = pkg->Param("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 u2; }; void ProblemGenerator(Mesh *pm, parthenon::ParameterInput *pin, MeshData *md); diff --git a/example/diffusion/diffusion_hypre.cpp b/example/diffusion/diffusion_hypre.cpp new file mode 100644 index 0000000000000..e18ebf5209e6c --- /dev/null +++ b/example/diffusion/diffusion_hypre.cpp @@ -0,0 +1,1135 @@ +#include "diffusion_hypre.hpp" + +#include +#include +#include +#include +#include +#include +#include + +#include "HYPRE_parcsr_mv.h" +#include "basic_types.hpp" +#include "defs.hpp" +#include "diffusion_package.hpp" +#include "globals.hpp" +#include "kokkos_abstraction.hpp" +#include "mesh/mesh.hpp" +#include "mesh/meshblock.hpp" +#include "pack/make_pack_descriptor.hpp" +#include "parameter_input.hpp" +#include "utils/error_checking.hpp" + +namespace diffusion_package { +using Real = parthenon::Real; + +namespace { + +int FaceFromOffsets(const parthenon::CellCentOffsets &ofs) { + if (ofs(parthenon::X1DIR) == -1) return parthenon::BoundaryFace::inner_x1; + if (ofs(parthenon::X1DIR) == 1) return parthenon::BoundaryFace::outer_x1; + if (ofs(parthenon::X2DIR) == -1) return parthenon::BoundaryFace::inner_x2; + if (ofs(parthenon::X2DIR) == 1) return parthenon::BoundaryFace::outer_x2; + if (ofs(parthenon::X3DIR) == -1) return parthenon::BoundaryFace::inner_x3; + if (ofs(parthenon::X3DIR) == 1) return parthenon::BoundaryFace::outer_x3; + return parthenon::BoundaryFace::undef; +} + +int FaceAxis(const int face) { + if (face == parthenon::BoundaryFace::inner_x1 || + face == parthenon::BoundaryFace::outer_x1) { + return 0; + } + if (face == parthenon::BoundaryFace::inner_x2 || + face == parthenon::BoundaryFace::outer_x2) { + return 1; + } + if (face == parthenon::BoundaryFace::inner_x3 || + face == parthenon::BoundaryFace::outer_x3) { + return 2; + } + return -1; +} + +int FaceSide(const int face) { + if (face == parthenon::BoundaryFace::inner_x1 || + face == parthenon::BoundaryFace::inner_x2 || + face == parthenon::BoundaryFace::inner_x3) { + return -1; + } + if (face == parthenon::BoundaryFace::outer_x1 || + face == parthenon::BoundaryFace::outer_x2 || + face == parthenon::BoundaryFace::outer_x3) { + return 1; + } + return 0; +} + +int DfcComponentFromGlobal(const int axis, const int gi, const int gj, const int gk, + const int ndim) { + if (axis == 0) { + const int comp1 = gj & 1; + const int comp2 = (ndim > 2) ? (gk & 1) : 0; + return comp1 + 2 * comp2; + } + if (axis == 1) { + const int comp1 = (ndim > 2) ? (gk & 1) : 0; + const int comp2 = gi & 1; + return comp1 + 2 * comp2; + } + const int comp1 = gi & 1; + const int comp2 = gj & 1; + return comp1 + 2 * comp2; +} + +void NeighborFaceBounds(const std::array &lo, const std::array &hi, + const int ndim, const parthenon::NeighborBlock &nb, + const bool neighbor_is_fine, int &is, int &ie, int &js, int &je, + int &ks, int &ke, int &axis, int &side) { + const int face = FaceFromOffsets(nb.offsets); + axis = FaceAxis(face); + side = FaceSide(face); + + is = lo[0]; + ie = hi[0]; + js = lo[1]; + je = hi[1]; + ks = lo[2]; + ke = hi[2]; + + if (axis == 0) { + if (side < 0) + ie = is; + else + is = ie; + } else if (axis == 1) { + if (side < 0) + je = js; + else + js = je; + } else { + if (side < 0) + ke = ks; + else + ks = ke; + } + + // If neighbor is finer, this neighbor only covers a half-face (2D) or quarter-face + // (3D). Restrict the tangential ranges using fi1/fi2. + if (neighbor_is_fine) { + auto split_half = [](int &s, int &e, int fi) { + const int n = e - s + 1; + const int h = n / 2; + s += fi * h; + e = s + h - 1; + }; + + if (axis == 0) { + split_half(js, je, nb.fi1); + if (ndim > 2) split_half(ks, ke, nb.fi2); + } else if (axis == 1) { + if (ndim > 2) { + split_half(ks, ke, nb.fi1); + split_half(is, ie, nb.fi2); + } else { + split_half(is, ie, nb.fi1); + } + } else { + split_half(is, ie, nb.fi1); + split_half(js, je, nb.fi2); + } + } +} + +} // namespace + +HypreSolver::HypreSolver(parthenon::ParameterInput *pin) { + // Solver type + solver_type = pin->GetOrAddString("hypre", "solver_type", "bicgstab", + "Hypre outer solver: pcg or bicgstab"); + std::transform(solver_type.begin(), solver_type.end(), solver_type.begin(), + [](unsigned char c) { return static_cast(std::tolower(c)); }); + PARTHENON_REQUIRE(solver_type == "pcg" || solver_type == "bicgstab", + "hypre/solver_type must be 'pcg' or 'bicgstab'."); + preconditioner = pin->GetOrAddString("hypre", "preconditioner", "amg", + "Hypre preconditioner: amg or none"); + std::transform(preconditioner.begin(), preconditioner.end(), preconditioner.begin(), + [](unsigned char c) { return static_cast(std::tolower(c)); }); + PARTHENON_REQUIRE(preconditioner == "amg" || preconditioner == "none", + "hypre/preconditioner must be 'amg' or 'none'."); + tol = pin->GetOrAddReal("hypre", "tol", 1e-12, "Relative convergence tolerance"); + absolute_tol = + pin->GetOrAddReal("hypre", "absolute_tol", 0.0, + "Absolute convergence tolerance (0 disables absolute stop)"); + max_iter = pin->GetOrAddInteger("hypre", "max_iter", 50, "Maximum solver iterations"); + print_level = pin->GetOrAddInteger("hypre", "print_level", 1, "Solver print verbosity"); + internal_print_level = pin->GetOrAddInteger( + "hypre", "internal_print_level", 0, + "Hypre internal solver print level (0 disables iteration table)"); + pcg_use_two_norm = + pin->GetOrAddBoolean("hypre", "pcg_use_two_norm", true, + "Use true 2-norm residual in PCG stopping criterion"); + pcg_recompute_residual = + pin->GetOrAddBoolean("hypre", "pcg_recompute_residual", false, + "Recompute residual explicitly in PCG for robustness"); + + // BoomerAMG preconditioner settings + amg_coarsen_type = + pin->GetOrAddInteger("hypre", "amg_coarsen_type", 10, "AMG coarsening type (HMIS)"); + amg_interp_type = pin->GetOrAddInteger("hypre", "amg_interp_type", 6, + "AMG interpolation type (ext+i)"); + amg_relax_type = pin->GetOrAddInteger("hypre", "amg_relax_type", 6, + "AMG relaxation type (symmetric GS)"); + amg_strong_threshold = pin->GetOrAddReal("hypre", "amg_strong_threshold", 0.25, + "AMG strong threshold (0.25 for 2D)"); + amg_num_sweeps = + pin->GetOrAddInteger("hypre", "amg_num_sweeps", 1, "AMG sweeps per level"); + + print_matrix = pin->GetOrAddBoolean("hypre", "print_matrix", false, + "Print Hypre matrix after assemble"); + print_vectors = pin->GetOrAddBoolean("hypre", "print_vectors", false, + "Print Hypre vectors after assemble"); + print_matrix_all = pin->GetOrAddBoolean("hypre", "print_matrix_all", false, + "Print full matrix (all=1)"); + // Cache problem parameters + diagonal_alpha = pin->GetReal("diffusion", "diagonal_alpha"); + auto u_bounds = + pin->GetOrAddVector("diffusion", "boundary_u", {0.0}, "Boundary us."); + if (u_bounds.size() == 1) u_bounds = std::vector(6, u_bounds[0]); + PARTHENON_REQUIRE(u_bounds.size() == 6, + "diffusion/boundary_u must have exactly 1 or 6 entries."); + for (int f = 0; f < 6; ++f) + boundary_u[f] = u_bounds[f]; + + // Determine dimensionality + const int nx3 = pin->GetInteger("parthenon/mesh", "nx3"); + ndim = (nx3 > 1) ? 3 : 2; + nstencil = (ndim == 1) ? 3 : (ndim == 2) ? 5 : 7; +} + +HypreSolver::~HypreSolver() { DestroyGrid(); } + +void HypreSolver::DestroyGrid() { + if (solver_handle) { + if (solver_type == "pcg") { + HYPRE_ParCSRPCGDestroy(solver_handle); + } else { + HYPRE_ParCSRBiCGSTABDestroy(solver_handle); + } + solver_handle = nullptr; + } + if (precond_handle) { + HYPRE_BoomerAMGDestroy(precond_handle); + precond_handle = nullptr; + } + if (A) { + HYPRE_SStructMatrixDestroy(A); + A = nullptr; + } + if (b) { + HYPRE_SStructVectorDestroy(b); + b = nullptr; + } + if (x) { + HYPRE_SStructVectorDestroy(x); + x = nullptr; + } + if (graph) { + HYPRE_SStructGraphDestroy(graph); + graph = nullptr; + } + if (stencil) { + HYPRE_SStructStencilDestroy(stencil); + stencil = nullptr; + } + if (grid) { + HYPRE_SStructGridDestroy(grid); + grid = nullptr; + } + + block_part.clear(); + block_ilower.clear(); + block_iupper.clear(); + block_neighbor_level.clear(); + block_is_domain_boundary.clear(); + level_to_part.clear(); + part_periodic.clear(); + + nparts = 0; + min_active_level = -1; + max_active_level = -1; + grid_is_setup = false; + solver_is_setup = false; +} + +parthenon::TaskStatus HypreSolver::BuildMatrixVector(HypreSolver *solver, int b, + parthenon::MeshBlock *pmb, + const Real dt) { + using namespace parthenon; + using TE = parthenon::TopologicalElement; + (void)dt; + + PARTHENON_REQUIRE(solver->grid_is_setup, + "BuildMatrixVector called before Hypre grid setup."); + PARTHENON_REQUIRE(b >= 0 && b < static_cast(solver->block_part.size()), + "BuildMatrixVector called with invalid block index."); + + auto *pmbd = pmb->meshblock_data.Get().get(); + auto desc = parthenon::MakePackDescriptor(pmbd); + auto pack = desc.GetPack(pmbd); + PARTHENON_REQUIRE(pack.GetNBlocks() == 1, + "BuildMatrixVector expects exactly one block in pack."); + constexpr int pb = 0; + + const auto ib = pmb->cellbounds.GetBoundsI(IndexDomain::interior); + const auto jb = pmb->cellbounds.GetBoundsJ(IndexDomain::interior); + const auto kb = pmb->cellbounds.GetBoundsK(IndexDomain::interior); + + const int ni = ib.e - ib.s + 1; + const int nj = jb.e - jb.s + 1; + const int nk = kb.e - kb.s + 1; + const int ncell = ni * nj * nk; + + const auto &il = solver->block_ilower[b]; + const auto &iu = solver->block_iupper[b]; + const int part = solver->block_part[b]; + const int legacy_root_level = pmb->pmy_mesh->GetLegacyTreeRootLevel(); + const auto legacy_loc = pmb->pmy_mesh->Forest().GetLegacyTreeLocation(pmb->loc); + const int lev = static_cast(legacy_loc.level()) - legacy_root_level; + + std::vector stencil_entries(solver->nstencil); + for (int e = 0; e < solver->nstencil; ++e) + stencil_entries[e] = e; + + // Hypre SStruct matrix/vector APIs consume HYPRE_Complex typed arrays. + std::vector matvals(static_cast(ncell * solver->nstencil), + 0.0); + std::vector rhsvals(static_cast(ncell), 0.0); + std::vector xvals(static_cast(ncell), 0.0); + + auto lin_idx = [&](const int k, const int j, const int i) { + return (k - kb.s) * nj * ni + (j - jb.s) * ni + (i - ib.s); + }; + auto A = [&](const int lin, const int ent) -> HYPRE_Complex & { + return matvals[static_cast(lin * solver->nstencil + ent)]; + }; + auto local_from_global = [&](const int gk, const int gj, const int gi) { + const int li = gi - il[0]; + const int lj = gj - il[1]; + const int lk = gk - il[2]; + PARTHENON_REQUIRE(li >= 0 && li < ni && lj >= 0 && lj < nj && lk >= 0 && lk < nk, + "Global-to-local mapping out of block bounds."); + return std::array{ib.s + li, jb.s + lj, kb.s + lk}; + }; + + auto face_conductance = [&](const int axis, const int side, const int k, const int j, + const int i) { + if (axis == 0) { + const Real d = pmb->coords.Dxc(k, j, i); + const Real area = (side < 0) ? pmb->coords.Volume(k, j, i) + : pmb->coords.Volume(k, j, i + 1); + const Real Dface = (side < 0) + ? pack(pb, TE::F1, diffusion_package::D(), k, j, i) + : pack(pb, TE::F1, diffusion_package::D(), k, j, i + 1); + return Dface * area / d; + } else if (axis == 1) { + const Real d = pmb->coords.Dxc(k, j, i); + const Real area = (side < 0) ? pmb->coords.Volume(k, j, i) + : pmb->coords.Volume(k, j + 1, i); + const Real Dface = (side < 0) + ? pack(pb, TE::F2, diffusion_package::D(), k, j, i) + : pack(pb, TE::F2, diffusion_package::D(), k, j + 1, i); + return Dface * area / d; + } + const Real d = pmb->coords.Dxc(k, j, i); + const Real area = (side < 0) ? pmb->coords.Volume(k, j, i) + : pmb->coords.Volume(k + 1, j, i); + const Real Dface = (side < 0) ? pack(pb, TE::F3, diffusion_package::D(), k, j, i) + : pack(pb, TE::F3, diffusion_package::D(), k + 1, j, i); + return Dface * area / d; + }; + + auto zero_face_stencil = [&](const int face, const int lin) { + if (face == BoundaryFace::inner_x1) + A(lin, 1) = 0.0; + else if (face == BoundaryFace::outer_x1) + A(lin, 2) = 0.0; + else if (face == BoundaryFace::inner_x2) + A(lin, 3) = 0.0; + else if (face == BoundaryFace::outer_x2) + A(lin, 4) = 0.0; + else if (face == BoundaryFace::inner_x3) + A(lin, 5) = 0.0; + else if (face == BoundaryFace::outer_x3) + A(lin, 6) = 0.0; + }; + + auto face_stencil_entry = [&](const int face) { + if (face == BoundaryFace::inner_x1) return 1; + if (face == BoundaryFace::outer_x1) return 2; + if (face == BoundaryFace::inner_x2) return 3; + if (face == BoundaryFace::outer_x2) return 4; + if (face == BoundaryFace::inner_x3) return 5; + if (face == BoundaryFace::outer_x3) return 6; + return -1; + }; + + parthenon::par_for( + parthenon::loop_pattern_mdrange_tag, "build_matrix_rows", + parthenon::HostExecSpace(), kb.s, kb.e, jb.s, jb.e, ib.s, ib.e, + [&](const int k, const int j, const int i) { + const int lin = lin_idx(k, j, i); + + const Real cell_vol = pmb->coords.Volume(k, j, i); + const Real kxm = face_conductance(0, -1, k, j, i); + const Real kxp = face_conductance(0, +1, k, j, i); + const Real kym = (solver->ndim > 1) ? face_conductance(1, -1, k, j, i) : 0.0; + const Real kyp = (solver->ndim > 1) ? face_conductance(1, +1, k, j, i) : 0.0; + const Real kzm = (solver->ndim > 2) ? face_conductance(2, -1, k, j, i) : 0.0; + const Real kzp = (solver->ndim > 2) ? face_conductance(2, +1, k, j, i) : 0.0; + + A(lin, 0) = solver->diagonal_alpha * cell_vol + kxm + kxp + kym + kyp + kzm + kzp; + A(lin, 1) = -kxm; + A(lin, 2) = -kxp; + if (solver->ndim > 1) { + A(lin, 3) = -kym; + A(lin, 4) = -kyp; + } + if (solver->ndim > 2) { + A(lin, 5) = -kzm; + A(lin, 6) = -kzp; + } + + const Real u0 = pack(pb, diffusion_package::u(), k, j, i); + rhsvals[lin] = solver->diagonal_alpha * cell_vol * u0; + xvals[lin] = u0; + }); + + // Phase B: physical Dirichlet boundary corrections. + for (int face = 0; face < 2 * solver->ndim; ++face) { + if (!solver->block_is_domain_boundary[b][face]) continue; + const int axis = FaceAxis(face); + const int side = FaceSide(face); + int gis = il[0], gie = iu[0]; + int gjs = il[1], gje = iu[1]; + int gks = il[2], gke = iu[2]; + if (axis == 0) { + if (side < 0) + gie = gis; + else + gis = gie; + } else if (axis == 1) { + if (side < 0) + gje = gjs; + else + gjs = gje; + } else { + if (side < 0) + gke = gks; + else + gks = gke; + } + + const Real bc = solver->boundary_u[face]; + parthenon::par_for(parthenon::loop_pattern_mdrange_tag, "bc_fixup", + parthenon::HostExecSpace(), gks, gke, gjs, gje, gis, gie, + [&](const int gk, const int gj, const int gi) { + auto lidx = local_from_global(gk, gj, gi); + const int i = lidx[0], j = lidx[1], k = lidx[2]; + const int lin = lin_idx(k, j, i); + const Real kface = face_conductance(axis, side, k, j, i); + A(lin, 0) += kface; + zero_face_stencil(face, lin); + rhsvals[lin] += 2.0 * kface * bc; + }); + } + + // Set non-stencil graph couplings at fine-coarse boundaries. + std::vector row_graph_count(static_cast(ncell), 0); + for (const auto &nb : pmb->GetNeighbors()) { + const int ax = std::abs(nb.offsets(parthenon::X1DIR)); + const int ay = std::abs(nb.offsets(parthenon::X2DIR)); + const int az = std::abs(nb.offsets(parthenon::X3DIR)); + if (ax + ay + az != 1) continue; + + const int face = FaceFromOffsets(nb.offsets); + if (face == parthenon::BoundaryFace::undef) continue; + + const auto nlegacy_loc = pmb->pmy_mesh->Forest().GetLegacyTreeLocation(nb.origin_loc); + const int nlev = static_cast(nlegacy_loc.level()) - legacy_root_level; + const int neighbor_level_relation = (nlev > lev) ? 1 : ((nlev < lev) ? -1 : 0); + if (neighbor_level_relation == 0) continue; + + int is, ie, js, je, ks, ke, axis, side; + NeighborFaceBounds(il, iu, solver->ndim, nb, neighbor_level_relation > 0, is, ie, js, + je, ks, ke, axis, side); + + for (int gk = ks; gk <= ke; ++gk) { + for (int gj = js; gj <= je; ++gj) { + for (int gi = is; gi <= ie; ++gi) { + const int lin = (gk - il[2]) * nj * ni + (gj - il[1]) * ni + (gi - il[0]); + auto lidx = local_from_global(gk, gj, gi); + const int i = lidx[0], j = lidx[1], k = lidx[2]; + + auto set_graph_value = [&](const Real value) { + int entry = solver->nstencil + row_graph_count[lin]; + int index[3] = {gi, gj, gk}; + HYPRE_Complex hval = value; + HYPRE_SStructMatrixSetValues(solver->A, part, index, 0, 1, &entry, &hval); + row_graph_count[lin] += 1; + }; + + const int face_entry = face_stencil_entry(face); + PARTHENON_REQUIRE(face_entry >= 0, + "Invalid face when applying FC graph/stencil corrections."); + const Real kface = face_conductance(axis, side, k, j, i); + const Real d_self = pmb->coords.Dxc(axis + 1, k, j, i); + const Real d_neighbor = + (neighbor_level_relation < 0) ? (2.0 * d_self) : (0.5 * d_self); + const Real d_eff = 0.5 * (d_self + d_neighbor); + zero_face_stencil(face, lin); + + Real diag_correction = 0.0; + + if (neighbor_level_relation < 0) { + // coarser neighbor, only 1 graph entry + const Real cond_face = kface * d_self; + const Real fc_face_diag = cond_face / d_eff; + diag_correction = fc_face_diag - kface; + A(lin, 0) += diag_correction; + const Real gval = -fc_face_diag; + set_graph_value(gval); + } else { + // finer neighbor, 1 graph through each subface + const int nsub = (solver->ndim == 2) ? 2 : 4; + const Real face_area = pmb->coords.FaceArea(axis + 1, k, j, i); + const Real subface_area = face_area / static_cast(nsub); + + std::array fine_face_anchor = {2 * gi, 2 * gj, 2 * gk}; + fine_face_anchor[axis] += (side < 0) ? -1 : 2; + const int tan_axis0 = (axis + 1) % solver->ndim; + const int tan_axis1 = (axis + 2) % solver->ndim; + + // Periodic wrapping info for the fine-level part, used to canonicalize + // fine `to` indices before DfcComponentFromGlobal parity check. + const int fine_part = solver->level_to_part[nlev]; + const auto &fpp = solver->part_periodic[fine_part]; + + auto wrap_to = [&](std::array &idx) { + for (int d = 0; d < solver->ndim; ++d) { + if (fpp[d] > 0) { + idx[d] = ((idx[d] % fpp[d]) + fpp[d]) % fpp[d]; + } + } + }; + + Real sum_cond_sub = 0.0; + if (solver->ndim == 2) { + for (int s = 0; s < 2; ++s) { + std::array to = fine_face_anchor; + to[tan_axis0] = 2 * ((tan_axis0 == 0) ? gi : gj) + s; + wrap_to(to); + const int comp = + DfcComponentFromGlobal(axis, to[0], to[1], to[2], solver->ndim); + const Real Dsub = + (axis == 0) + ? ((side < 0) + ? pack(pb, TE::F1, diffusion_package::Dfc(comp), k, j, i) + : pack(pb, TE::F1, diffusion_package::Dfc(comp), k, j, + i + 1)) + : ((side < 0) + ? pack(pb, TE::F2, diffusion_package::Dfc(comp), k, j, i) + : pack(pb, TE::F2, diffusion_package::Dfc(comp), k, j + 1, + i)); + const Real cond_sub = Dsub * subface_area; + sum_cond_sub += cond_sub; + const Real gval = -cond_sub / d_eff; + set_graph_value(gval); + } + } else { + for (int s0 = 0; s0 < 2; ++s0) { + for (int s1 = 0; s1 < 2; ++s1) { + std::array to = fine_face_anchor; + to[tan_axis0] = + 2 * ((tan_axis0 == 0) ? gi : ((tan_axis0 == 1) ? gj : gk)) + s0; + to[tan_axis1] = + 2 * ((tan_axis1 == 0) ? gi : ((tan_axis1 == 1) ? gj : gk)) + s1; + wrap_to(to); + const int comp = + DfcComponentFromGlobal(axis, to[0], to[1], to[2], solver->ndim); + const Real Dsub = + (axis == 0) + ? ((side < 0) + ? pack(pb, TE::F1, diffusion_package::Dfc(comp), k, j, i) + : pack(pb, TE::F1, diffusion_package::Dfc(comp), k, j, + i + 1)) + : ((axis == 1) + ? ((side < 0) + ? pack(pb, TE::F2, diffusion_package::Dfc(comp), + k, j, i) + : pack(pb, TE::F2, diffusion_package::Dfc(comp), + k, j + 1, i)) + : ((side < 0) + ? pack(pb, TE::F3, diffusion_package::Dfc(comp), + k, j, i) + : pack(pb, TE::F3, diffusion_package::Dfc(comp), + k + 1, j, i))); + const Real cond_sub = Dsub * subface_area; + sum_cond_sub += cond_sub; + const Real gval = -cond_sub / d_eff; + set_graph_value(gval); + } + } + } + + const Real fc_face_diag = sum_cond_sub / d_eff; + diag_correction = fc_face_diag - kface; + A(lin, 0) += diag_correction; + } + } + } + } + } + + HYPRE_SStructMatrixSetBoxValues(solver->A, part, const_cast(il.data()), + const_cast(iu.data()), 0, solver->nstencil, + stencil_entries.data(), matvals.data()); + + HYPRE_SStructVectorSetBoxValues(solver->b, part, const_cast(il.data()), + const_cast(iu.data()), 0, rhsvals.data()); + HYPRE_SStructVectorSetBoxValues(solver->x, part, const_cast(il.data()), + const_cast(iu.data()), 0, xvals.data()); + + return parthenon::TaskStatus::complete; +} + +parthenon::TaskStatus HypreSolver::Solve(HypreSolver *solver) { + solver->solve_call_count += 1; + HYPRE_SStructMatrixAssemble(solver->A); + HYPRE_SStructVectorAssemble(solver->b); + HYPRE_SStructVectorAssemble(solver->x); + + if (solver->print_matrix || solver->print_vectors) { + std::ostringstream base; + base << "hypre_debug_step" << solver->solve_call_count << "_rank" + << parthenon::Globals::my_rank; + if (solver->print_matrix) { + std::string mat_name = base.str() + ".A"; + HYPRE_SStructMatrixPrint(mat_name.c_str(), solver->A, + solver->print_matrix_all ? 1 : 0); + } + if (solver->print_vectors) { + std::string b_name = base.str() + ".b"; + std::string x_name = base.str() + ".x"; + HYPRE_SStructVectorPrint(b_name.c_str(), solver->b, + solver->print_matrix_all ? 1 : 0); + HYPRE_SStructVectorPrint(x_name.c_str(), solver->x, + solver->print_matrix_all ? 1 : 0); + } + } + + HYPRE_ParCSRMatrix parA; + HYPRE_ParVector parb, parx; + HYPRE_SStructMatrixGetObject(solver->A, reinterpret_cast(&parA)); + HYPRE_SStructVectorGetObject(solver->b, reinterpret_cast(&parb)); + HYPRE_SStructVectorGetObject(solver->x, reinterpret_cast(&parx)); + + if (!solver->solver_is_setup) { + solver->SetupSolver(); + } + + HYPRE_Int niter = 0; + HYPRE_Real rnorm = 0.0; + + if (solver->solver_type == "pcg") { + HYPRE_ParCSRPCGSetup(solver->solver_handle, parA, parb, parx); + HYPRE_ParCSRPCGSolve(solver->solver_handle, parA, parb, parx); + HYPRE_ParCSRPCGGetNumIterations(solver->solver_handle, &niter); + HYPRE_ParCSRPCGGetFinalRelativeResidualNorm(solver->solver_handle, &rnorm); + } else { + HYPRE_ParCSRBiCGSTABSetup(solver->solver_handle, parA, parb, parx); + HYPRE_ParCSRBiCGSTABSolve(solver->solver_handle, parA, parb, parx); + HYPRE_ParCSRBiCGSTABGetNumIterations(solver->solver_handle, &niter); + HYPRE_ParCSRBiCGSTABGetFinalRelativeResidualNorm(solver->solver_handle, &rnorm); + } + + HYPRE_SStructVectorGather(solver->x); + solver->niter = niter; + solver->rnorm = rnorm; + return parthenon::TaskStatus::complete; +} + +parthenon::TaskStatus HypreSolver::UpdateSolution(HypreSolver *solver, int b, + parthenon::MeshBlock *pmb) { + using namespace parthenon; + + const auto ib = pmb->cellbounds.GetBoundsI(IndexDomain::interior); + const auto jb = pmb->cellbounds.GetBoundsJ(IndexDomain::interior); + const auto kb = pmb->cellbounds.GetBoundsK(IndexDomain::interior); + + const int ni = ib.e - ib.s + 1; + const int nj = jb.e - jb.s + 1; + const int nk = kb.e - kb.s + 1; + const int ncell = ni * nj * nk; + + const int part = solver->block_part[b]; + const auto &il = solver->block_ilower[b]; + const auto &iu = solver->block_iupper[b]; + + std::vector soln(static_cast(ncell), 0.0); + HYPRE_SStructVectorGetBoxValues(solver->x, part, const_cast(il.data()), + const_cast(iu.data()), 0, soln.data()); + + auto &uvar = pmb->meshblock_data.Get()->Get(diffusion_package::u::name()).data; + auto lin_idx = [&](const int k, const int j, const int i) { + return (k - kb.s) * nj * ni + (j - jb.s) * ni + (i - ib.s); + }; + + for (int k = kb.s; k <= kb.e; ++k) { + for (int j = jb.s; j <= jb.e; ++j) { + for (int i = ib.s; i <= ib.e; ++i) { + uvar(k, j, i) = soln[lin_idx(k, j, i)]; + } + } + } + + return parthenon::TaskStatus::complete; +} + +void HypreSolver::SetupSolver() { + if (solver_handle) { + if (solver_type == "pcg") { + HYPRE_ParCSRPCGDestroy(solver_handle); + } else { + HYPRE_ParCSRBiCGSTABDestroy(solver_handle); + } + solver_handle = nullptr; + } + if (precond_handle) { + HYPRE_BoomerAMGDestroy(precond_handle); + precond_handle = nullptr; + } + + const bool use_amg_preconditioner = (preconditioner == "amg"); + if (use_amg_preconditioner) { + HYPRE_BoomerAMGCreate(&precond_handle); + HYPRE_BoomerAMGSetTol(precond_handle, 0.0); + HYPRE_BoomerAMGSetMaxIter(precond_handle, 1); + HYPRE_BoomerAMGSetCoarsenType(precond_handle, amg_coarsen_type); + HYPRE_BoomerAMGSetInterpType(precond_handle, amg_interp_type); + HYPRE_BoomerAMGSetRelaxType(precond_handle, amg_relax_type); + HYPRE_BoomerAMGSetStrongThreshold(precond_handle, amg_strong_threshold); + HYPRE_BoomerAMGSetNumSweeps(precond_handle, amg_num_sweeps); + HYPRE_BoomerAMGSetPrintLevel(precond_handle, 0); + } + + if (solver_type == "pcg") { + HYPRE_ParCSRPCGCreate(MPI_COMM_WORLD, &solver_handle); + HYPRE_ParCSRPCGSetTol(solver_handle, tol); + if (absolute_tol > 0.0) { + HYPRE_ParCSRPCGSetAbsoluteTol(solver_handle, absolute_tol); + } + HYPRE_ParCSRPCGSetMaxIter(solver_handle, max_iter); + if (pcg_use_two_norm) { + HYPRE_ParCSRPCGSetTwoNorm(solver_handle, 1); + } + if (pcg_recompute_residual) { + HYPRE_PCGSetRecomputeResidual(solver_handle, 1); + } + HYPRE_ParCSRPCGSetPrintLevel(solver_handle, internal_print_level); + HYPRE_ParCSRPCGSetLogging(solver_handle, (internal_print_level > 0) ? 1 : 0); + if (use_amg_preconditioner) { + HYPRE_ParCSRPCGSetPrecond(solver_handle, HYPRE_BoomerAMGSolve, HYPRE_BoomerAMGSetup, + precond_handle); + } + } else { + HYPRE_ParCSRBiCGSTABCreate(MPI_COMM_WORLD, &solver_handle); + HYPRE_ParCSRBiCGSTABSetTol(solver_handle, tol); + if (absolute_tol > 0.0) { + HYPRE_ParCSRBiCGSTABSetAbsoluteTol(solver_handle, absolute_tol); + } + HYPRE_ParCSRBiCGSTABSetMaxIter(solver_handle, max_iter); + HYPRE_ParCSRBiCGSTABSetPrintLevel(solver_handle, internal_print_level); + HYPRE_ParCSRBiCGSTABSetLogging(solver_handle, (internal_print_level > 0) ? 1 : 0); + if (use_amg_preconditioner) { + HYPRE_ParCSRBiCGSTABSetPrecond(solver_handle, HYPRE_BoomerAMGSolve, + HYPRE_BoomerAMGSetup, precond_handle); + } + } + + solver_is_setup = true; +} + +void HypreSolver::SetupGrid(parthenon::Mesh *pmesh) { + if (grid_is_setup) return; + + auto &blocks = pmesh->block_list; + const int nblocks = static_cast(blocks.size()); + const int legacy_root_level = pmesh->GetLegacyTreeRootLevel(); + + if (nblocks == 0) { + PARTHENON_FAIL("SetupGrid called with empty block list."); + } + + block_part.resize(nblocks, -1); + block_ilower.resize(nblocks); + block_iupper.resize(nblocks); + block_neighbor_level.resize(nblocks); + block_is_domain_boundary.resize(nblocks); + std::vector, std::array>>> + global_part_boxes; + + // Determine globally active refinement levels. + int local_max_level = -1; + int local_min_level = std::numeric_limits::max(); + for (const auto &pmb : blocks) { + const auto legacy_loc = pmesh->Forest().GetLegacyTreeLocation(pmb->loc); + const int lev = static_cast(legacy_loc.level()) - legacy_root_level; + local_max_level = std::max(local_max_level, lev); + local_min_level = std::min(local_min_level, lev); + } + + int max_level = -1; + int min_level = std::numeric_limits::max(); + MPI_Allreduce(&local_max_level, &max_level, 1, MPI_INT, MPI_MAX, MPI_COMM_WORLD); + MPI_Allreduce(&local_min_level, &min_level, 1, MPI_INT, MPI_MIN, MPI_COMM_WORLD); + + std::vector local_level_present(std::max(max_level + 1, 0), 0); + for (const auto &pmb : blocks) { + const auto legacy_loc = pmesh->Forest().GetLegacyTreeLocation(pmb->loc); + local_level_present[static_cast(legacy_loc.level()) - legacy_root_level] = 1; + } + + std::vector global_level_present(std::max(max_level + 1, 0), 0); + if (max_level >= 0) { + MPI_Allreduce(local_level_present.data(), global_level_present.data(), max_level + 1, + MPI_INT, MPI_MAX, MPI_COMM_WORLD); + } + + // map from refinement level to hypre part + level_to_part.assign(std::max(max_level + 1, 0), -1); + std::vector part_to_level; + nparts = 0; + min_active_level = -1; + max_active_level = -1; + for (int lev = 0; lev <= max_level; ++lev) { + if (global_level_present[lev]) { + level_to_part[lev] = nparts++; + part_to_level.push_back(lev); + if (min_active_level < 0) min_active_level = lev; + max_active_level = lev; + } + } + + PARTHENON_REQUIRE(nparts > 0, "HYPRE SetupGrid found no active levels."); + PARTHENON_REQUIRE(min_active_level == min_level, + "HYPRE SetupGrid min active level mismatch."); + PARTHENON_REQUIRE(max_active_level == max_level, + "HYPRE SetupGrid max active level mismatch."); + + HYPRE_SStructGridCreate(MPI_COMM_WORLD, ndim, nparts, &grid); + global_part_boxes.resize(nparts); + + // Add block extents and cache per-block metadata. + for (int b = 0; b < nblocks; ++b) { + auto *pmb = blocks[b].get(); + const auto legacy_loc = pmesh->Forest().GetLegacyTreeLocation(pmb->loc); + const int lev = static_cast(legacy_loc.level()) - legacy_root_level; + const int part = level_to_part[lev]; + block_part[b] = part; + + const int nx1 = pmb->block_size.nx(parthenon::X1DIR); + const int nx2 = pmb->block_size.nx(parthenon::X2DIR); + const int nx3 = (ndim == 3) ? pmb->block_size.nx(parthenon::X3DIR) : 1; + + const int i0 = static_cast(legacy_loc.lx1()) * nx1; + const int j0 = static_cast(legacy_loc.lx2()) * nx2; + const int k0 = (ndim == 3) ? static_cast(legacy_loc.lx3()) * nx3 : 0; + + block_ilower[b] = {i0, j0, k0}; + block_iupper[b] = {i0 + nx1 - 1, j0 + nx2 - 1, (ndim == 3) ? (k0 + nx3 - 1) : 0}; + + HYPRE_SStructGridSetExtents(grid, part, block_ilower[b].data(), + block_iupper[b].data()); + + std::array nbr_level; + nbr_level.fill(parthenon::CellLevel::same); + block_neighbor_level[b] = nbr_level; + + std::array is_domain; + for (int face = 0; face < 6; ++face) { + const auto bf = pmb->boundary_flag[face]; + const bool is_user = bf == parthenon::BoundaryFlag::user; + const bool is_block = bf == parthenon::BoundaryFlag::block; + const bool is_periodic = bf == parthenon::BoundaryFlag::periodic; + PARTHENON_REQUIRE(is_user || is_block || is_periodic, + "HYPRE SetupGrid encountered unsupported BoundaryFlag."); + // This solver currently treats user boundaries as physical Dirichlet boundaries. + is_domain[face] = is_user; + } + block_is_domain_boundary[b] = is_domain; + + for (const auto &nb : pmb->GetNeighbors()) { + const int ax = std::abs(nb.offsets(parthenon::X1DIR)); + const int ay = std::abs(nb.offsets(parthenon::X2DIR)); + const int az = std::abs(nb.offsets(parthenon::X3DIR)); + if (ax + ay + az != 1) continue; + + const int face = FaceFromOffsets(nb.offsets); + if (face == parthenon::BoundaryFace::undef) continue; + + const auto nlegacy_loc = pmesh->Forest().GetLegacyTreeLocation(nb.origin_loc); + const int nlev = static_cast(nlegacy_loc.level()) - legacy_root_level; + if (nlev > lev) { + block_neighbor_level[b][face] = parthenon::CellLevel::fine; + } else if (nlev < lev) { + block_neighbor_level[b][face] = parthenon::CellLevel::coarse; + } else { + block_neighbor_level[b][face] = parthenon::CellLevel::same; + } + } + } + + { + const int local_nboxes = nblocks; + std::vector counts(parthenon::Globals::nranks, 0); + MPI_Allgather(&local_nboxes, 1, MPI_INT, counts.data(), 1, MPI_INT, MPI_COMM_WORLD); + + std::vector displs(parthenon::Globals::nranks, 0); + int total_nboxes = 0; + for (int r = 0; r < parthenon::Globals::nranks; ++r) { + displs[r] = total_nboxes; + total_nboxes += counts[r]; + } + + std::vector sendbuf(static_cast(local_nboxes * 7), 0); + for (int b = 0; b < nblocks; ++b) { + sendbuf[7 * b + 0] = block_part[b]; + sendbuf[7 * b + 1] = block_ilower[b][0]; + sendbuf[7 * b + 2] = block_ilower[b][1]; + sendbuf[7 * b + 3] = block_ilower[b][2]; + sendbuf[7 * b + 4] = block_iupper[b][0]; + sendbuf[7 * b + 5] = block_iupper[b][1]; + sendbuf[7 * b + 6] = block_iupper[b][2]; + } + + std::vector recv_counts(parthenon::Globals::nranks, 0); + std::vector recv_displs(parthenon::Globals::nranks, 0); + for (int r = 0; r < parthenon::Globals::nranks; ++r) { + recv_counts[r] = counts[r] * 7; + recv_displs[r] = displs[r] * 7; + } + + std::vector recvbuf(static_cast(total_nboxes * 7), 0); + MPI_Allgatherv(sendbuf.data(), local_nboxes * 7, MPI_INT, recvbuf.data(), + recv_counts.data(), recv_displs.data(), MPI_INT, MPI_COMM_WORLD); + + for (int n = 0; n < total_nboxes; ++n) { + const int part = recvbuf[7 * n + 0]; + std::array lo_g{recvbuf[7 * n + 1], recvbuf[7 * n + 2], recvbuf[7 * n + 3]}; + std::array hi_g{recvbuf[7 * n + 4], recvbuf[7 * n + 5], recvbuf[7 * n + 6]}; + PARTHENON_REQUIRE(part >= 0 && part < nparts, + "Invalid part while constructing global_part_boxes."); + global_part_boxes[part].push_back({lo_g, hi_g}); + } + } + + // Variables and periodicity for each part. + HYPRE_SStructVariable cell_var = HYPRE_SSTRUCT_VARIABLE_CELL; + part_periodic.resize(nparts, {0, 0, 0}); + for (int part = 0; part < nparts; ++part) { + HYPRE_SStructGridSetVariables(grid, part, 1, &cell_var); + + std::array periodic{0, 0, 0}; + const int lev = part_to_level[part]; + if (pmesh->mesh_bcs[parthenon::BoundaryFace::inner_x1] == + parthenon::BoundaryFlag::periodic && + pmesh->mesh_bcs[parthenon::BoundaryFace::outer_x1] == + parthenon::BoundaryFlag::periodic) { + periodic[0] = pmesh->mesh_size.nx(parthenon::X1DIR) * (1 << lev); + } + if (ndim > 1 && + pmesh->mesh_bcs[parthenon::BoundaryFace::inner_x2] == + parthenon::BoundaryFlag::periodic && + pmesh->mesh_bcs[parthenon::BoundaryFace::outer_x2] == + parthenon::BoundaryFlag::periodic) { + periodic[1] = pmesh->mesh_size.nx(parthenon::X2DIR) * (1 << lev); + } + if (ndim > 2 && + pmesh->mesh_bcs[parthenon::BoundaryFace::inner_x3] == + parthenon::BoundaryFlag::periodic && + pmesh->mesh_bcs[parthenon::BoundaryFace::outer_x3] == + parthenon::BoundaryFlag::periodic) { + periodic[2] = pmesh->mesh_size.nx(parthenon::X3DIR) * (1 << lev); + } + + // Passing zeros is the Hypre convention for non-periodic directions. + HYPRE_SStructGridSetPeriodic(grid, part, periodic.data()); + part_periodic[part] = periodic; + } + + HYPRE_SStructGridAssemble(grid); + + // 3-point (1D), 5-point (2D) or 7-point (3D) stencil. + HYPRE_SStructStencilCreate(ndim, nstencil, &stencil); + int var = 0; + std::array off{0, 0, 0}; + HYPRE_SStructStencilSetEntry(stencil, 0, off.data(), var); + off = {-1, 0, 0}; + HYPRE_SStructStencilSetEntry(stencil, 1, off.data(), var); + off = {1, 0, 0}; + if (ndim > 1) { + HYPRE_SStructStencilSetEntry(stencil, 2, off.data(), var); + off = {0, -1, 0}; + HYPRE_SStructStencilSetEntry(stencil, 3, off.data(), var); + off = {0, 1, 0}; + HYPRE_SStructStencilSetEntry(stencil, 4, off.data(), var); + } + if (ndim > 2) { + off = {0, 0, -1}; + HYPRE_SStructStencilSetEntry(stencil, 5, off.data(), var); + off = {0, 0, 1}; + HYPRE_SStructStencilSetEntry(stencil, 6, off.data(), var); + } + + HYPRE_SStructGraphCreate(MPI_COMM_WORLD, grid, &graph); + HYPRE_SStructGraphSetObjectType(graph, HYPRE_PARCSR); + for (int part = 0; part < nparts; ++part) { + HYPRE_SStructGraphSetStencil(graph, part, 0, stencil); + } + + // Add non-stencil graph entries across fine-coarse boundaries. + for (int b = 0; b < nblocks; ++b) { + const auto *pmb = blocks[b].get(); + const auto legacy_loc = pmesh->Forest().GetLegacyTreeLocation(pmb->loc); + const int lev = static_cast(legacy_loc.level()) - legacy_root_level; + const int part = block_part[b]; + const auto &lo = block_ilower[b]; + const auto &hi = block_iupper[b]; + + auto add_entry = [&](const std::array &from, std::array to, + int to_part) { + // Wrap periodic target indices into [0, period) before validation and Hypre call. + // HYPRE_SStructGraphAddEntries does NOT auto-wrap periodic indices. + const auto &pp = part_periodic[to_part]; + for (int d = 0; d < ndim; ++d) { + if (pp[d] > 0) { + to[d] = ((to[d] % pp[d]) + pp[d]) % pp[d]; + } + } + + bool valid_to = false; + for (const auto &bx : global_part_boxes[to_part]) { + const auto &lbx = bx.first; + const auto &ubx = bx.second; + if (to[0] >= lbx[0] && to[0] <= ubx[0] && to[1] >= lbx[1] && to[1] <= ubx[1] && + to[2] >= lbx[2] && to[2] <= ubx[2]) { + valid_to = true; + break; + } + } + if (!valid_to) { + const int lev_to = part_to_level[to_part]; + std::stringstream msg; + msg << "SetupGrid graph target index not found in any registered box: from=(" + << from[0] << "," << from[1] << "," << from[2] << ") to=(" << to[0] << "," + << to[1] << "," << to[2] << ") to_part=" << to_part << " to_level=" << lev_to; + PARTHENON_FAIL(msg); + } + HYPRE_SStructGraphAddEntries(graph, part, const_cast(from.data()), 0, + to_part, to.data(), 0); + }; + + for (const auto &nb : pmb->GetNeighbors()) { + const int ax = std::abs(nb.offsets(parthenon::X1DIR)); + const int ay = std::abs(nb.offsets(parthenon::X2DIR)); + const int az = std::abs(nb.offsets(parthenon::X3DIR)); + if (ax + ay + az != 1) continue; + + const int face = FaceFromOffsets(nb.offsets); + if (face == parthenon::BoundaryFace::undef) continue; + + const auto nlegacy_loc = pmesh->Forest().GetLegacyTreeLocation(nb.origin_loc); + const int nlev = static_cast(nlegacy_loc.level()) - legacy_root_level; + const int relative_nbr_level = (nlev > lev) ? 1 : ((nlev < lev) ? -1 : 0); + if (relative_nbr_level == 0) continue; + + const int to_level = lev + relative_nbr_level; + PARTHENON_REQUIRE(to_level >= 0 && + to_level < static_cast(level_to_part.size()), + "Invalid neighbor level mapping in SetupGrid."); + const int to_part = level_to_part[to_level]; + PARTHENON_REQUIRE(to_part >= 0, "Invalid neighbor part mapping in SetupGrid."); + + int is, ie, js, je, ks, ke, axis, side; + NeighborFaceBounds(lo, hi, ndim, nb, relative_nbr_level > 0, is, ie, js, je, ks, ke, + axis, side); + PARTHENON_REQUIRE(axis >= 0 && axis < ndim, + "Invalid face axis in SetupGrid graph construction."); + + const int ni = ie - is + 1; + const int nj = je - js + 1; + const int nk = ke - ks + 1; + const int nface_cells = ni * nj * nk; + std::vector> from_cells(static_cast(nface_cells)); + + parthenon::seq_for(ks, ke, js, je, is, ie, + [&](const int k, const int j, const int i) { + const int lin = (k - ks) * nj * ni + (j - js) * ni + (i - is); + from_cells[lin] = {i, j, k}; + }); + + for (const auto &from : from_cells) { + if (relative_nbr_level < 0) { + std::array to{from[0] / 2, from[1] / 2, from[2] / 2}; + to[axis] = (from[axis] + ((side < 0) ? -1 : 1)) / 2; + add_entry(from, to, to_part); + continue; + } + + std::array base = {2 * from[0], 2 * from[1], 2 * from[2]}; + base[axis] += (side < 0) ? -1 : 2; + const int t0 = (axis + 1) % ndim; + const int t1 = (axis + 2) % ndim; + + if (ndim == 2) { + for (int s = 0; s < 2; ++s) { + std::array to = base; + to[t0] = 2 * from[t0] + s; + add_entry(from, to, to_part); + } + } else { + for (int s0 = 0; s0 < 2; ++s0) { + for (int s1 = 0; s1 < 2; ++s1) { + std::array to = base; + to[t0] = 2 * from[t0] + s0; + to[t1] = 2 * from[t1] + s1; + add_entry(from, to, to_part); + } + } + } + } + } + } + + HYPRE_SStructGraphAssemble(graph); + + HYPRE_SStructMatrixCreate(MPI_COMM_WORLD, graph, &A); + HYPRE_SStructMatrixSetObjectType(A, HYPRE_PARCSR); + HYPRE_SStructMatrixInitialize(A); + + HYPRE_SStructVectorCreate(MPI_COMM_WORLD, grid, &b); + HYPRE_SStructVectorSetObjectType(b, HYPRE_PARCSR); + HYPRE_SStructVectorInitialize(b); + + HYPRE_SStructVectorCreate(MPI_COMM_WORLD, grid, &x); + HYPRE_SStructVectorSetObjectType(x, HYPRE_PARCSR); + HYPRE_SStructVectorInitialize(x); + + grid_is_setup = true; + needs_grid_setup = false; + solver_is_setup = false; +} + +} // namespace diffusion_package diff --git a/example/diffusion/diffusion_hypre.hpp b/example/diffusion/diffusion_hypre.hpp new file mode 100644 index 0000000000000..53251d749a47f --- /dev/null +++ b/example/diffusion/diffusion_hypre.hpp @@ -0,0 +1,160 @@ +#ifndef EXAMPLE_DIFFUSION_DIFFUSION_HYPRE_ +#define EXAMPLE_DIFFUSION_DIFFUSION_HYPRE_ +#include "basic_types.hpp" +#include "mesh/mesh.hpp" +#include "mesh/meshblock.hpp" +#include "tasks/tasks.hpp" +#ifdef DIFFUSION_WITH_HYPRE + +#include +#include +#include + +#include "HYPRE_krylov.h" +#include "HYPRE_parcsr_ls.h" +#include "HYPRE_sstruct_ls.h" +#include "HYPRE_sstruct_mv.h" +#include "parameter_input.hpp" + +namespace diffusion_package { + +struct HypreSolver { + // class to hold information we need for the hypre solves + // main hypre data members: + // * hypre solvers -- preconditioner and actual solver + // * hypre matrix -- matrix we will build, A, for solving Ax=b + // * hypre vectors -- solution (x) and rhs (b) + // * hypre grid -- we will use the SStruct interface to map parthenon's tree mesh to + // hypre + // * hypre stencil -- used for same level couplings + // + // we will also require some cached information about the grid for mapping parthenon + // meshblocks to thier sstruct hypre counterparts + // + // * Basic idea is to treat each AMR refinement level as a unique hypre part. + // * We use parthenon's legacy logical location to map blocks to unique locations on the + // part + // using a corner ID that maps the meshblock's lower left corner cell to the (k,j,i) + // index of that cell if the entire domain was at that meshblock's refinement level + // * we use stencil couplings within a meshblock as well as between meshblocks at the + // same level + // * at F-C boundaries we create graph entries that map to our coarse/fine neighbors to + // replace the stencil + // couplings that would have gone to the other meshblock + // + // per meshblock we will need + // * part number + (k,j,i) corner ID + // * neighbor information (relative refinement level of our face neighbors), can use + // the parthenon::CellLevel enum + + // --------------------------------------------------------------------------- + // Per-block metadata cached at grid setup time (SoA layout) + // Indexed by block list position (same order as pmesh->block_list) + // --------------------------------------------------------------------------- + std::vector block_part; // Hypre part number + std::vector> block_ilower; // lower corner ID (global cell index) + std::vector> block_iupper; // upper corner ID (ilower + NX - 1) + // Relative refinement level of each face neighbor per block: + // CellLevel::same (0), fine (+1), coarse (-1) + // Inner index: BoundaryFace (inner_x1=0 .. outer_x3=5) + std::vector> block_neighbor_level; + // Whether each face is a physical (domain) boundary + std::vector> block_is_domain_boundary; + + // --------------------------------------------------------------------------- + // Hypre object handles + // --------------------------------------------------------------------------- + HYPRE_SStructGrid grid = nullptr; + HYPRE_SStructStencil stencil = nullptr; + HYPRE_SStructGraph graph = nullptr; + HYPRE_SStructMatrix A = nullptr; + HYPRE_SStructVector b = nullptr; // RHS vector + HYPRE_SStructVector x = nullptr; // solution vector + HYPRE_Solver solver_handle = nullptr; + HYPRE_Solver precond_handle = nullptr; + + // --------------------------------------------------------------------------- + // Grid / solver state + // --------------------------------------------------------------------------- + bool grid_is_setup = false; + bool solver_is_setup = false; + bool needs_grid_setup = true; + + int ndim = 2; // number of spatial dimensions (2 for this problem) + int nparts = 0; // number of distinct AMR levels with leaf blocks + int nstencil = 5; // stencil size (5 for 2D, 7 for 3D) + int min_active_level = -1; + int max_active_level = -1; + + // Map from AMR refinement level -> Hypre part index (0-based) + // Sized to max_level+1, indexed directly by level + std::vector level_to_part; + + // Per-part periodic lengths: part_periodic[part][d] = total cells at that level + // if direction d is periodic, else 0. Used to wrap FC graph target indices. + std::vector> part_periodic; + + // --------------------------------------------------------------------------- + // Solver configuration (read from [hypre] input block) + // --------------------------------------------------------------------------- + std::string solver_type; // "pcg" or "bicgstab" + std::string preconditioner; // "amg" or "none" + parthenon::Real tol; // relative convergence tolerance + parthenon::Real absolute_tol; // optional absolute convergence tolerance (0 disables) + int max_iter; // maximum solver iterations + int print_level; // solver verbosity + int internal_print_level; // Hypre internal solver print level + bool pcg_use_two_norm; // use true (2-) norm in PCG stopping criterion + bool pcg_recompute_residual; // recompute residual explicitly in PCG + + // BoomerAMG preconditioner settings + int amg_coarsen_type; // HMIS coarsening + int amg_interp_type; // ext+i interpolation + int amg_relax_type; // symmetric Gauss-Seidel + parthenon::Real amg_strong_threshold; // strong threshold (0.25 for 2D) + int amg_num_sweeps; // sweeps per AMG level + + // Debug controls + bool print_matrix = false; + bool print_vectors = false; + bool print_matrix_all = false; + int solve_call_count = 0; + + // Problem parameters cached from package + parthenon::Real diagonal_alpha; + std::array boundary_u{}; + + int niter{0}; + parthenon::Real rnorm{0.0}; + // --------------------------------------------------------------------------- + // Methods + // --------------------------------------------------------------------------- + + // initialize all our settings from the parameter input + HypreSolver(parthenon::ParameterInput *pin); + + // cleanup hypre objects + ~HypreSolver(); + + // adds all the meshblocks to the hypre grid via the sstruct interface + void SetupGrid(parthenon::Mesh *pmesh); + void DestroyGrid(); + void MarkGridDirty() { needs_grid_setup = true; } + + void SetupSolver(); + + using Real = parthenon::Real; + + // calls hypre api to set matrix rows for a single mesh block + // as well as the initial guess and rhs vectors + static parthenon::TaskStatus + BuildMatrixVector(HypreSolver *solver, int b, parthenon::MeshBlock *pmb, const Real dt); + // call the hypre solve (maybe we even setup the solvers here) + static parthenon::TaskStatus Solve(HypreSolver *solver); + static parthenon::TaskStatus UpdateSolution(HypreSolver *solver, int b, + parthenon::MeshBlock *pmb); +}; + +} // namespace diffusion_package +#endif // DIFFUSION_WITH_HYPRE +#endif // EXAMPLE_DIFFUSION_DIFFUSION_HYPRE_ diff --git a/example/diffusion/diffusion_package.cpp b/example/diffusion/diffusion_package.cpp index 5eeebbd457b38..b833129c5120c 100644 --- a/example/diffusion/diffusion_package.cpp +++ b/example/diffusion/diffusion_package.cpp @@ -29,6 +29,7 @@ #include "defs.hpp" #include "diffusion_equation.hpp" +#include "diffusion_hypre.hpp" #include "diffusion_package.hpp" #include "kokkos_abstraction.hpp" @@ -116,6 +117,10 @@ std::shared_ptr Initialize(ParameterInput *pin) { Real t0 = pin->GetOrAddReal("diffusion", "t0", 0.001); pkg->AddParam<>("t0", t0); + Real rel_res = + pin->GetOrAddReal("diffusion/solver_params", "relative_residual_tolerance", 0.0); + pkg->AddParam<>("rel_res", rel_res); + bool constant_coeff = pin->GetOrAddBoolean("diffusion", "constant_coefficient", true); pkg->AddParam<>("constant_coefficient", constant_coeff); @@ -129,7 +134,7 @@ std::shared_ptr Initialize(ParameterInput *pin) { using PoissEq = diffusion_package::DiffusionEquation; pkg->AddParam<>("diffusion_equation", std::make_shared(pin, "diffusion")); - + bool report_timings = pin->GetOrAddBoolean("diffusion", "report_timings", false, "Report different timings of the diffusion solver " @@ -191,6 +196,25 @@ std::shared_ptr Initialize(ParameterInput *pin) { // timestep pkg->EstimateTimestepMesh = EstimateTimestep; + bool use_hypre = pin->GetOrAdd("hypre", "use_hypre", false, + "Use HYPRE solvers, requires -DDIFFUSION_WITH_HYPRE"); + pkg->AddParam("use_hypre", use_hypre && WithHypre()); +#ifdef DIFFUSION_WITH_HYPRE + { + auto hypre_solver = std::make_shared(pin); + pkg->AddParam("hypre_solver", hypre_solver); + + if (use_hypre) { + // add a flux variable with 4 components to communicate fine face D up to the + // coarser blocks + auto m = Metadata( + {Metadata::Flux, Metadata::Face, Metadata::OneCopy, Metadata::CellMemAligned}, + std::vector{4}); + pkg->AddField(m); + } + } +#endif // DIFFUSION_WITH_HYPRE + return pkg; } @@ -295,5 +319,56 @@ parthenon::TaskStatus SetDiffusionCoefficient(std::shared_ptr> md } return TaskStatus::complete; } // SetRHS +parthenon::TaskStatus SetDiffusionCoefficientHypre(std::shared_ptr> md, + const Real dt) { +#ifdef DIFFUSION_WITH_HYPRE + { + const int ndim = md->GetMeshPointer()->ndim; + auto pkg = md->GetMeshPointer()->packages.Get("diffusion_package"); + auto desc = parthenon::MakePackDescriptor(md.get()); + auto pack = desc.GetPack(md.get()); + + const auto profile_D = pkg->Param("diffusion_coefficient"); + + for (auto te : {TE::F1, TE::F2, TE::F3}) { + IndexRange ib = md->GetBoundsI(IndexDomain::interior, te); + IndexRange jb = md->GetBoundsJ(IndexDomain::interior, te); + IndexRange kb = md->GetBoundsK(IndexDomain::interior, te); + + int offset_x1 = te == TE::F1 ? 1 : 0; + int offset_x2 = te == TE::F2 && (ndim > 1) ? 1 : 0; + int offset_x3 = te == TE::F3 && (ndim > 2) ? 1 : 0; + + parthenon::par_for( + "SetDiffusionCoefficient", 0, pack.GetNBlocks() - 1, kb.s, kb.e, jb.s, jb.e, + ib.s, ib.e, KOKKOS_LAMBDA(const int b, const int k, const int j, const int i) { + const Real u = 0.5 * (pack(b, TE::CC, diffusion_package::u(), k - offset_x3, + j - offset_x2, i - offset_x1) + + pack(b, TE::CC, diffusion_package::u(), k, j, i)); + + const auto &coords = pack.GetCoordinates(b); + const Real x1 = + offset_x1 ? coords.Xc<1>(k, j, i) : coords.X<1, TE::F1>(k, j, i); + const Real x2 = + offset_x2 ? coords.Xc<2>(k, j, i) : coords.X<2, TE::F2>(k, j, i); + const Real x3 = + offset_x3 ? coords.Xc<3>(k, j, i) : coords.X<3, TE::F3>(k, j, i); + pack(b, te, diffusion_package::D(), k, j, i) = dt * profile_D(x1, x2, x3); + + auto face_comp = [&](const int k, const int j, const int i) { + const int comp1 = offset_x1 ? j % 2 : (offset_x2 ? k % 2 : i % 2); + const int comp2 = offset_x1 ? k % 2 : (offset_x2 ? i % 2 : j % 2); + return comp1 + 2 * comp2; + }; + + pack(b, te, diffusion_package::Dfc(face_comp(k, j, i)), k, j, i) = + pack(b, te, diffusion_package::D(), k, j, i); + }); + } + } +#endif // DIFFUSION_WITH_HYPRE + return TaskStatus::complete; +} // SetDiffusionCoefficientHypre } // namespace diffusion_package diff --git a/example/diffusion/diffusion_package.hpp b/example/diffusion/diffusion_package.hpp index 5a0a42e339ca8..d584f38658b1a 100644 --- a/example/diffusion/diffusion_package.hpp +++ b/example/diffusion/diffusion_package.hpp @@ -30,10 +30,18 @@ } namespace diffusion_package { +constexpr bool WithHypre() { +#ifdef DIFFUSION_WITH_HYPRE + return true; +#else + return false; +#endif +} using namespace parthenon::package::prelude; VARIABLE(diffusion, D); VARIABLE(diffusion, u); +VARIABLE(diffusion, Dfc); struct DiffusionCoefficient { Real Dright{1.0}; @@ -70,6 +78,8 @@ TaskStatus SetRHS(std::shared_ptr> md, std::shared_ptr> md_rhs); parthenon::TaskStatus SetDiffusionCoefficient(std::shared_ptr> md, const Real dt); +parthenon::TaskStatus SetDiffusionCoefficientHypre(std::shared_ptr> md, + const Real dt); Real EstimateTimestep(MeshData *md); } // namespace diffusion_package diff --git a/example/diffusion/main.cpp b/example/diffusion/main.cpp index d51c726c5c201..adc2c32232d24 100644 --- a/example/diffusion/main.cpp +++ b/example/diffusion/main.cpp @@ -11,6 +11,7 @@ // the public, perform publicly and display publicly, and to permit others to do so. //======================================================================================== +#include "mesh/mesh.hpp" #include "parthenon_manager.hpp" #include "diffusion_driver.hpp" diff --git a/example/diffusion/parthinput.diffusion b/example/diffusion/parthinput.diffusion index 0898cf2ed6a98..e8ad30dbf8b84 100644 --- a/example/diffusion/parthinput.diffusion +++ b/example/diffusion/parthinput.diffusion @@ -100,7 +100,24 @@ absolute_residual_tolerance = 1.e-16 relative_residual_tolerance = 1.e-8 print_per_step = false presmoother = SRJ1 -postsmoother = SRJ3 +postsmoother = SRJ2 do_FAS = true -block_interior_prolongation = Kwak +block_interior_prolongation = Constant volume_weight = true + + +use_hypre = false +solver_type = bicgstab +preconditioner = amg +tol = 1e-8 +max_iter = 100 +print_level = 1 +internal_print_level = 0 +amg_coarsen_type = 10 +amg_interp_type = 6 +amg_relax_type = 6 +amg_strong_threshold = 0.25 +amg_num_sweeps = 1 +print_matrix = false +print_vectors = false +print_matrix_all = false diff --git a/example/diffusion/plot_steps.py b/example/diffusion/plot_steps.py new file mode 100644 index 0000000000000..130a4568cb67f --- /dev/null +++ b/example/diffusion/plot_steps.py @@ -0,0 +1,218 @@ +# ======================================================================================== +# (C) (or copyright) 2026. Triad National Security, LLC. All rights reserved. +# +# This program was produced under U.S. Government contract 89233218CNA000001 for Los +# Alamos National Laboratory (LANL), which is operated by Triad National Security, LLC +# for the U.S. Department of Energy/National Nuclear Security Administration. All rights +# in the program are reserved by Triad National Security, LLC, and the U.S. Department +# of Energy/National Nuclear Security Administration. The Government is granted for +# itself and others acting on its behalf a nonexclusive, paid-up, irrevocable worldwide +# license in this material to reproduce, prepare derivative works, distribute copies to +# the public, perform publicly and display publicly, and to permit others to do so. +# ======================================================================================== + +#!/usr/bin/env python3 +""" +[This code was generated with the help of generative AI] +Plot zone-cycles/wsec_step and v-cycles versus step for one or more Parthenon log files. + +Usage: + python plot_steps.py run1.log run2.log + python plot_steps.py *.out --output perf.png + +Plot style: + - Left y-axis: zone-cycles/wsec_step (solid) + - Right y-axis: v-cycles (dashed) + - One color per run + - Compact legend: + * one legend for run colors + * one legend for line styles +""" + +from __future__ import annotations + +import argparse +import pathlib +import re +import sys +from dataclasses import dataclass + +import matplotlib.pyplot as plt +from matplotlib.lines import Line2D + +plt.rcParams["axes.prop_cycle"] = plt.cycler( + color=[ + "#0072B2", + "#E69F00", + "#009E73", + "#CC79A7", + "#56B4E9", + "#D55E00", + "#F0E442", + "#000000", + ] +) + +CYCLE_RE = re.compile( + r""" + \bcycle=(?P\d+) + .*? + \bzone-cycles/wsec_step=(?P[0-9.eE+\-]+) + .*? + \bv-cycles=(?P\d+) + """, + re.VERBOSE, +) + + +@dataclass +class RunData: + label: str + cycles: list[int] + zc_per_sec: list[float] + vcycles: list[int] + + +def parse_log(path: pathlib.Path, include_zero: bool = False) -> RunData: + cycles: list[int] = [] + zc_per_sec: list[float] = [] + vcycles: list[int] = [] + + with path.open("r", encoding="utf-8", errors="replace") as f: + for line in f: + m = CYCLE_RE.search(line) + if not m: + continue + + cycle = int(m.group("cycle")) + if cycle == 0 and not include_zero: + continue + + cycles.append(cycle) + zc_per_sec.append(float(m.group("zc"))) + vcycles.append(int(m.group("vc"))) + + if not cycles: + raise ValueError(f"No matching cycle lines found in {path}") + + return RunData( + label=path.stem, + cycles=cycles, + zc_per_sec=zc_per_sec, + vcycles=vcycles, + ) + + +def make_plot( + runs: list[RunData], output: str | None = None, title: str | None = None +) -> None: + fig, ax1 = plt.subplots(figsize=(10, 6)) + ax2 = ax1.twinx() + + colors = plt.rcParams["axes.prop_cycle"].by_key()["color"] + + run_handles: list[Line2D] = [] + run_labels: list[str] = [] + + for i, run in enumerate(runs): + color = colors[i % len(colors)] + + ax1.semilogy( + run.cycles, + run.zc_per_sec, + color=color, + linewidth=2.0, + linestyle="-", + ) + ax2.plot( + run.cycles, + run.vcycles, + color=color, + linewidth=1.8, + linestyle="--", + ) + + run_handles.append(Line2D([0], [0], color=color, lw=2.5)) + run_labels.append(run.label) + + ax1.set_xlabel("Step") + ax1.set_ylabel("zone-cycles / s") + ax2.set_ylabel("v-cycles") + + if title: + ax1.set_title(title) + else: + ax1.set_title("Performance vs. step") + + ax1.grid(True, which="major", alpha=0.3) + + # Legend 1: colors = runs + legend_runs = ax1.legend( + run_handles, + run_labels, + title="Runs", + loc="upper left", + fontsize=9, + title_fontsize=10, + ) + ax1.add_artist(legend_runs) + + # Legend 2: line style meaning + style_handles = [ + Line2D([0], [0], color="black", lw=2.0, linestyle="-"), + Line2D([0], [0], color="black", lw=1.8, linestyle="--"), + ] + style_labels = [ + "zone-cycles / wsec_step", + "v-cycles", + ] + ax1.legend( + style_handles, + style_labels, + title="Line style", + loc="upper right", + fontsize=9, + title_fontsize=10, + ) + + fig.tight_layout() + + if output: + fig.savefig(output, dpi=200, bbox_inches="tight") + print(f"Wrote {output}") + else: + plt.show() + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("logs", nargs="+", help="Log files to read") + parser.add_argument( + "--output", "-o", help="Write figure to this file instead of showing it" + ) + parser.add_argument("--title", help="Custom plot title") + parser.add_argument( + "--include-zero", + action="store_true", + help="Include cycle=0 in the plot", + ) + args = parser.parse_args() + + runs: list[RunData] = [] + for log in args.logs: + path = pathlib.Path(log) + if not path.exists(): + print(f"File not found: {path}", file=sys.stderr) + return 1 + try: + runs.append(parse_log(path, include_zero=args.include_zero)) + except Exception as e: + print(f"Failed to parse {path}: {e}", file=sys.stderr) + return 1 + + make_plot(runs, output=args.output, title=args.title) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/solvers/bicgstab_solver.hpp b/src/solvers/bicgstab_solver.hpp index 38577b7ddebda..86bb110572087 100644 --- a/src/solvers/bicgstab_solver.hpp +++ b/src/solvers/bicgstab_solver.hpp @@ -38,32 +38,11 @@ namespace solvers { enum class Preconditioner { None, Diagonal, Multigrid }; struct BiCGSTABParams { MGParams mg_params; - int max_iters = 1000; - std::shared_ptr residual_tolerance = std::make_shared(1.e-12); - std::shared_ptr relative_residual_tolerance = std::make_shared(0.0); - std::shared_ptr absolute_residual_tolerance = std::make_shared(0.0); + Preconditioner precondition_type = Preconditioner::Multigrid; - bool print_per_step = false; - bool relative_residual = false; bool volume_weight = false; BiCGSTABParams() = default; BiCGSTABParams(ParameterInput *pin, const std::string &input_block) { - max_iters = pin->GetOrAddInteger(input_block, "max_iterations", max_iters); - *residual_tolerance = - pin->GetOrAddReal(input_block, "residual_tolerance", *residual_tolerance); - relative_residual = - pin->GetOrAddBoolean(input_block, "relative_residual", relative_residual); - if (relative_residual) { - *relative_residual_tolerance = pin->GetOrAddReal( - input_block, "relative_residual_tolerance", *residual_tolerance); - *absolute_residual_tolerance = - pin->GetOrAddReal(input_block, "absolute_residual_tolerance", 0.0); - } else { - *relative_residual_tolerance = - pin->GetOrAddReal(input_block, "relative_residual_tolerance", 0.0); - *absolute_residual_tolerance = pin->GetOrAddReal( - input_block, "absolute_residual_tolerance", *residual_tolerance); - } bool precondition = pin->GetOrAddBoolean(input_block, "precondition", true); std::string precondition_str = pin->GetOrAddString(input_block, "preconditioner", "Multigrid"); @@ -74,7 +53,6 @@ struct BiCGSTABParams { } else { precondition_type = Preconditioner::None; } - print_per_step = pin->GetOrAddBoolean(input_block, "print_per_step", print_per_step); mg_params = MGParams(pin, input_block); volume_weight = pin->GetOrAddBoolean(input_block, "volume_weight", volume_weight, @@ -108,8 +86,8 @@ class BiCGSTABSolver : public SolverBase, BiCGSTABSolverCounter { const std::string &input_block, equations_t eq_in = equations_t()) : preconditioner(container_base, container_u, container_rhs, pin, input_block, eq_in), - SolverBase(container_base, container_u, container_rhs), params_(pin, input_block), - iter_counter(0), eqs_(eq_in) { + SolverBase(container_base, container_u, container_rhs, pin, input_block), + params_(pin, input_block), iter_counter(0), eqs_(eq_in) { FieldTL::IterateTypes( [this](auto t) { this->sol_fields.push_back(decltype(t)::name()); }); std::string solver_id = "bicgstab" + std::to_string(id++); @@ -215,7 +193,7 @@ class BiCGSTABSolver : public SolverBase, BiCGSTABSolverCounter { TaskQualifier::once_per_region, initialize, "print to screen", [&](BiCGSTABSolver *solver, std::shared_ptr abs_res_tol, std::shared_ptr rel_res_tol, Mesh *pm) { - if (Globals::my_rank == 0 && params_.print_per_step) { + if (Globals::my_rank == 0 && print_per_step) { Real res_tol = *rel_res_tol * std::sqrt(solver->rhs2 / pm->GetTotalCells()); printf("# [0] v-cycle\n# [1] rms-residual (abs_tol = %e, rel_tol = %e) \n# " "[2] rms-error\n", @@ -224,11 +202,10 @@ class BiCGSTABSolver : public SolverBase, BiCGSTABSolverCounter { } return TaskStatus::complete; }, - this, params_.absolute_residual_tolerance, params_.relative_residual_tolerance, - pmesh); + this, absolute_residual_tolerance, relative_residual_tolerance, pmesh); // BEGIN ITERATIVE TASKS - auto [itl, solver_id] = tl.AddSublist(initialize, {1, params_.max_iters}); + auto [itl, solver_id] = tl.AddSublist(initialize, {1, max_iters}); auto sync = itl.AddTask(TaskQualifier::local_sync, none, []() { return TaskStatus::complete; }); @@ -294,7 +271,7 @@ class BiCGSTABSolver : public SolverBase, BiCGSTABSolverCounter { this, md_r, md_v, md_s); // Check and print out residual - if (params_.print_per_step) { + if (print_per_step) { auto get_res = DotProduct(correct_s, itl, &residual, md_s, md_s, params_.volume_weight); @@ -302,7 +279,7 @@ class BiCGSTABSolver : public SolverBase, BiCGSTABSolverCounter { TaskQualifier::once_per_region, get_res, [&](BiCGSTABSolver *solver, Mesh *pmesh) { Real rms_res = std::sqrt(solver->residual.val / pmesh->GetTotalCells()); - if (Globals::my_rank == 0 && solver->params_.print_per_step) + if (Globals::my_rank == 0 && solver->print_per_step) printf("%i %e\n", solver->iter_counter * 2 + 1, rms_res); return TaskStatus::complete; }, @@ -372,7 +349,7 @@ class BiCGSTABSolver : public SolverBase, BiCGSTABSolverCounter { TaskQualifier::once_per_region, get_res2_rhat0r, [&](BiCGSTABSolver *solver, Mesh *pmesh) { Real rms_err = std::sqrt(solver->res_rhat0r.val[0] / pmesh->GetTotalCells()); - if (Globals::my_rank == 0 && solver->params_.print_per_step) + if (Globals::my_rank == 0 && solver->print_per_step) printf("%i %e\n", solver->iter_counter * 2 + 2, rms_err); return TaskStatus::complete; }, @@ -411,8 +388,7 @@ class BiCGSTABSolver : public SolverBase, BiCGSTABSolverCounter { } return TaskStatus::iterate; }, - this, pmesh, params_.max_iters, params_.absolute_residual_tolerance, - params_.relative_residual_tolerance); + this, pmesh, max_iters, absolute_residual_tolerance, relative_residual_tolerance); timer_res->StopCollectingTasks(); return tl.AddTask(solver_id, TF(CopyData), md_x, md_u); } diff --git a/src/solvers/cg_solver.hpp b/src/solvers/cg_solver.hpp index 9cde01ff10460..075b351744d28 100644 --- a/src/solvers/cg_solver.hpp +++ b/src/solvers/cg_solver.hpp @@ -39,21 +39,11 @@ namespace solvers { struct CGParams { MGParams mg_params; - int max_iters = 1000; - std::shared_ptr residual_tolerance = std::make_shared(1.e-12); bool precondition = true; - bool print_per_step = false; - bool relative_residual = false; CGParams() = default; CGParams(ParameterInput *pin, const std::string &input_block) { - max_iters = pin->GetOrAddInteger(input_block, "max_iterations", max_iters); - *residual_tolerance = - pin->GetOrAddReal(input_block, "residual_tolerance", *residual_tolerance); precondition = pin->GetOrAddBoolean(input_block, "precondition", precondition); - print_per_step = pin->GetOrAddBoolean(input_block, "print_per_step", print_per_step); mg_params = MGParams(pin, input_block); - relative_residual = - pin->GetOrAddBoolean(input_block, "relative_residual", relative_residual); } }; @@ -80,7 +70,7 @@ class CGSolver : public SolverBase, CGSolverCounter { CGSolver(const std::string &container_base, const std::string &container_u, const std::string &container_rhs, ParameterInput *pin, const std::string &input_block, const equations_t &eq_in = equations_t()) - : SolverBase(container_base, container_u, container_rhs), + : SolverBase(container_base, container_u, container_rhs, pin, input_block), preconditioner(container_base, container_u, container_rhs, pin, input_block, eq_in), params_(pin, input_block), iter_counter(0), eqs_(eq_in) { @@ -138,7 +128,7 @@ class CGSolver : public SolverBase, CGSolverCounter { auto zero_p = tl.AddTask(dependence, TF(SetToZero), md_p); auto copy_r = tl.AddTask(dependence, TF(CopyData), md_rhs, md_r); auto get_rhs2 = none; - if (params_.relative_residual || params_.print_per_step) + if (relative_residual || print_per_step) get_rhs2 = DotProduct(dependence, tl, &rhs2, md_rhs, md_rhs); auto initialize = tl.AddTask( TaskQualifier::once_per_region | TaskQualifier::local_sync, @@ -150,7 +140,7 @@ class CGSolver : public SolverBase, CGSolverCounter { }, this); - if (params_.print_per_step && Globals::my_rank == 0) { + if (print_per_step && Globals::my_rank == 0) { initialize = tl.AddTask( TaskQualifier::once_per_region, initialize, "print to screen", [&](CGSolver *solver, std::shared_ptr res_tol, bool relative_residual, @@ -163,11 +153,11 @@ class CGSolver : public SolverBase, CGSolverCounter { printf("\t0 %e\n", std::sqrt(solver->rhs2.val / pm->GetTotalCells())); return TaskStatus::complete; }, - this, params_.residual_tolerance, params_.relative_residual, pmesh); + this, residual_tolerance, relative_residual, pmesh); } // BEGIN ITERATIVE TASKS - auto [itl, solver_id] = tl.AddSublist(initialize, {1, params_.max_iters}); + auto [itl, solver_id] = tl.AddSublist(initialize, {1, max_iters}); auto sync = itl.AddTask(TaskQualifier::local_sync, none, []() { return TaskStatus::complete; }); @@ -239,7 +229,7 @@ class CGSolver : public SolverBase, CGSolverCounter { TaskQualifier::once_per_region, get_res, [&](CGSolver *solver, Mesh *pmesh) { Real rms_res = std::sqrt(solver->residual.val / pmesh->GetTotalCells()); - if (Globals::my_rank == 0 && solver->params_.print_per_step) + if (Globals::my_rank == 0 && solver->print_per_step) printf("\t%i %e\n", solver->iter_counter + 1, rms_res); return TaskStatus::complete; }, @@ -262,8 +252,7 @@ class CGSolver : public SolverBase, CGSolverCounter { } return TaskStatus::iterate; }, - this, pmesh, params_.max_iters, params_.residual_tolerance, - params_.relative_residual); + this, pmesh, max_iters, residual_tolerance, relative_residual); return tl.AddTask(solver_id, TF(CopyData), md_x, md_u); } diff --git a/src/solvers/mg_solver.hpp b/src/solvers/mg_solver.hpp index 421f3b0b15a66..9730cd6e9cae1 100644 --- a/src/solvers/mg_solver.hpp +++ b/src/solvers/mg_solver.hpp @@ -39,8 +39,6 @@ namespace parthenon { namespace solvers { struct MGParams { - int max_iters = 1000; - Real residual_tolerance = 1.e-12; bool do_FAS = true; std::string presmoother = "SRJ2"; std::string postsmoother = "SRJ2"; @@ -49,9 +47,6 @@ struct MGParams { MGParams() = default; MGParams(ParameterInput *pin, const std::string &input_block) { - max_iters = pin->GetOrAddInteger(input_block, "max_iterations", max_iters); - residual_tolerance = - pin->GetOrAddReal(input_block, "residual_tolerance", residual_tolerance); do_FAS = pin->GetOrAddBoolean( input_block, "do_FAS", do_FAS, "Use the full approximation scheme in multigrid, required for amr."); @@ -102,20 +97,14 @@ class MGSolver : public SolverBase, MGSolverCounter { bool initial_guess_is_zero; bool constant_prolongation; + MGSolver(const std::string &container_base, const std::string &container_u, const std::string &container_rhs, ParameterInput *pin, const std::string &input_block, equations_t eq_in = equations_t()) - : MGSolver(container_base, container_u, container_rhs, MGParams(pin, input_block), - eq_in, prolongator_t(pin, input_block), restrictor_t(pin, input_block)) { - } - - MGSolver(const std::string &container_base, const std::string &container_u, - const std::string &container_rhs, MGParams params_in, - equations_t eq_in = equations_t(), prolongator_t prol_in = prolongator_t(), - restrictor_t rest_in = restrictor_t()) - : SolverBase(container_base, container_u, container_rhs), params_(params_in), - iter_counter(0), eqs_(eq_in), prolongator_(prol_in), initial_guess_is_zero{false}, - constant_prolongation{false}, restrictor_(rest_in) { + : SolverBase(container_base, container_u, container_rhs, pin, input_block), + params_(pin, input_block), iter_counter(0), eqs_(eq_in), + prolongator_(pin, input_block), initial_guess_is_zero{false}, + constant_prolongation{false}, restrictor_(pin, input_block) { FieldTL::IterateTypes( [this](auto t) { this->sol_fields.push_back(decltype(t)::name()); }); std::string solver_id = "mg" + std::to_string(id++); @@ -160,7 +149,7 @@ class MGSolver : public SolverBase, MGSolverCounter { Mesh *pmesh) { using namespace utils; TaskID none; - auto [itl, solve_id] = tl.AddSublist(dependence, {1, this->params_.max_iters}); + auto [itl, solve_id] = tl.AddSublist(dependence, {1, this->max_iters}); iter_counter = -1; auto update_iter = itl.AddTask( TaskQualifier::local_sync | TaskQualifier::once_per_region, none, "print", @@ -204,7 +193,7 @@ class MGSolver : public SolverBase, MGSolverCounter { printf("%i %e\n", solver->iter_counter, rms_res); solver->final_residual = rms_res; solver->final_iteration = solver->iter_counter; - if (rms_res > solver->params_.residual_tolerance) return TaskStatus::iterate; + if (rms_res > *(solver->residual_tolerance)) return TaskStatus::iterate; return TaskStatus::complete; }, this, pmesh); @@ -392,6 +381,9 @@ class MGSolver : public SolverBase, MGSolverCounter { // Damping factors from Yang & Mittal (2017) const std::array, 3> omega{ {{0.8723, 0.5395}, {1.3895, 0.5617}, {1.7319, 0.5695}}}; + // Chebyshev for diffusion + // const std::array, 3> omega{ + // {{0.576896, 2.159946}, {0.576896, 2.159946}, {0.576896, 2.159946}}}; auto jacobi1 = AddJacobiIteration(tl, depends_on, omega[ndim - 1][0], partition, pmesh, in_is_zero); return AddJacobiIteration(tl, jacobi1, omega[ndim - 1][1], partition, pmesh, false); @@ -399,6 +391,10 @@ class MGSolver : public SolverBase, MGSolverCounter { // Damping factors from Yang & Mittal (2017) const std::array, 3> omega{ {{0.9372, 0.6667, 0.5173}, {1.6653, 0.8000, 0.5264}, {2.2473, 0.8571, 0.5296}}}; + // Chebyshev for diffusion + // const std::array, 3> omega{ + // {{0.532079, 0.909091, 3.119832}, {0.532079, 0.909091, 3.119832}, {0.532079, + // 0.909091, 3.119832}}}; auto jacobi1 = AddJacobiIteration(tl, depends_on, omega[ndim - 1][0], partition, pmesh, in_is_zero); auto jacobi2 = diff --git a/src/solvers/solver_base.hpp b/src/solvers/solver_base.hpp index 29c182b2d02fd..fafdef996fdf9 100644 --- a/src/solvers/solver_base.hpp +++ b/src/solvers/solver_base.hpp @@ -51,9 +51,32 @@ struct has_RestrictGetOrAddBoolean( + input_block, "relative_residual", + false)), // Eventually needs to be deprecated + residual_tolerance( + std::make_shared(0.0)), // Eventually needs to be deprecated + absolute_residual_tolerance(std::make_shared(0.0)), + relative_residual_tolerance(std::make_shared(0.0)), + max_iters(pin->GetOrAddInteger(input_block, "max_iterations", 1000)), + print_per_step(pin->GetOrAddBoolean(input_block, "print_per_step", false)) { + *residual_tolerance = + pin->GetOrAddReal(input_block, "residual_tolerance", *residual_tolerance); + if (relative_residual) { + *relative_residual_tolerance = pin->GetOrAddReal( + input_block, "relative_residual_tolerance", *residual_tolerance); + *absolute_residual_tolerance = + pin->GetOrAddReal(input_block, "absolute_residual_tolerance", 0.0); + } else { + *relative_residual_tolerance = + pin->GetOrAddReal(input_block, "relative_residual_tolerance", 0.0); + *absolute_residual_tolerance = pin->GetOrAddReal( + input_block, "absolute_residual_tolerance", *residual_tolerance); + } + } virtual ~SolverBase() {} @@ -106,6 +129,13 @@ class SolverBase { static inline TimingAccumulatorDictionary solver_timings; + bool relative_residual; + std::shared_ptr residual_tolerance; + std::shared_ptr absolute_residual_tolerance; + std::shared_ptr relative_residual_tolerance; + int max_iters; + bool print_per_step; + protected: // Labels of all fields included in the vector std::vector sol_fields; diff --git a/src/solvers/tridiag_solver.hpp b/src/solvers/tridiag_solver.hpp index 8e168845febde..f7575279ed22d 100644 --- a/src/solvers/tridiag_solver.hpp +++ b/src/solvers/tridiag_solver.hpp @@ -58,8 +58,8 @@ class TridiagSolver : public SolverBase, TridiagSolverCounter { TridiagSolver(const std::string &container_base, const std::string &container_u, const std::string &container_rhs, ParameterInput *pin, const std::string &input_block, const equations &eq_in = equations()) - : SolverBase(container_base, container_u, container_rhs), iter_counter(0), - eqs_(eq_in), + : SolverBase(container_base, container_u, container_rhs, pin, input_block), + iter_counter(0), eqs_(eq_in), print_solution_(pin->GetOrAddBoolean(input_block, "print_solution", false)) { FieldTL::IterateTypes( [this](auto t) { this->sol_fields.push_back(decltype(t)::name()); });