Merge Add distributed L1 smoother

MarcelKoch · web-flow · commit e66d6d1d28f7 · 2025-03-19T15:31:26.000+01:00
This PR adds a distributed version of #1310. It only uses the non-local matrix to create the additional diagonal matrix which is added to the local matrix. Related PR: #1806
diff --git a/common/unified/matrix/csr_kernels.cpp b/common/unified/matrix/csr_kernels.cpp
@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -313,6 +313,27 @@ void benchmark_lookup(std::shared_ptr<const DefaultExecutor> exec,
 GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_CSR_BENCHMARK_LOOKUP_KERNEL);
 
 
+template <typename ValueType, typename IndexType>
+void row_wise_absolute_sum(std::shared_ptr<const DefaultExecutor> exec,
+                           const matrix::Csr<ValueType, IndexType>* orig,
+                           array<ValueType>& sum)
+{
+    run_kernel(
+        exec,
+        [] GKO_KERNEL(auto row, auto row_ptrs, auto value_ptr, auto sum_ptr) {
+            sum_ptr[row] = zero<device_type<ValueType>>();
+            for (auto k = row_ptrs[row]; k < row_ptrs[row + 1]; ++k) {
+                sum_ptr[row] += abs(value_ptr[k]);
+            }
+        },
+        sum.get_num_elems(), orig->get_const_row_ptrs(),
+        orig->get_const_values(), sum.get_data());
+}
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+    GKO_DECLARE_CSR_ROW_WISE_ABSOLUTE_SUM);
+
+
 }  // namespace csr
 }  // namespace GKO_DEVICE_NAMESPACE
 }  // namespace kernels
diff --git a/core/device_hooks/common_kernels.inc.cpp b/core/device_hooks/common_kernels.inc.cpp
@@ -767,6 +767,7 @@ GKO_STUB_VALUE_AND_INDEX_TYPE(
 GKO_STUB_INDEX_TYPE(GKO_DECLARE_CSR_BUILD_LOOKUP_OFFSETS_KERNEL);
 GKO_STUB_INDEX_TYPE(GKO_DECLARE_CSR_BUILD_LOOKUP_KERNEL);
 GKO_STUB_INDEX_TYPE(GKO_DECLARE_CSR_BENCHMARK_LOOKUP_KERNEL);
+GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_CSR_ROW_WISE_ABSOLUTE_SUM);
 
 template <typename ValueType, typename IndexType>
 GKO_DECLARE_CSR_SCALE_KERNEL(ValueType, IndexType)
diff --git a/core/distributed/preconditioner/schwarz.cpp b/core/distributed/preconditioner/schwarz.cpp
@@ -17,17 +17,26 @@
 #include <ginkgo/core/distributed/matrix.hpp>
 #include <ginkgo/core/matrix/csr.hpp>
 #include <ginkgo/core/matrix/dense.hpp>
+#include <ginkgo/core/matrix/identity.hpp>
 
 #include "core/base/utils.hpp"
 #include "core/config/config_helper.hpp"
 #include "core/config/dispatch.hpp"
 #include "core/distributed/helpers.hpp"
