From e988e0e7e0c692d0684c8d0191746c1210d62222 Mon Sep 17 00:00:00 2001
From: Pratik Nayak <pratikvn@protonmail.com>
Date: Fri, 31 Jan 2025 13:55:07 +0100
Subject: [PATCH] [precond] move csol to a VectorCache

---
 core/distributed/preconditioner/schwarz.cpp   | 26 +++++++++----------
 .../distributed/preconditioner/schwarz.hpp    |  4 +--
 2 files changed, 15 insertions(+), 15 deletions(-)

diff --git a/core/distributed/preconditioner/schwarz.cpp b/core/distributed/preconditioner/schwarz.cpp
index 77f13784e1a..230a8a5494a 100644
--- a/core/distributed/preconditioner/schwarz.cpp
+++ b/core/distributed/preconditioner/schwarz.cpp
@@ -102,9 +102,19 @@ void Schwarz<ValueType, LocalIndexType, GlobalIndexType>::apply_dense_impl(
                            ->get_prolong_op();
         GKO_ASSERT(this->half_ != nullptr);
 
-        restrict->apply(dense_b, this->csol_);
-        this->coarse_solver_->apply(this->csol_, this->csol_);
-        prolong->apply(this->half_, this->csol_, this->half_, dense_x);
+        // Coarse solve vector cache init
+        // Should allocare only in the first apply call.
+        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 comm = coarse->get_communicator();
+        csol_cache_.init(exec, comm, cs_global_size, cs_local_size);
+
+        restrict->apply(dense_b, csol_cache_);
+        this->coarse_solver_->apply(csol_cache_, csol_cache_);
+        prolong->apply(this->half_, csol_cache_, this->half_, dense_x);
     }
 }
 
@@ -181,16 +191,6 @@ void Schwarz<ValueType, LocalIndexType, GlobalIndexType>::generate(
             auto comm = coarse->get_communicator();
             this->coarse_solver_ =
                 share(parameters_.coarse_solver->generate(coarse));
-            // TODO: Set correct rhs and stride.
-            auto cs_ncols = 1;  // 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);
-            this->csol_ = gko::share(
-                dist_vec::create(exec, comm, cs_global_size, cs_local_size,
-                                 1 /*dense_x->get_stride()*/));
-            // this->temp_ = this->csol->clone();
             this->half_ = gko::share(gko::initialize<Vector>({0.5}, exec));
         }
     }
diff --git a/include/ginkgo/core/distributed/preconditioner/schwarz.hpp b/include/ginkgo/core/distributed/preconditioner/schwarz.hpp
index c190684cb32..e9a52857c16 100644
--- a/include/ginkgo/core/distributed/preconditioner/schwarz.hpp
+++ b/include/ginkgo/core/distributed/preconditioner/schwarz.hpp
@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -179,10 +179,10 @@ class Schwarz
     std::shared_ptr<const LinOp> local_solver_;
 
     detail::VectorCache<ValueType> cache_;
+    detail::VectorCache<ValueType> csol_cache_;
 
     std::shared_ptr<const LinOp> galerkin_ops_;
     std::shared_ptr<const LinOp> coarse_solver_;
-    std::shared_ptr<LinOp> csol_;
     std::shared_ptr<const LinOp> half_;
 };