[prec] implement SOR preconditioner reference kernels

only reference kernels are available

common/unified/CMakeLists.txt

@@ -19,6 +19,7 @@ set(UNIFIED_SOURCES
     matrix/diagonal_kernels.cpp
     multigrid/pgm_kernels.cpp
     preconditioner/jacobi_kernels.cpp
+    preconditioner/sor_kernels.cpp
     solver/bicg_kernels.cpp
     solver/bicgstab_kernels.cpp
     solver/cg_kernels.cpp

common/unified/preconditioner/sor_kernels.cpp

@@ -0,0 +1,44 @@
+// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include "core/preconditioner/sor_kernels.hpp"
+
+#include <ginkgo/core/base/math.hpp>
+#include <ginkgo/core/matrix/csr.hpp>
+
+#include "common/unified/base/kernel_launch.hpp"
+
+
+namespace gko {
+namespace kernels {
+namespace GKO_DEVICE_NAMESPACE {
+namespace sor {
+
+
+template <typename ValueType, typename IndexType>
+void initialize_weighted_l(
+    std::shared_ptr<const DefaultExecutor> exec,
+    const matrix::Csr<ValueType, IndexType>* system_matrix,
+    remove_complex<ValueType> weight,
+    matrix::Csr<ValueType, IndexType>* l_mtx) GKO_NOT_IMPLEMENTED;
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+    GKO_DECLARE_SOR_INITIALIZE_WEIGHTED_L);
+
+
+template <typename ValueType, typename IndexType>
+void initialize_weighted_l_u(
+    std::shared_ptr<const DefaultExecutor> exec,
+    const matrix::Csr<ValueType, IndexType>* system_matrix,
+    remove_complex<ValueType> weight, matrix::Csr<ValueType, IndexType>* l_mtx,
+    matrix::Csr<ValueType, IndexType>* u_mtx) GKO_NOT_IMPLEMENTED;
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(
+    GKO_DECLARE_SOR_INITIALIZE_WEIGHTED_L_U);
+
+
+}  // namespace sor
+}  // namespace GKO_DEVICE_NAMESPACE
+}  // namespace kernels
+}  // namespace gko

core/CMakeLists.txt

@@ -84,6 +84,7 @@ target_sources(${ginkgo_core}
     multigrid/pgm.cpp
     multigrid/fixed_coarsening.cpp
     preconditioner/batch_jacobi.cpp
+        preconditioner/sor.cpp
     preconditioner/ic.cpp
     preconditioner/ilu.cpp
     preconditioner/isai.cpp

core/device_hooks/common_kernels.inc.cpp

@@ -48,6 +48,7 @@
 #include "core/preconditioner/batch_jacobi_kernels.hpp"
 #include "core/preconditioner/isai_kernels.hpp"
 #include "core/preconditioner/jacobi_kernels.hpp"
+#include "core/preconditioner/sor_kernels.hpp"
 #include "core/reorder/rcm_kernels.hpp"
 #include "core/solver/batch_bicgstab_kernels.hpp"
 #include "core/solver/batch_cg_kernels.hpp"
@@ -818,6 +819,16 @@ GKO_STUB(GKO_DECLARE_JACOBI_INITIALIZE_PRECISIONS_KERNEL);
 }  // namespace jacobi
 
 