+#include "core/matrix/csr_kernels.hpp"
 
 
 namespace gko {
 namespace experimental {
 namespace distributed {
 namespace preconditioner {
+namespace {
+
+
+GKO_REGISTER_OPERATION(row_wise_absolute_sum, csr::row_wise_absolute_sum);
+
+
+}
 
 
 template <typename ValueType, typename LocalIndexType, typename GlobalIndexType>
@@ -36,7 +45,7 @@ Schwarz<ValueType, LocalIndexType, GlobalIndexType>::parse(
     const config::pnode& config, const config::registry& context,
     const config::type_descriptor& td_for_child)
 {
-    auto params = Schwarz<ValueType, LocalIndexType, GlobalIndexType>::build();
+    auto params = Schwarz::build();
 
     if (auto& obj = config.get("generated_local_solver")) {
         params.with_generated_local_solver(
@@ -47,7 +56,9 @@ Schwarz<ValueType, LocalIndexType, GlobalIndexType>::parse(
             gko::config::parse_or_get_factory<const LinOpFactory>(
                 obj, context, td_for_child));
     }
-
+    if (auto& obj = config.get("l1_smoother")) {
+        params.with_l1_smoother(obj.get_boolean());
+    }
     return params;
 }
 
@@ -76,7 +87,6 @@ template <typename VectorType>
 void Schwarz<ValueType, LocalIndexType, GlobalIndexType>::apply_dense_impl(
     const VectorType* dense_b, VectorType* dense_x) const
 {
-    using Vector = matrix::Dense<ValueType>;
     auto exec = this->get_executor();
     if (this->local_solver_ != nullptr) {
         this->local_solver_->apply(gko::detail::get_local(dense_b),
@@ -130,14 +140,48 @@ void Schwarz<ValueType, LocalIndexType, GlobalIndexType>::generate(
             "Requires either a generated solver or an solver factory");
     }
 
-    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())));
+    if (parameters_.generated_local_solver) {
+        this->set_solver(parameters_.generated_local_solver);
+        return;
+    }
 
+    auto local_matrix =
+        as<Matrix<ValueType, LocalIndexType, GlobalIndexType>>(system_matrix)
+            ->get_local_matrix();
+
+    if (parameters_.l1_smoother) {
+        auto exec = this->get_executor();
+
+        using Csr = matrix::Csr<ValueType, LocalIndexType>;
+        auto local_matrix_copy = share(Csr::create(exec));
+        as<ConvertibleTo<Csr>>(local_matrix)->convert_to(local_matrix_copy);
+
+        auto non_local_matrix = copy_and_convert_to<Csr>(
+            exec, as<Matrix<ValueType, LocalIndexType, GlobalIndexType>>(
+                      system_matrix)
+                      ->get_non_local_matrix());
+
+        array<ValueType> l1_diag_arr{exec, local_matrix->get_size()[0]};
+
+        exec->run(
+            make_row_wise_absolute_sum(non_local_matrix.get(), l1_diag_arr));
+
+        // compute local_matrix_copy <- diag(l1) + local_matrix_copy
+        auto l1_diag = matrix::Diagonal<ValueType>::create(
+            exec, local_matrix->get_size()[0], std::move(l1_diag_arr));
+        auto l1_diag_csr = Csr::create(exec);
+        l1_diag->move_to(l1_diag_csr);
+        auto id = matrix::Identity<ValueType>::create(
+            exec, local_matrix->get_size()[0]);
+        auto one = initialize<matrix::Dense<ValueType>>(
+            {::gko::one<ValueType>()}, exec);
+        l1_diag_csr->apply(one, id, one, local_matrix_copy);
+
+        this->set_solver(
+            gko::share(parameters_.local_solver->generate(local_matrix_copy)));
     } else {
-        this->set_solver(parameters_.generated_local_solver);
+        this->set_solver(
+            gko::share(parameters_.local_solver->generate(local_matrix)));
     }
 }
 
diff --git a/core/matrix/csr_kernels.hpp b/core/matrix/csr_kernels.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
 
@@ -258,6 +258,12 @@ namespace kernels {
                           IndexType sample_size, IndexType* result)
 
 
+#define GKO_DECLARE_CSR_ROW_WISE_ABSOLUTE_SUM(ValueType, IndexType)           \
+    void row_wise_absolute_sum(std::shared_ptr<const DefaultExecutor> exec,   \
+                               const matrix::Csr<ValueType, IndexType>* orig, \
+                               array<ValueType>& sum)
+
+
 #define GKO_DECLARE_ALL_AS_TEMPLATES                                        \
     template <typename MatrixValueType, typename InputValueType,            \
               typename OutputValueType, typename IndexType>                 \
@@ -336,7 +342,9 @@ namespace kernels {
     template <typename IndexType>                                           \
     GKO_DECLARE_CSR_BUILD_LOOKUP_KERNEL(IndexType);                         \
     template <typename IndexType>                                           \
-    GKO_DECLARE_CSR_BENCHMARK_LOOKUP_KERNEL(IndexType)
+    GKO_DECLARE_CSR_BENCHMARK_LOOKUP_KERNEL(IndexType);                     \
+    template <typename ValueType, typename IndexType>                       \
+    GKO_DECLARE_CSR_ROW_WISE_ABSOLUTE_SUM(ValueType, IndexType)
 
 
 GKO_DECLARE_FOR_ALL_EXECUTOR_NAMESPACES(csr, GKO_DECLARE_ALL_AS_TEMPLATES);
diff --git a/core/test/config/preconditioner.cpp b/core/test/config/preconditioner.cpp
@@ -437,6 +437,8 @@ struct Schwarz
         config_map["generated_local_solver"] = pnode{"linop"};
         param.with_generated_local_solver(
             detail::registry_accessor::get_data<gko::LinOp>(reg, "linop"));
+        config_map["l1_smoother"] = pnode{true};
+        param.with_l1_smoother(true);
     }
 
     template <bool from_reg, typename AnswerType>
diff --git 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
 
@@ -89,6 +89,16 @@ class Schwarz
          */
         std::shared_ptr<const LinOp> GKO_FACTORY_PARAMETER_SCALAR(
             generated_local_solver, nullptr);
+
+        /**
+         * Enable l1 smoother.
+         *
+         * This creates a diagonal matrix from the row-wise absolute
+         * sum of the non-local matrix entries. The diagonal matrix
+         * is then added to the system matrix when generating the
+         * local solver.
+         */
+        bool GKO_FACTORY_PARAMETER_SCALAR(l1_smoother, false);
     };
     GKO_ENABLE_LIN_OP_FACTORY(Schwarz, parameters, Factory);
     GKO_ENABLE_BUILD_METHOD(Factory);
diff --git a/reference/matrix/csr_kernels.cpp b/reference/matrix/csr_kernels.cpp
@@ -1458,6 +1458,28 @@ void benchmark_lookup(std::shared_ptr<const DefaultExecutor> exec,
 GKO_INSTANTIATE_FOR_EACH_INDEX_TYPE(GKO_DECLARE_CSR_BENCHMARK_LOOKUP_KERNEL);
 
 
+template <typename ValueType, typename IndexType>
+void row_wise_absolute_sum(std::shared_ptr<const DefaultExecutor> exec,
+                           const matrix::Csr<ValueType, IndexType>* orig,
+                           array<ValueType>& sum)
+{
+    auto row_ptrs = orig->get_const_row_ptrs();
+    auto value_ptr = orig->get_const_values();
+    auto sum_ptr = sum.get_data();
+
+    for (size_type row = 0; row < orig->get_size()[0]; ++row) {
+        sum_ptr[row] = zero<ValueType>();
+        for (size_type k = row_ptrs[row];
+             k < static_cast<size_type>(row_ptrs[row + 1]); ++k) {
+            sum_ptr[row] += abs(value_ptr[k]);
+        }
+    }
+}
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+    GKO_DECLARE_CSR_ROW_WISE_ABSOLUTE_SUM);
+
+
 }  // namespace csr
 }  // namespace reference
 }  // namespace kernels
