Additive Schwarz coarse solver

core/distributed/preconditioner/schwarz.cpp

@@ -44,11 +44,37 @@ void Schwarz<ValueType, LocalIndexType, GlobalIndexType>::apply_dense_impl(
     const VectorType* dense_b, VectorType* dense_x) const
 {
     using Vector = matrix::Dense<ValueType>;
+    using dist_vec = experimental::distributed::Vector<ValueType>;
     auto exec = this->get_executor();
+
     if (this->local_solver_ != nullptr) {
         this->local_solver_->apply(gko::detail::get_local(dense_b),
                                    gko::detail::get_local(dense_x));
     }
+
+    if (this->coarse_solver_ != nullptr && this->galerkin_ops_ != nullptr) {
+        auto restrict = this->galerkin_ops_->get_restrict_op();
+        auto prolong = this->galerkin_ops_->get_prolong_op();
+        auto coarse =
+            as<experimental::distributed::Matrix<ValueType, LocalIndexType,
+                                                 GlobalIndexType>>(
+                this->galerkin_ops_->get_coarse_op());
+        auto comm = coarse->get_communicator();
+
+        auto cs_ncols = dense_x->get_size()[1];
+        auto cs_local_nrows = coarse->get_local_matrix()->get_size()[0];
+        auto cs_global_nrows = coarse->get_size()[0];
+        auto cs_local_size = dim<2>(cs_local_nrows, cs_ncols);
+        auto cs_global_size = dim<2>(cs_global_nrows, cs_ncols);
+        auto csol = dist_vec::create(exec, comm, cs_global_size, cs_local_size,
+                                     dense_x->get_stride());
+        restrict->apply(dense_b, csol);
+        auto tmp = csol->clone();
+        this->coarse_solver_->apply(csol, tmp);
+        auto one = gko::initialize<Vector>({0.5}, exec);
+        auto zero = gko::initialize<Vector>({0.5}, exec);
+        prolong->apply(one, tmp, zero, dense_x);
+    }
 }
 
 
@@ -94,23 +120,23 @@ void Schwarz<ValueType, LocalIndexType, GlobalIndexType>::generate(
         GKO_INVALID_STATE(
             "Requires either a generated solver or an solver factory");
     }
+    auto dist_mat =
+        as<experimental::distributed::Matrix<ValueType, LocalIndexType,
+                                             GlobalIndexType>>(system_matrix);
 
     if (parameters_.local_solver) {
-        this->set_solver(gko::share(parameters_.local_solver->generate(
-            as<experimental::distributed::Matrix<
-                ValueType, LocalIndexType, GlobalIndexType>>(system_matrix)
-                ->get_local_matrix())));
+        this->set_solver(gko::share(
+            parameters_.local_solver->generate(dist_mat->get_local_matrix())));
     } else {
         this->set_solver(parameters_.generated_local_solver);
     }
 
-    auto dist_mat =
-        as<experimental::distributed::Matrix<ValueType, LocalIndexType,
-                                             GlobalIndexType>>(system_matrix);
 
-    if (parameters_.coarse_solver_factory) {
-        this->coarse_solver_ = as<multigrid::MultigridLevel>(
-            share(parameters_.coarse_solver_factory->generate(dist_mat)));
+    if (parameters_.galerkin_ops_factory && parameters_.coarse_solver_factory) {
+        this->galerkin_ops_ = as<multigrid::MultigridLevel>(
+            share(parameters_.galerkin_ops_factory->generate(dist_mat)));
+        this->coarse_solver_ = parameters_.coarse_solver_factory->generate(
+            this->galerkin_ops_->get_coarse_op());
     }
 }
 

core/test/mpi/distributed/preconditioner/schwarz.cpp

@@ -9,6 +9,7 @@
 #include <ginkgo/core/matrix/csr.hpp>
 #include <ginkgo/core/multigrid/pgm.hpp>
 #include <ginkgo/core/preconditioner/jacobi.hpp>
+#include <ginkgo/core/solver/cg.hpp>
 
 #include "core/test/utils.hpp"
 