+namespace sor {
+
+
+GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SOR_INITIALIZE_WEIGHTED_L);
+GKO_STUB_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SOR_INITIALIZE_WEIGHTED_L_U);
+
+
+}  // namespace sor
+
+
 namespace isai {
 
 

core/preconditioner/sor.cpp

@@ -0,0 +1,133 @@
+// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include <ginkgo/core/base/array.hpp>
+#include <ginkgo/core/base/executor.hpp>
+#include <ginkgo/core/base/precision_dispatch.hpp>
+#include <ginkgo/core/matrix/csr.hpp>
+#include <ginkgo/core/matrix/diagonal.hpp>
+#include <ginkgo/core/preconditioner/sor.hpp>
+#include <ginkgo/core/solver/triangular.hpp>
+
+#include "core/base/array_access.hpp"
+#include "core/base/utils.hpp"
+#include "core/factorization/factorization_kernels.hpp"
+#include "core/matrix/csr_builder.hpp"
+#include "core/preconditioner/sor_kernels.hpp"
+
+namespace gko {
+namespace preconditioner {
+namespace {
+
+
+GKO_REGISTER_OPERATION(initialize_row_ptrs_l,
+                       factorization::initialize_row_ptrs_l);
+GKO_REGISTER_OPERATION(initialize_row_ptrs_l_u,
+                       factorization::initialize_row_ptrs_l_u);
+GKO_REGISTER_OPERATION(initialize_weighted_l, sor::initialize_weighted_l);
+GKO_REGISTER_OPERATION(initialize_weighted_l_u, sor::initialize_weighted_l_u);
+
+
+}  // namespace
+
+
+template <typename ValueType, typename IndexType>
+std::unique_ptr<typename Sor<ValueType, IndexType>::composition_type>
+Sor<ValueType, IndexType>::generate(
+    std::shared_ptr<const LinOp> system_matrix) const
+{
+    auto product =
+        std::unique_ptr<composition_type>(static_cast<composition_type*>(
+            this->LinOpFactory::generate(std::move(system_matrix)).release()));
+    return product;
+}
+
+
+template <typename ValueType, typename IndexType>
+std::unique_ptr<LinOp> Sor<ValueType, IndexType>::generate_impl(
+    std::shared_ptr<const LinOp> system_matrix) const
+{
+    using Csr = matrix::Csr<ValueType, IndexType>;
+    using LTrs = solver::LowerTrs<value_type, index_type>;
+    using UTrs = solver::UpperTrs<value_type, index_type>;
+
+    GKO_ASSERT_IS_SQUARE_MATRIX(system_matrix);
+
+    auto exec = this->get_executor();
+    auto size = system_matrix->get_size();
+
+    auto csr_matrix = convert_to_with_sorting<Csr>(exec, system_matrix,
+                                                   parameters_.skip_sorting);
+
+    auto l_trs_factory =
+        parameters_.l_solver ? parameters_.l_solver : LTrs::build().on(exec);
+    auto u_trs_factory =
+        parameters_.u_solver ? parameters_.u_solver : UTrs::build().on(exec);
+
+    if (parameters_.symmetric) {
+        array<index_type> l_row_ptrs{exec, size[0] + 1};
+        array<index_type> u_row_ptrs{exec, size[0] + 1};
+        exec->run(make_initialize_row_ptrs_l_u(
+            csr_matrix.get(), l_row_ptrs.get_data(), u_row_ptrs.get_data()));
+        const auto l_nnz =
+            static_cast<size_type>(get_element(l_row_ptrs, size[0]));
+        const auto u_nnz =
+            static_cast<size_type>(get_element(u_row_ptrs, size[0]));
+
+        // create matrices
+        auto l_mtx =
+            Csr::create(exec, size, array<value_type>{exec, l_nnz},
+                        array<index_type>{exec, l_nnz}, std::move(l_row_ptrs));
+        auto u_mtx =
+            Csr::create(exec, size, array<value_type>{exec, u_nnz},
+                        array<index_type>{exec, u_nnz}, std::move(u_row_ptrs));
+
+        // fill l_mtx with 1/w (D + wL)
+        // fill u_mtx with 1/(1-w) (D + wU)
+        exec->run(make_initialize_weighted_l_u(csr_matrix.get(),
+                                               parameters_.relaxation_factor,
+                                               l_mtx.get(), u_mtx.get()));
+
+        // scale u_mtx with D^-1
+        auto diag = csr_matrix->extract_diagonal();
+        diag->inverse_apply(u_mtx, u_mtx);
+
+        // invert the triangular matrices with triangular solvers
+        auto l_trs = l_trs_factory->generate(std::move(l_mtx));
+        auto u_trs = u_trs_factory->generate(std::move(u_mtx));
+
+        // return (1/(w * (1 - w)) (D + wL) D^-1 (D + wU))^-1
+        // because of the inversion, the factor order is switched
+        return composition_type::create(std::move(u_trs), std::move(l_trs));
+    } else {
+        array<index_type> l_row_ptrs{exec, size[0] + 1};
+        exec->run(make_initialize_row_ptrs_l(csr_matrix.get(),
+                                             l_row_ptrs.get_data()));
+        const auto l_nnz =
+            static_cast<size_type>(get_element(l_row_ptrs, size[0]));
+
+        // create matrices
+        auto l_mtx =
+            Csr::create(exec, size, array<value_type>{exec, l_nnz},
+                        array<index_type>{exec, l_nnz}, std::move(l_row_ptrs));
+
+        // fill l_mtx with 1/w * (D + wL)
+        exec->run(make_initialize_weighted_l(
+            csr_matrix.get(), parameters_.relaxation_factor, l_mtx.get()));
+
+        // invert the triangular matrices with triangular solvers
+        auto l_trs = l_trs_factory->generate(std::move(l_mtx));
+
+        return composition_type::create(std::move(l_trs));
+    }
+}
+
+
+#define GKO_DECLARE_SOR(ValueType, IndexType) class Sor<ValueType, IndexType>
+
+GKO_INSTANTIATE_FOR_EACH_VALUE_AND_INDEX_TYPE(GKO_DECLARE_SOR);
+
+
+}  // namespace preconditioner
+}  // namespace gko

core/preconditioner/sor_kernels.hpp

@@ -0,0 +1,50 @@
+// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#ifndef GKO_CORE_PRECONDITIONER_SOR_KERNELS_HPP_
+#define GKO_CORE_PRECONDITIONER_SOR_KERNELS_HPP_
+
+
+#include <ginkgo/core/matrix/csr.hpp>
+#include <ginkgo/core/preconditioner/sor.hpp>
+
+#include "core/base/kernel_declaration.hpp"
+
+
+namespace gko {
+namespace kernels {
+
+
+#define GKO_DECLARE_SOR_INITIALIZE_WEIGHTED_L(_vtype, _itype) \
+    void initialize_weighted_l(                               \
+        std::shared_ptr<const DefaultExecutor> exec,          \
+        const matrix::Csr<_vtype, _itype>* system_matrix,     \
+        remove_complex<_vtype> weight, matrix::Csr<_vtype, _itype>* l_factor)
+
+
+#define GKO_DECLARE_SOR_INITIALIZE_WEIGHTED_L_U(_vtype, _itype)               \
+    void initialize_weighted_l_u(                                             \
+        std::shared_ptr<const DefaultExecutor> exec,                          \
+        const matrix::Csr<_vtype, _itype>* system_matrix,                     \
+        remove_complex<_vtype> weight, matrix::Csr<_vtype, _itype>* l_factor, \
+        matrix::Csr<_vtype, _itype>* u_factor)
+
+
+#define GKO_DECLARE_ALL_AS_TEMPLATES                             \
+    template <typename ValueType, typename IndexType>            \
+    GKO_DECLARE_SOR_INITIALIZE_WEIGHTED_L(ValueType, IndexType); \
+    template <typename ValueType, typename IndexType>            \
+    GKO_DECLARE_SOR_INITIALIZE_WEIGHTED_L_U(ValueType, IndexType)
+
+
+GKO_DECLARE_FOR_ALL_EXECUTOR_NAMESPACES(sor, GKO_DECLARE_ALL_AS_TEMPLATES);
+
+
+#undef GKO_DECLARE_ALL_AS_TEMPLATES
+
+
+}  // namespace kernels
+}  // namespace gko
+
+#endif  // GKO_CORE_PRECONDITIONER_SOR_KERNELS_HPP_

include/ginkgo/core/preconditioner/sor.hpp

@@ -0,0 +1,125 @@
+// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
+//
+// SPDX-License-Identifier: BSD-3-Clause
+
+#ifndef GKO_PUBLIC_CORE_PRECONDITIONER_SOR_HPP_
+#define GKO_PUBLIC_CORE_PRECONDITIONER_SOR_HPP_
+
+
+#include <vector>
+
+#include <ginkgo/core/base/abstract_factory.hpp>
+#include <ginkgo/core/base/composition.hpp>
+#include <ginkgo/core/base/lin_op.hpp>
+#include <ginkgo/core/base/polymorphic_object.hpp>
+
+
+namespace gko {
+namespace preconditioner {
+
+
+/**
+ * This class generates the (S)SOR preconditioner.
+ *
+ * The SOR preconditioner starts from a splitting of the the matrix $A$ into
+ * $A = D + L + U$, where $L$ contains all entries below the diagonal, and $U$
+ * contains all entries above the diagonal. The application of the
+ * preconditioner is then defined as solving $M x = y$ with
+ * $$
+ * M = \frac{1}{\omega} (D + \omega L), \quad 0 < \omega < 2.
+ * $$
+ * $\omega$ is known as the relaxation factor.
+ * The preconditioner can be made symmetric, leading to the SSOR preconitioner.
+ * Here, $M$ is defined as
+ * $$
+ * M = \frac{1}{\omega (2 - \omega)} (D + \omega L) D^{-1} (D + \omega U) ,
+ * \quad 0 < \omega < 2.
+ * $$
+ *
+ * This class is a factory, which will only generate the preconditioner. The
+ * resulting LinOp will represent the application of $M^{-1}$.
+ *
+ * @tparam ValueType  The value type of the internally used CSR matrix
+ * @tparam IndexType  The index type of the internally used CSR matrix
+ */
+template <typename ValueType = default_precision, typename IndexType = int32>
+class Sor
+    : public EnablePolymorphicObject<Sor<ValueType, IndexType>, LinOpFactory>,
+      public EnablePolymorphicAssignment<Sor<ValueType, IndexType>> {
+    friend class EnablePolymorphicObject<Sor, LinOpFactory>;
+
+public:
+    struct parameters_type;
+    friend class enable_parameters_type<parameters_type, Sor>;
+
+    using value_type = ValueType;
+    using index_type = IndexType;
+    using composition_type = Composition<ValueType>;
+
+    struct parameters_type
+        : public enable_parameters_type<parameters_type, Sor> {
+        // skip sorting of input matrix
+        bool GKO_FACTORY_PARAMETER_SCALAR(skip_sorting, true);
+
+        // determines if SOR or SSOR should be used
+        bool GKO_FACTORY_PARAMETER_SCALAR(symmetric, false);
+
+        // has to be between 0.0 and 2.0
+        remove_complex<value_type> GKO_FACTORY_PARAMETER_SCALAR(
+            relaxation_factor, remove_complex<value_type>(1.2));
+
+        // factory for the lower triangular factor solver
+        std::shared_ptr<const LinOpFactory> GKO_DEFERRED_FACTORY_PARAMETER(
+            l_solver);
+
+        // factory for the upper triangular factor solver, unused if symmetric
+        // is false
+        std::shared_ptr<const LinOpFactory> GKO_DEFERRED_FACTORY_PARAMETER(
+            u_solver);
+    };
+
+    /**
+     * Returns the parameters used to construct the factory.
+     *
+     * @return the parameters used to construct the factory.
+     */
+    const parameters_type& get_parameters() { return parameters_; }
+
+    /**
+     * @copydoc get_parameters
+     */
+    const parameters_type& get_parameters() const { return parameters_; }
+
+    /**
+     * @copydoc LinOpFactory::generate
+     * @note This function overrides the default LinOpFactory::generate to
+     *       return a Factorization instead of a generic LinOp, which would need
+     *       to be cast to Factorization again to access its factors.
+     *       It is only necessary because smart pointers aren't covariant.
+     */
+    std::unique_ptr<composition_type> generate(
+        std::shared_ptr<const LinOp> system_matrix) const;
+
+    /** Creates a new parameter_type to set up the factory. */
+    static parameters_type build() { return {}; }
+
+protected:
+    explicit Sor(std::shared_ptr<const Executor> exec,
+                 const parameters_type& params = {})
+        : EnablePolymorphicObject<Sor, LinOpFactory>(exec), parameters_(params)
+    {
+        GKO_ASSERT(parameters_.relaxation_factor > 0.0 &&
+                   parameters_.relaxation_factor < 2.0);
+    }
+
+    std::unique_ptr<LinOp> generate_impl(
+        std::shared_ptr<const LinOp> system_matrix) const override;
+
+private:
+    parameters_type parameters_;
+};
+}  // namespace preconditioner
+}  // namespace gko
+
+
+#endif  // GKO_PUBLIC_CORE_PRECONDITIONER_SOR_HPP_