From 3ee8e1ec9dfb9f0cb356c941fc2353a4c747a32c Mon Sep 17 00:00:00 2001
From: Pratik Nayak <pratikvn@protonmail.com>
Date: Mon, 5 Aug 2024 10:37:56 +0200
Subject: [PATCH] [bench] add distributed preconds

---
 benchmark/utils/preconditioners.hpp | 137 +++++++++++++++++++++++++++-
 1 file changed, 134 insertions(+), 3 deletions(-)

diff --git a/benchmark/utils/preconditioners.hpp b/benchmark/utils/preconditioners.hpp
index 63fd22708e6..c095f0af7e2 100644
--- a/benchmark/utils/preconditioners.hpp
+++ b/benchmark/utils/preconditioners.hpp
@@ -22,7 +22,11 @@ DEFINE_string(preconditioners, "none",
               "A comma-separated list of preconditioners to use. "
               "Supported values are: none, jacobi, paric, parict, parilu, "
               "parilut, ic, ilu, paric-isai, parict-isai, parilu-isai, "
-              "parilut-isai, ic-isai, ilu-isai, overhead");
+              "parilut-isai, ic-isai, ilu-isai, overhead"
+#ifdef GINKGO_BUILD_MPI
+              ", schwarz-jacobi, schwarz-ilu, schwarz-ic, schwarz-lu, dist-mg"
+#endif
+              "");
 
 DEFINE_uint32(parilu_iterations, 5,
               "The number of iterations for ParIC(T)/ParILU(T)");
@@ -292,12 +296,139 @@ const std::map<std::string, std::function<std::unique_ptr<gko::LinOpFactory>(
                  .with_sparsity_power(FLAGS_isai_power)
                  .on(exec);
          }},