@@ -29,6 +30,7 @@ class SchwarzFactory : public ::testing::Test {
         value_type, local_index_type, global_index_type>;
     using Jacobi = gko::preconditioner::Jacobi<value_type, local_index_type>;
     using Pgm = gko::multigrid::Pgm<value_type, local_index_type>;
+    using Cg = gko::solver::Cg<value_type>;
     using Mtx =
         gko::experimental::distributed::Matrix<value_type, local_index_type,
                                                global_index_type>;
@@ -37,10 +39,12 @@ class SchwarzFactory : public ::testing::Test {
         : exec(gko::ReferenceExecutor::create()),
           jacobi_factory(Jacobi::build().on(exec)),
           pgm_factory(Pgm::build().on(exec)),
+          cg_factory(Cg::build().on(exec)),
           mtx(Mtx::create(exec, MPI_COMM_WORLD))
     {
         schwarz = Schwarz::build()
                       .with_local_solver(jacobi_factory)
+                      .with_galerkin_ops_factory(pgm_factory)
                       .with_coarse_solver_factory(pgm_factory)
                       .on(exec)
                       ->generate(mtx);
@@ -59,6 +63,8 @@ class SchwarzFactory : public ::testing::Test {
         ASSERT_EQ(a->get_size(), b->get_size());
         ASSERT_EQ(a->get_parameters().local_solver,
                   b->get_parameters().local_solver);
+        ASSERT_EQ(a->get_parameters().galerkin_ops_factory,
+                  b->get_parameters().galerkin_ops_factory);
         ASSERT_EQ(a->get_parameters().coarse_solver_factory,
                   b->get_parameters().coarse_solver_factory);
     }
@@ -67,6 +73,7 @@ class SchwarzFactory : public ::testing::Test {
     std::unique_ptr<Schwarz> schwarz;
     std::shared_ptr<typename Jacobi::Factory> jacobi_factory;
     std::shared_ptr<typename Pgm::Factory> pgm_factory;
+    std::shared_ptr<typename Cg::Factory> cg_factory;
     std::shared_ptr<Mtx> mtx;
 };
 
@@ -87,10 +94,17 @@ TYPED_TEST(SchwarzFactory, CanSetLocalFactory)
 }
 
 
+TYPED_TEST(SchwarzFactory, CanSetGalerkinOpsFactory)
+{
+    ASSERT_EQ(this->schwarz->get_parameters().galerkin_ops_factory,
+              this->pgm_factory);
+}
+
+
 TYPED_TEST(SchwarzFactory, CanSetCoarseSolverFactory)
 {
     ASSERT_EQ(this->schwarz->get_parameters().coarse_solver_factory,
-              this->pgm_factory);
+              this->cg_factory);
 }
 
 
@@ -106,13 +120,16 @@ TYPED_TEST(SchwarzFactory, CanBeCopied)
 {
     using Jacobi = typename TestFixture::Jacobi;
     using Pgm = typename TestFixture::Pgm;
+    using Cg = typename TestFixture::Cg;
     using Schwarz = typename TestFixture::Schwarz;
     using Mtx = typename TestFixture::Mtx;
     auto bj = gko::share(Jacobi::build().on(this->exec));
     auto pgm = gko::share(Pgm::build().on(this->exec));
+    auto cg = gko::share(Cg::build().on(this->exec));
     auto copy = Schwarz::build()
                     .with_local_solver(bj)
-                    .with_coarse_solver_factory(pgm)
+                    .with_galerkin_ops_factory(pgm)
+                    .with_coarse_solver_factory(cg)
                     .on(this->exec)
                     ->generate(Mtx::create(this->exec, MPI_COMM_WORLD));
 
@@ -126,14 +143,17 @@ TYPED_TEST(SchwarzFactory, CanBeMoved)
 {
     using Jacobi = typename TestFixture::Jacobi;
     using Pgm = typename TestFixture::Pgm;
+    using Cg = typename TestFixture::Cg;
     using Schwarz = typename TestFixture::Schwarz;
     using Mtx = typename TestFixture::Mtx;
     auto tmp = clone(this->schwarz);
     auto bj = gko::share(Jacobi::build().on(this->exec));
     auto pgm = gko::share(Pgm::build().on(this->exec));
+    auto cg = gko::share(Cg::build().on(this->exec));
     auto copy = Schwarz::build()
                     .with_local_solver(bj)
-                    .with_coarse_solver_factory(pgm)
+                    .with_galerkin_ops_factory(pgm)
+                    .with_coarse_solver_factory(cg)
                     .on(this->exec)
                     ->generate(Mtx::create(this->exec, MPI_COMM_WORLD));
 
@@ -149,6 +169,7 @@ TYPED_TEST(SchwarzFactory, CanBeCleared)
 
     ASSERT_EQ(this->schwarz->get_size(), gko::dim<2>(0, 0));
     ASSERT_EQ(this->schwarz->get_parameters().local_solver, nullptr);
+    ASSERT_EQ(this->schwarz->get_parameters().galerkin_ops_factory, nullptr);
     ASSERT_EQ(this->schwarz->get_parameters().coarse_solver_factory, nullptr);
 }
 

examples/distributed-solver/distributed-solver.cpp

@@ -53,6 +53,7 @@ int main(int argc, char* argv[])
     using schwarz = gko::experimental::distributed::preconditioner::Schwarz<
         ValueType, LocalIndexType, GlobalIndexType>;
     using bj = gko::preconditioner::Jacobi<ValueType, LocalIndexType>;
+    using pgm = gko::multigrid::Pgm<ValueType, LocalIndexType>;
 
     // Create an MPI communicator get the rank of the calling process.
     const auto comm = gko::experimental::mpi::communicator(MPI_COMM_WORLD);
@@ -188,20 +189,36 @@ int main(int argc, char* argv[])
 
     // Setup the local block diagonal solver factory.
     auto local_solver = gko::share(bj::build().on(exec));
+    auto coarse_solver = gko::share(
+        solver::build()
+            .with_criteria(
+                gko::stop::Iteration::build().with_max_iters(100).on(exec),
+                gko::stop::ResidualNorm<ValueType>::build()
+                    .with_reduction_factor(1e-6)
+                    .on(exec))
+            .on(exec));
+    auto pgm_fac = gko::share(pgm::build().on(exec));
 
     // Setup the stopping criterion and logger
     const gko::remove_complex<ValueType> reduction_factor{1e-8};
     std::shared_ptr<const gko::log::Convergence<ValueType>> logger =
         gko::log::Convergence<ValueType>::create();
-    auto Ainv = solver::build()
-                    .with_preconditioner(
-                        schwarz::build().with_local_solver(local_solver))
-                    .with_criteria(
-                        gko::stop::Iteration::build().with_max_iters(num_iters),
-                        gko::stop::ResidualNorm<ValueType>::build()
-                            .with_reduction_factor(reduction_factor))
-                    .on(exec)
-                    ->generate(A);
+    auto Ainv =
+        solver::build()
+            .with_preconditioner(
+                schwarz::build()
+                    .with_local_solver_factory(local_solver)
+                    // .with_galerkin_ops_factory(pgm_fac)
+                    // .with_coarse_solver_factory(coarse_solver)
+                    .on(exec))
+            .with_criteria(
+                gko::stop::Iteration::build().with_max_iters(num_iters).on(
+                    exec),
+                gko::stop::ResidualNorm<ValueType>::build()
+                    .with_reduction_factor(reduction_factor)
+                    .on(exec))
+            .on(exec)
+            ->generate(A);
     // Add logger to the generated solver to log the iteration count and
     // residual norm
     Ainv->add_logger(logger);
@@ -219,16 +236,16 @@ int main(int argc, char* argv[])
     ValueType t_end = gko::experimental::mpi::get_walltime();
 
     // Get the residual.
-    auto res_norm = gko::clone(exec->get_master(),
-                               gko::as<vec>(logger->get_residual_norm()));
+    auto res_norm = gko::as<vec>(logger->get_residual_norm());
+    auto host_res = gko::make_temporary_clone(exec->get_master(), res_norm);
 
     // @sect3{Printing Results}
     // Print the achieved residual norm and timings on rank 0.
     if (comm.rank() == 0) {
         // clang-format off
         std::cout << "\nNum rows in matrix: " << num_rows
                   << "\nNum ranks: " << comm.size()
-                  << "\nFinal Res norm: " << res_norm->at(0, 0)
+                  << "\nFinal Res norm: " << *host_res->get_const_values()
                   << "\nIteration count: " << logger->get_num_iterations()
                   << "\nInit time: " << t_init_end - t_init
                   << "\nRead time: " << t_read_setup_end - t_init

include/ginkgo/core/distributed/preconditioner/schwarz.hpp

@@ -76,6 +76,13 @@ class Schwarz
         std::shared_ptr<const LinOp> GKO_FACTORY_PARAMETER_SCALAR(
             generated_local_solver, nullptr);
 
+        /**
+         * Operator factory to generate the triplet (prolong_op, coarse_op,
+         * restrict_op).
+         */
+        std::shared_ptr<const LinOpFactory> GKO_FACTORY_PARAMETER_SCALAR(
+            galerkin_ops_factory, nullptr);
+
         /**
          * Coarse solver factory.
          */
@@ -134,7 +141,8 @@ class Schwarz
     void set_solver(std::shared_ptr<const LinOp> new_solver);
 
     std::shared_ptr<const LinOp> local_solver_;
-    std::shared_ptr<const multigrid::MultigridLevel> coarse_solver_;
+    std::shared_ptr<const multigrid::MultigridLevel> galerkin_ops_;
+    std::shared_ptr<const LinOp> coarse_solver_;
 };