ginkgo-project · greole · Oct 24, 2023 · Oct 13, 2023 · Oct 13, 2023 · Oct 16, 2023
diff --git a/core/distributed/preconditioner/schwarz.cpp b/core/distributed/preconditioner/schwarz.cpp
@@ -98,18 +98,42 @@ void Schwarz<ValueType, LocalIndexType, GlobalIndexType>::apply_impl(
 }
 
 
+template <typename ValueType, typename LocalIndexType, typename GlobalIndexType>
+void Schwarz<ValueType, LocalIndexType, GlobalIndexType>::set_solver(
+    std::shared_ptr<const LinOp> new_solver)
+{
+    auto exec = this->get_executor();
+    if (new_solver) {
+        if (new_solver->get_executor() != exec) {
+            new_solver = gko::clone(exec, new_solver);
+        }
+    }
+    this->local_solver_ = new_solver;
+}
+
+
 template <typename ValueType, typename LocalIndexType, typename GlobalIndexType>
 void Schwarz<ValueType, LocalIndexType, GlobalIndexType>::generate(
     std::shared_ptr<const LinOp> system_matrix)
 {
+    if (parameters_.local_solver && parameters_.generated_local_solver) {
+        GKO_INVALID_STATE(
+            "Provided both a generated solver and a solver factory");
+    }
+
+    if (!parameters_.local_solver && !parameters_.generated_local_solver) {
+        GKO_INVALID_STATE(
+            "Requires either a generated solver or an solver factory");
+    }
+
     if (parameters_.local_solver) {
-        this->local_solver_ = 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(
+            as<experimental::distributed::Matrix<
+                ValueType, LocalIndexType, GlobalIndexType>>(system_matrix)
+                ->get_local_matrix())));
+
     } else {
-        GKO_NOT_IMPLEMENTED;
+        this->set_solver(parameters_.generated_local_solver);
     }
 }
 

diff --git a/include/ginkgo/core/distributed/preconditioner/schwarz.hpp b/include/ginkgo/core/distributed/preconditioner/schwarz.hpp
@@ -95,6 +95,12 @@ class Schwarz
          * Local solver factory.
          */
         GKO_DEFERRED_FACTORY_PARAMETER(local_solver, LinOpFactory);
+
+        /**
+         * Generated Inner solvers.
+         */
+        std::shared_ptr<const LinOp> GKO_FACTORY_PARAMETER_SCALAR(
+            generated_local_solver, nullptr);
     };
     GKO_ENABLE_LIN_OP_FACTORY(Schwarz, parameters, Factory);
     GKO_ENABLE_BUILD_METHOD(Factory);
@@ -130,7 +136,6 @@ class Schwarz
      */
     void generate(std::shared_ptr<const LinOp> system_matrix);
 
-
     void apply_impl(const LinOp* b, LinOp* x) const override;
 
     template <typename VectorType>
@@ -140,6 +145,13 @@ class Schwarz
                     LinOp* x) const override;
 
 private:
+    /**
+     * Sets the solver operator used as the local solver.
+     *
+     * @param new_solver  the new local solver
+     */
+    void set_solver(std::shared_ptr<const LinOp> new_solver);
+
     std::shared_ptr<const LinOp> local_solver_;
 };
 

diff --git a/test/mpi/preconditioner/schwarz.cpp b/test/mpi/preconditioner/schwarz.cpp
@@ -178,6 +178,36 @@ class SchwarzPreconditioner : public CommonMpiTestFixture {
 TYPED_TEST_SUITE(SchwarzPreconditioner, gko::test::ValueLocalGlobalIndexTypes,
                  TupleTypenameNameGenerator);
 
+TYPED_TEST(SchwarzPreconditioner, GenerateFailsIfInvalidState)
+{
+    using value_type = typename TestFixture::value_type;
+    using local_index_type = typename TestFixture::local_index_type;
+    using local_prec_type =
+        gko::preconditioner::Jacobi<value_type, local_index_type>;
+    using prec = typename TestFixture::dist_prec_type;
+
+    auto local_solver = gko::share(local_prec_type::build()
+                                       .with_max_block_size(1u)
+                                       .on(this->exec)
+                                       ->generate(this->non_dist_mat));
+    auto schwarz = prec::build()
+                       .with_local_solver(this->local_solver_factory)
+                       .with_generated_local_solver(local_solver)
+                       .on(this->exec);
+
+    ASSERT_THROW(schwarz->generate(this->dist_mat), gko::InvalidStateError);
+}
+
+
+TYPED_TEST(SchwarzPreconditioner, GenerateFailsIfNoSolverProvided)
+{
+    using prec = typename TestFixture::dist_prec_type;
+    auto schwarz_no_solver = prec::build().on(this->exec);
+
+    ASSERT_THROW(schwarz_no_solver->generate(this->dist_mat),
+                 gko::InvalidStateError);
+}
+
 
 TYPED_TEST(SchwarzPreconditioner, CanApplyPreconditionedSolver)
 {
@@ -218,13 +248,46 @@ TYPED_TEST(SchwarzPreconditioner, CanApplyPreconditionedSolver)
 }
 
 
-TYPED_TEST(SchwarzPreconditioner, CanApplyPreconditioner)
+TYPED_TEST(SchwarzPreconditioner, CanApplyPreconditionedSolverWithPregenSolver)
 {
     using value_type = typename TestFixture::value_type;
+    using local_index_type = typename TestFixture::local_index_type;
+    using local_prec_type =
+        gko::preconditioner::Jacobi<value_type, local_index_type>;
     using csr = typename TestFixture::local_matrix_type;
     using cg = typename TestFixture::solver_type;
     using prec = typename TestFixture::dist_prec_type;
 
+    auto local_solver =
+        gko::share(local_prec_type::build()
+                       .with_max_block_size(1u)
+                       .on(this->exec)
+                       ->generate(this->dist_mat->get_local_matrix()));
+    auto precond = prec::build()
+                       .with_local_solver(this->local_solver_factory)
+                       .on(this->exec)
+                       ->generate(this->dist_mat);
+    auto precond_pregen = prec::build()
+                              .with_generated_local_solver(local_solver)
+                              .on(this->exec)
+                              ->generate(this->dist_mat);
+    auto dist_x = gko::share(this->dist_x->clone());
+    auto dist_x_pregen = gko::share(this->dist_x->clone());
+
+    precond->apply(this->dist_b.get(), dist_x.get());
+    precond_pregen->apply(this->dist_b.get(), dist_x_pregen.get());
+
+    GKO_ASSERT_MTX_NEAR(dist_x->get_local_vector(),
+                        dist_x_pregen->get_local_vector(),
+                        r<value_type>::value);
+}
+
+
+TYPED_TEST(SchwarzPreconditioner, CanApplyPreconditioner)
+{
+    using value_type = typename TestFixture::value_type;
+    using prec = typename TestFixture::dist_prec_type;
+
     auto precond_factory = prec::build()
                                .with_local_solver(this->local_solver_factory)
                                .on(this->exec);