-        {"overhead", [](std::shared_ptr<const gko::Executor> exec) {
+        {"overhead",
+         [](std::shared_ptr<const gko::Executor> exec) {
              return gko::Overhead<etype>::build()
                  .with_criteria(gko::stop::ResidualNorm<etype>::build()
                                     .with_reduction_factor(rc_etype{}))
                  .on(exec);
-         }}};
+         }}
+#ifdef GINKGO_BUILD_MPI
+        ,
+        {"schwarz-jacobi",
+         [](std::shared_ptr<const gko::Executor> exec) {
+             return gko::experimental::distributed::preconditioner::Schwarz<
+                        etype>::build()
+                 .with_local_solver(
+                     gko::preconditioner::Jacobi<etype>::build()
+                         .with_max_block_size(FLAGS_jacobi_max_block_size)
+                         .with_storage_optimization(
+                             parse_storage_optimization(FLAGS_jacobi_storage))
+                         .with_accuracy(
+                             static_cast<rc_etype>(FLAGS_jacobi_accuracy))
+                         .with_skip_sorting(true)
+                         .on(exec))
+                 .on(exec);
+         }},
+        {"schwarz-general-isai",
+         [](std::shared_ptr<const gko::Executor> exec) {
+             return gko::experimental::distributed::preconditioner::Schwarz<
+                        etype, itype>::build()
+                 .with_local_solver(
+                     gko::preconditioner::GeneralIsai<etype, itype>::build()
+                         .with_sparsity_power(FLAGS_isai_power)
+                         .on(exec))
+                 .on(exec);
+         }},
+        {"schwarz-spd-isai",
+         [](std::shared_ptr<const gko::Executor> exec) {
+             return gko::experimental::distributed::preconditioner::Schwarz<
+                        etype, itype>::build()
+                 .with_local_solver(
+                     gko::preconditioner::SpdIsai<etype, itype>::build()
+                         .with_sparsity_power(FLAGS_isai_power)
+                         .on(exec))
+                 .on(exec);
+         }},
+        {"schwarz-ilu",
+         [](std::shared_ptr<const gko::Executor> exec) {
+             auto fact =
+                 gko::share(gko::factorization::Ilu<etype, itype>::build()
+                                .with_skip_sorting(FLAGS_skip_sorting)
+                                .on(exec));
+             return gko::experimental::distributed::preconditioner::Schwarz<
+                        etype, itype>::build()
+                 .with_local_solver(gko::preconditioner::Ilu<
+                                        gko::solver::LowerTrs<etype, itype>,
+                                        gko::solver::UpperTrs<etype, itype>,
+                                        false, itype>::build()
+                                        .with_factorization(fact)
+                                        .on(exec))
+                 .on(exec);
+         }},
+        {"schwarz-ic",
+         [](std::shared_ptr<const gko::Executor> exec) {
+             auto fact =
+                 gko::share(gko::factorization::Ic<etype, itype>::build()
+                                .with_skip_sorting(FLAGS_skip_sorting)
+                                .on(exec));
+             return gko::experimental::distributed::preconditioner::Schwarz<
+                        etype, itype>::build()
+                 .with_local_solver(
+                     gko::preconditioner::Ic<
+                         gko::solver::LowerTrs<etype, itype>, itype>::build()
+                         .with_factorization(fact)
+                         .on(exec))
+                 .on(exec);
+         }},
+        {"schwarz-lu",
+         [](std::shared_ptr<const gko::Executor> exec) {
+             auto fact = gko::share(
+                 gko::experimental::factorization::Lu<etype, itype>::build().on(
+                     exec));
+             return gko::experimental::distributed::preconditioner::Schwarz<
+                        etype, itype>::build()
+                 .with_local_solver(
+                     gko::experimental::solver::Direct<etype, itype>::build()
+                         .with_factorization(fact)
+                         .on(exec))
+                 .on(exec);
+         }},
+        {"dist-mg",
+         [](std::shared_ptr<const gko::Executor> exec) {
+             using ir = gko::solver::Ir<etype>;
+             using schwarz =
+                 gko::experimental::distributed::preconditioner::Schwarz<etype,
+                                                                         itype>;
+             using bj = gko::preconditioner::Jacobi<etype, itype>;
+             auto iter_stop = gko::share(gko::stop::Iteration::build()
+                                             .with_max_iters(FLAGS_mg_max_iters)
+                                             .on(exec));
+             auto tol_stop =
+                 gko::share(gko::stop::ResidualNorm<etype>::build()
+                                .with_baseline(gko::stop::mode::absolute)
+                                .with_reduction_factor(FLAGS_mg_tolerance)
+                                .on(exec));
+             auto coarsest_solver_gen = gko::share(
+                 ir::build()
+                     .with_solver(schwarz::build().with_local_solver(
+                         bj::build().with_max_block_size(1u).with_skip_sorting(
+                             true)))
+                     .with_relaxation_factor(static_cast<etype>(0.9))
+                     .with_criteria(
+                         gko::stop::Iteration::build().with_max_iters(4u))
+                     .on(exec));
+             auto pre_smoother = gko::share(
+                 ir::build()
+                     .with_solver(schwarz::build().with_local_solver(
+                         bj::build().with_max_block_size(1u).with_skip_sorting(
+                             true)))
+                     .with_relaxation_factor(static_cast<etype>(0.9))
+                     .with_criteria(
+                         gko::stop::Iteration::build().with_max_iters(1u))
+                     .on(exec));
+             return gko::solver::Multigrid::build()
+                 .with_mg_level(gko::multigrid::Pgm<etype, itype>::build()
+                                    .with_deterministic(FLAGS_mg_deterministic))
+                 .with_criteria(iter_stop, tol_stop)
+                 .with_pre_smoother(pre_smoother)
+                 .with_post_uses_pre(true)
+                 .with_max_levels(FLAGS_mg_num_levels)
+                 .with_coarsest_solver(coarsest_solver_gen)
+                 .on(exec);
+         }}
+#endif
+    };
 
 
 #endif  // GKO_BENCHMARK_UTILS_PRECONDITIONERS_HPP_