diff --git a/reference/test/matrix/csr_kernels.cpp b/reference/test/matrix/csr_kernels.cpp
@@ -2558,6 +2558,21 @@ TYPED_TEST(Csr, CanGetSubmatrixWithIndexSet)
 }
 
 
+TYPED_TEST(Csr, CanComputeRowWiseAbsoluteSum)
+{
+    using value_type = typename TestFixture::value_type;
+    gko::array<value_type> sum(this->exec, this->mtx3_sorted->get_size()[0]);
+    this->create_mtx3(this->mtx3_sorted.get(), this->mtx3_unsorted.get());
+    this->mtx3_sorted->scale(gko::initialize<gko::matrix::Dense<value_type>>(
+        {-gko::one<value_type>()}, this->exec));
+
+    gko::kernels::reference::csr::row_wise_absolute_sum(
+        this->exec, this->mtx3_sorted.get(), sum);
+
+    gko::array<value_type> sum_result(this->exec, {3, 12, 5});
+    GKO_ASSERT_ARRAY_EQ(sum, sum_result);
+}
+
 template <typename ValueIndexType>
 class CsrLookup : public ::testing::Test {
 protected:
diff --git a/test/matrix/csr_kernels.cpp b/test/matrix/csr_kernels.cpp
@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
 //
 // SPDX-License-Identifier: BSD-3-Clause
 
@@ -78,6 +78,20 @@ TEST_F(Csr, InvScaleIsEquivalentToRef)
 }
 
 
+TEST_F(Csr, RowWiseSumIsEquivalentToRef)
+{
+    set_up_apply_data();
+    gko::array<value_type> sum{ref, x->get_size()[0]};
+    gko::array<value_type> dsum{exec, dx->get_size()[0]};
+
+    gko::kernels::reference::csr::row_wise_absolute_sum(ref, x.get(), sum);
+    gko::kernels::GKO_DEVICE_NAMESPACE::csr::row_wise_absolute_sum(
+        exec, dx.get(), dsum);
+
+    GKO_ASSERT_ARRAY_EQ(sum, dsum);
+}
+
+
 template <typename IndexType>
 class CsrLookup : public CommonTestFixture {
 public:
diff --git a/test/mpi/preconditioner/schwarz.cpp b/test/mpi/preconditioner/schwarz.cpp
@@ -32,13 +32,6 @@
 #include "test/utils/mpi/common_fixture.hpp"
 
 
-#if GINKGO_DPCPP_SINGLE_MODE
-using solver_value_type = float;
-#else
-using solver_value_type = double;
-#endif  // GINKGO_DPCPP_SINGLE_MODE
-
-
 template <typename ValueLocalGlobalIndexType>
 class SchwarzPreconditioner : public CommonMpiTestFixture {
 protected:
@@ -296,3 +289,29 @@ TYPED_TEST(SchwarzPreconditioner, CanAdvancedApplyPreconditioner)
     this->assert_equal_to_non_distributed_vector(this->dist_x,
                                                  this->non_dist_x);
 }
+
+
+TYPED_TEST(SchwarzPreconditioner, CanApplyPreconditionerWithL1Smoother)
+{
+    using value_type = typename TestFixture::value_type;
+    using prec = typename TestFixture::dist_prec_type;
+    using local_matrix_type = typename TestFixture::local_matrix_type;
+    auto non_dist_diag_with_l1 =
+        gko::share(gko::matrix::Diagonal<value_type>::create(
+            this->exec, 8u,
+            gko::array<value_type>(this->exec, {2, 3, 3, 3, 3, 2, 2, 2})));
+    auto precond_factory = prec::build()
+                               .with_local_solver(this->local_solver_factory)
+                               .with_l1_smoother(true)
+                               .on(this->exec);
+    auto local_precond = this->local_solver_factory->generate(
+        gko::copy_and_convert_to<local_matrix_type>(this->exec,
+                                                    non_dist_diag_with_l1));
+    auto precond = precond_factory->generate(this->dist_mat);
+
+    precond->apply(this->dist_b.get(), this->dist_x.get());
+    local_precond->apply(this->non_dist_b.get(), this->non_dist_x.get());
+
+    this->assert_equal_to_non_distributed_vector(this->dist_x,
+                                                 this->non_dist_x);
+}

Original file line number	Diff line number	Diff line change
`@@ -437,6 +437,8 @@ struct Schwarz`
`437`	`437`	`config_map["generated_local_solver"] = pnode{"linop"};`
`438`	`438`	`param.with_generated_local_solver(`
`439`	`439`	`detail::registry_accessor::get_data<gko::LinOp>(reg, "linop"));`
	`440`	`+ config_map["l1_smoother"] = pnode{true};`
	`441`	`+ param.with_l1_smoother(true);`
`440`	`442`	`}`
`441`	`443`
`442`	`444`	`template <bool from_reg, typename AnswerType>`