csr.hpp

// SPDX-FileCopyrightText: 2017 - 2026 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_MATRIX_CSR_HPP_
#define GKO_PUBLIC_CORE_MATRIX_CSR_HPP_
 
 
#include <ginkgo/core/base/array.hpp>
#include <ginkgo/core/base/index_set.hpp>
#include <ginkgo/core/base/lin_op.hpp>
#include <ginkgo/core/base/math.hpp>
#include <ginkgo/core/matrix/device_views.hpp>
#include <ginkgo/core/matrix/permutation.hpp>
#include <ginkgo/core/matrix/scaled_permutation.hpp>
 
 
namespace gko {
namespace matrix {
 
 
template <typename ValueType>
class Dense;
 
template <typename ValueType>
class Diagonal;
 
template <typename ValueType, typename IndexType>
class Coo;
 
template <typename ValueType, typename IndexType>
class Ell;
 
template <typename ValueType, typename IndexType>
class Hybrid;
 
template <typename ValueType, typename IndexType>
class Sellp;
 
template <typename ValueType, typename IndexType>
class SparsityCsr;
 
template <typename ValueType, typename IndexType>
class Csr;
 
template <typename ValueType, typename IndexType>
class Fbcsr;
 
template <typename ValueType, typename IndexType>
class CsrBuilder;
 
 
namespace csr {
 
 
enum class spmv_strategy {
    automatic,
    load_balance,
    merge_path,
    classical,
    sparselib
};
 
 
}  // namespace csr
 
 
template <typename ValueType = default_precision, typename IndexType = int32>
class Csr : public LinOp,
            public EnableCloneable<Csr<ValueType, IndexType>>,
            public ConvertibleTo<Csr<next_precision<ValueType>, IndexType>>,
#if GINKGO_ENABLE_HALF || GINKGO_ENABLE_BFLOAT16
            public ConvertibleTo<Csr<next_precision<ValueType, 2>, IndexType>>,
#endif
#if GINKGO_ENABLE_HALF && GINKGO_ENABLE_BFLOAT16
            public ConvertibleTo<Csr<next_precision<ValueType, 3>, IndexType>>,
#endif
            public ConvertibleTo<Dense<ValueType>>,
            public ConvertibleTo<Coo<ValueType, IndexType>>,
            public ConvertibleTo<Ell<ValueType, IndexType>>,
            public ConvertibleTo<Fbcsr<ValueType, IndexType>>,
            public ConvertibleTo<Hybrid<ValueType, IndexType>>,
            public ConvertibleTo<Sellp<ValueType, IndexType>>,
            public ConvertibleTo<SparsityCsr<ValueType, IndexType>>,
            public DiagonalExtractable<ValueType>,
            public ReadableFromMatrixData<ValueType, IndexType>,
            public WritableToMatrixData<ValueType, IndexType>,
            public Transposable,
            public Permutable<IndexType>,
            public EnableAbsoluteComputation<
                remove_complex<Csr<ValueType, IndexType>>>,
            public ScaledIdentityAddable {
    friend class EnableCloneable<Csr>;
    friend class Coo<ValueType, IndexType>;
    friend class Dense<ValueType>;
    friend class Diagonal<ValueType>;
    friend class Ell<ValueType, IndexType>;
    friend class Hybrid<ValueType, IndexType>;
    friend class Sellp<ValueType, IndexType>;
    friend class SparsityCsr<ValueType, IndexType>;
    friend class Fbcsr<ValueType, IndexType>;
    friend class CsrBuilder<ValueType, IndexType>;
    friend class Csr<to_complex<ValueType>, IndexType>;
    GKO_ASSERT_SUPPORTED_VALUE_AND_INDEX_TYPE;
 
public:
    using EnableCloneable<Csr>::convert_to;
    using EnableCloneable<Csr>::move_to;
    using ConvertibleTo<Csr<next_precision<ValueType>, IndexType>>::convert_to;
    using ConvertibleTo<Csr<next_precision<ValueType>, IndexType>>::move_to;
    using ConvertibleTo<Dense<ValueType>>::convert_to;
    using ConvertibleTo<Dense<ValueType>>::move_to;
    using ConvertibleTo<Coo<ValueType, IndexType>>::convert_to;
    using ConvertibleTo<Coo<ValueType, IndexType>>::move_to;
    using ConvertibleTo<Ell<ValueType, IndexType>>::convert_to;
    using ConvertibleTo<Ell<ValueType, IndexType>>::move_to;
    using ConvertibleTo<Fbcsr<ValueType, IndexType>>::convert_to;
    using ConvertibleTo<Fbcsr<ValueType, IndexType>>::move_to;
    using ConvertibleTo<Hybrid<ValueType, IndexType>>::convert_to;
    using ConvertibleTo<Hybrid<ValueType, IndexType>>::move_to;
    using ConvertibleTo<Sellp<ValueType, IndexType>>::convert_to;
    using ConvertibleTo<Sellp<ValueType, IndexType>>::move_to;
    using ConvertibleTo<SparsityCsr<ValueType, IndexType>>::convert_to;
    using ConvertibleTo<SparsityCsr<ValueType, IndexType>>::move_to;
    using ReadableFromMatrixData<ValueType, IndexType>::read;
 
    using value_type = ValueType;
    using index_type = IndexType;
    using transposed_type = Csr<ValueType, IndexType>;
    using mat_data = matrix_data<ValueType, IndexType>;
    using device_mat_data = device_matrix_data<ValueType, IndexType>;
    using absolute_type = remove_complex<Csr>;
    using device_view = view::csr<value_type, index_type>;
    using const_device_view = view::csr<const value_type, const index_type>;
 
    class GKO_DEPRECATED(
        "please use enum gko::matrix::csr::spmv_strategy::<strategy>")
        strategy_type {
    public:
        virtual ~strategy_type() = default;
 
        // return the corresponding enum in incoming release
        virtual csr::spmv_strategy get_enum() const = 0;
    };
 
    class GKO_DEPRECATED(
        "please use enum gko::matrix::csr::spmv_strategy::classical") classical
        : public strategy_type {
    public:
        csr::spmv_strategy get_enum() const override
        {
            return csr::spmv_strategy::classical;
        }
    };
 
    class GKO_DEPRECATED(
        "please use enum gko::matrix::csr::spmv_strategy::merge_path")
        merge_path : public strategy_type {
    public:
        csr::spmv_strategy get_enum() const override
        {
            return csr::spmv_strategy::merge_path;
        }
    };
 
    class GKO_DEPRECATED(
        "please use enum gko::matrix::csr::spmv_strategy::sparselib") cusparse
        : public strategy_type {
    public:
        csr::spmv_strategy get_enum() const override
        {
            return csr::spmv_strategy::sparselib;
        }
    };
 
    class GKO_DEPRECATED(
        "please use enum gko::matrix::csr::spmv_strategy::sparselib") sparselib
        : public strategy_type {
    public:
        csr::spmv_strategy get_enum() const override
        {
            return csr::spmv_strategy::sparselib;
        }
    };
 
    class GKO_DEPRECATED(
        "please use enum gko::matrix::csr::spmv_strategy::load_balance")
        load_balance : public strategy_type {
    public:
        load_balance(std::shared_ptr<const Executor>) {}
 
        csr::spmv_strategy get_enum() const override
        {
            return csr::spmv_strategy::load_balance;
        }
    };
 
    class GKO_DEPRECATED(
        "please use enum gko::matrix::csr::spmv_strategy::automatic")
        automatical : public strategy_type {
    public:
        automatical(std::shared_ptr<const Executor>) {}
 
        csr::spmv_strategy get_enum() const override
        {
            return csr::spmv_strategy::automatic;
        }
    };
 
 
    friend class Csr<previous_precision<ValueType>, IndexType>;
 
    void convert_to(
        Csr<next_precision<ValueType>, IndexType>* result) const override;
 
    void move_to(Csr<next_precision<ValueType>, IndexType>* result) override;
 
#if GINKGO_ENABLE_HALF || GINKGO_ENABLE_BFLOAT16
    friend class Csr<previous_precision<ValueType, 2>, IndexType>;
    using ConvertibleTo<
        Csr<next_precision<ValueType, 2>, IndexType>>::convert_to;
    using ConvertibleTo<Csr<next_precision<ValueType, 2>, IndexType>>::move_to;
 
    void convert_to(
        Csr<next_precision<ValueType, 2>, IndexType>* result) const override;
 
    void move_to(Csr<next_precision<ValueType, 2>, IndexType>* result) override;
#endif
 
#if GINKGO_ENABLE_HALF && GINKGO_ENABLE_BFLOAT16
    friend class Csr<previous_precision<ValueType, 3>, IndexType>;
    using ConvertibleTo<
        Csr<next_precision<ValueType, 3>, IndexType>>::convert_to;
    using ConvertibleTo<Csr<next_precision<ValueType, 3>, IndexType>>::move_to;
 
    void convert_to(
        Csr<next_precision<ValueType, 3>, IndexType>* result) const override;
 
    void move_to(Csr<next_precision<ValueType, 3>, IndexType>* result) override;
#endif
 
    void convert_to(Dense<ValueType>* other) const override;
 
    void move_to(Dense<ValueType>* other) override;
 
    void convert_to(Coo<ValueType, IndexType>* result) const override;
 
    void move_to(Coo<ValueType, IndexType>* result) override;
 
    void convert_to(Ell<ValueType, IndexType>* result) const override;
 
    void move_to(Ell<ValueType, IndexType>* result) override;
 
    void convert_to(Fbcsr<ValueType, IndexType>* result) const override;
 
    void move_to(Fbcsr<ValueType, IndexType>* result) override;
 
    void convert_to(Hybrid<ValueType, IndexType>* result) const override;
 
    void move_to(Hybrid<ValueType, IndexType>* result) override;
 
    void convert_to(Sellp<ValueType, IndexType>* result) const override;
 
    void move_to(Sellp<ValueType, IndexType>* result) override;
 
    void convert_to(SparsityCsr<ValueType, IndexType>* result) const override;
 
    void move_to(SparsityCsr<ValueType, IndexType>* result) override;
 
    void read(const mat_data& data) override;
 
    void read(const device_mat_data& data) override;
 
    void read(device_mat_data&& data) override;
 
    void write(mat_data& data) const override;
 
    void write(device_mat_data& data) const override;
 
    std::unique_ptr<LinOp> transpose() const override;
 
    std::unique_ptr<LinOp> conj_transpose() const override;
 
    device_view get_device_view();
 
    const_device_view get_const_device_view() const;
 
    class multiply_reuse_info {
        friend class Csr;
 
    public:
        explicit multiply_reuse_info();
 
        ~multiply_reuse_info();
 
        multiply_reuse_info(const multiply_reuse_info&) = delete;
 
        multiply_reuse_info(multiply_reuse_info&&) noexcept;
 
        multiply_reuse_info& operator=(const multiply_reuse_info&) = delete;
 
        multiply_reuse_info& operator=(multiply_reuse_info&&) noexcept;
 
        void update_values(ptr_param<const Csr> mtx1, ptr_param<const Csr> mtx2,
                           ptr_param<Csr> out) const;
 
    private:
        struct lookup_data;
 
        explicit multiply_reuse_info(std::unique_ptr<lookup_data> data);
 
        std::unique_ptr<lookup_data> internal;
    };
 
    std::unique_ptr<Csr> multiply(ptr_param<const Csr> other) const;
 
    std::pair<std::unique_ptr<Csr>, multiply_reuse_info> multiply_reuse(
        ptr_param<const Csr> other) const;
 
    class multiply_add_reuse_info {
        friend class Csr;
 
    public:
        explicit multiply_add_reuse_info();
 
        ~multiply_add_reuse_info();
 
        multiply_add_reuse_info(const multiply_add_reuse_info&) = delete;
 
        multiply_add_reuse_info(multiply_add_reuse_info&&) noexcept;
 
        multiply_add_reuse_info& operator=(const multiply_add_reuse_info&) =
            delete;
 
        multiply_add_reuse_info& operator=(multiply_add_reuse_info&&) noexcept;
 
        void update_values(ptr_param<const Csr> mtx,
                           ptr_param<const Dense<value_type>> scale_mult,
                           ptr_param<const Csr> mtx_mult,
                           ptr_param<const Dense<value_type>> scale_add,
                           ptr_param<const Csr> mtx_add,
                           ptr_param<Csr> out) const;
 
    private:
        struct lookup_data;
 
        explicit multiply_add_reuse_info(std::unique_ptr<lookup_data> data);
 
        std::unique_ptr<lookup_data> internal;
    };
 
    std::unique_ptr<Csr> multiply_add(
        ptr_param<const Dense<value_type>> scale_mult,
        ptr_param<const Csr> mtx_mult,
        ptr_param<const Dense<value_type>> scale_add,
        ptr_param<const Csr> mtx_add) const;
 
    std::pair<std::unique_ptr<Csr>, multiply_add_reuse_info> multiply_add_reuse(
        ptr_param<const Dense<value_type>> scale_mult,
        ptr_param<const Csr> mtx_mult,
        ptr_param<const Dense<value_type>> scale_add,
        ptr_param<const Csr> mtx_add) const;
 
    class scale_add_reuse_info {
        friend class Csr;
 
    public:
        explicit scale_add_reuse_info();
 
        ~scale_add_reuse_info();
 
        scale_add_reuse_info(const scale_add_reuse_info&) = delete;
 
        scale_add_reuse_info(scale_add_reuse_info&&) noexcept;
 
        scale_add_reuse_info& operator=(const scale_add_reuse_info&) = delete;
 
        scale_add_reuse_info& operator=(scale_add_reuse_info&&) noexcept;
 
        void update_values(ptr_param<const Dense<value_type>> scale1,
                           ptr_param<const Csr> mtx1,
                           ptr_param<const Dense<value_type>> scale2,
                           ptr_param<const Csr> mtx2, ptr_param<Csr> out) const;
 
    private:
        struct lookup_data;
 
        explicit scale_add_reuse_info(std::unique_ptr<lookup_data> data);
 
        std::unique_ptr<lookup_data> internal;
    };
 
    std::unique_ptr<Csr> scale_add(
        ptr_param<const Dense<value_type>> scale_this,
        ptr_param<const Dense<value_type>> scale_other,
        ptr_param<const Csr> mtx_other) const;
 
    std::pair<std::unique_ptr<Csr>, scale_add_reuse_info> add_scale_reuse(
        ptr_param<const Dense<value_type>> scale_this,
        ptr_param<const Dense<value_type>> scale_other,
        ptr_param<const Csr> mtx_other) const;
 
    struct permuting_reuse_info {
        explicit permuting_reuse_info();
 
        explicit permuting_reuse_info(
            std::unique_ptr<Permutation<index_type>> value_permutation);
 
        void update_values(ptr_param<const Csr> input,
                           ptr_param<Csr> output) const;
 
        std::unique_ptr<Permutation<IndexType>> value_permutation;
    };
 
    std::pair<std::unique_ptr<Csr>, permuting_reuse_info> transpose_reuse()
        const;
 
    std::unique_ptr<Csr> permute(
        ptr_param<const Permutation<index_type>> permutation,
        permute_mode mode = permute_mode::symmetric) const;
 
    std::unique_ptr<Csr> permute(
        ptr_param<const Permutation<index_type>> row_permutation,
        ptr_param<const Permutation<index_type>> column_permutation,
        bool invert = false) const;
 
    std::pair<std::unique_ptr<Csr>, permuting_reuse_info> permute_reuse(
        ptr_param<const Permutation<index_type>> permutation,
        permute_mode mode = permute_mode::symmetric) const;
 
    std::pair<std::unique_ptr<Csr>, permuting_reuse_info> permute_reuse(
        ptr_param<const Permutation<index_type>> row_permutation,
        ptr_param<const Permutation<index_type>> column_permutation,
        bool invert = false) const;
 
    std::unique_ptr<Csr> scale_permute(
        ptr_param<const ScaledPermutation<value_type, index_type>> permutation,
        permute_mode = permute_mode::symmetric) const;
 
    std::unique_ptr<Csr> scale_permute(
        ptr_param<const ScaledPermutation<value_type, index_type>>
            row_permutation,
        ptr_param<const ScaledPermutation<value_type, index_type>>
            column_permutation,
        bool invert = false) const;
 
    std::unique_ptr<LinOp> permute(
        const array<IndexType>* permutation_indices) const override;
 
    std::unique_ptr<LinOp> inverse_permute(
        const array<IndexType>* inverse_permutation_indices) const override;
 
    std::unique_ptr<LinOp> row_permute(
        const array<IndexType>* permutation_indices) const override;
 
    std::unique_ptr<LinOp> column_permute(
        const array<IndexType>* permutation_indices) const override;
 
    std::unique_ptr<LinOp> inverse_row_permute(
        const array<IndexType>* inverse_permutation_indices) const override;
 
    std::unique_ptr<LinOp> inverse_column_permute(
        const array<IndexType>* inverse_permutation_indices) const override;
 
    std::unique_ptr<Diagonal<ValueType>> extract_diagonal() const override;
 
    std::unique_ptr<absolute_type> compute_absolute() const override;
 
    void compute_absolute_inplace() override;
 
    void sort_by_column_index();
 
    /*
     * Tests if all row entry pairs (value, col_idx) are sorted by column index
     *
     * @returns True if all row entry pairs (value, col_idx) are sorted by
     *          column index
     */
    bool is_sorted_by_column_index() const;
 
    value_type* get_values() noexcept { return values_.get_data(); }
 
    const value_type* get_const_values() const noexcept
    {
        return values_.get_const_data();
    }
 
    std::unique_ptr<Dense<ValueType>> create_value_view();
 
    std::unique_ptr<const Dense<ValueType>> create_const_value_view() const;
 
    index_type* get_col_idxs() noexcept { return col_idxs_.get_data(); }
 
    const index_type* get_const_col_idxs() const noexcept
    {
        return col_idxs_.get_const_data();
    }
 
    index_type* get_row_ptrs() noexcept { return row_ptrs_.get_data(); }
 
    const index_type* get_const_row_ptrs() const noexcept
    {
        return row_ptrs_.get_const_data();
    }
 
    index_type* get_srow() noexcept { return srow_.get_data(); }
 
    const index_type* get_const_srow() const noexcept
    {
        return srow_.get_const_data();
    }
 
    size_type get_num_srow_elements() const noexcept
    {
        return srow_.get_size();
    }
 
    size_type get_num_stored_elements() const noexcept
    {
        return values_.get_size();
    }
 
    csr::spmv_strategy get_strategy() const noexcept;
 
    void set_strategy(csr::spmv_strategy strategy)
    {
        strategy_ = strategy;
        this->make_srow();
    }
 
    void scale(ptr_param<const LinOp> alpha)
    {
        auto exec = this->get_executor();
        GKO_ASSERT_EQUAL_DIMENSIONS(alpha, dim<2>(1, 1));
        this->scale_impl(make_temporary_clone(exec, alpha).get());
    }
 
    void inv_scale(ptr_param<const LinOp> alpha)
    {
        auto exec = this->get_executor();
        GKO_ASSERT_EQUAL_DIMENSIONS(alpha, dim<2>(1, 1));
        this->inv_scale_impl(make_temporary_clone(exec, alpha).get());
    }
 
    void validate_data() const override;
 
    static std::unique_ptr<Csr> create(std::shared_ptr<const Executor> exec,
                                       csr::spmv_strategy strategy);
 
    static std::unique_ptr<Csr> create(
        std::shared_ptr<const Executor> exec, const dim<2>& size = {},
        size_type num_nonzeros = {},
        csr::spmv_strategy strategy = csr::spmv_strategy::automatic);
 
    static std::unique_ptr<Csr> create(
        std::shared_ptr<const Executor> exec, const dim<2>& size,
        array<value_type> values, array<index_type> col_idxs,
        array<index_type> row_ptrs,
        csr::spmv_strategy strategy = csr::spmv_strategy::automatic);
 
    template <typename InputValueType, typename InputColumnIndexType,
              typename InputRowPtrType>
    GKO_DEPRECATED(
        "explicitly construct the gko::array argument instead of passing "
        "initializer lists")
    static std::unique_ptr<Csr> create(
        std::shared_ptr<const Executor> exec, const dim<2>& size,
        std::initializer_list<InputValueType> values,
        std::initializer_list<InputColumnIndexType> col_idxs,
        std::initializer_list<InputRowPtrType> row_ptrs)
    {
        return create(exec, size, array<value_type>{exec, std::move(values)},
                      array<index_type>{exec, std::move(col_idxs)},
                      array<index_type>{exec, std::move(row_ptrs)});
    }
 
    static std::unique_ptr<const Csr> create_const(
        std::shared_ptr<const Executor> exec, const dim<2>& size,
        gko::detail::const_array_view<ValueType>&& values,
        gko::detail::const_array_view<IndexType>&& col_idxs,
        gko::detail::const_array_view<IndexType>&& row_ptrs,
        csr::spmv_strategy strategy = csr::spmv_strategy::automatic);
 
    GKO_BEGIN_DISABLE_DEPRECATION_WARNINGS
 
    [[deprecated("please use enum version")]] static std::unique_ptr<Csr>
    create(std::shared_ptr<const Executor> exec,
           std::shared_ptr<strategy_type> strategy);
 
    [[deprecated("please use enum version")]] static std::unique_ptr<Csr>
    create(std::shared_ptr<const Executor> exec, const dim<2>& size,
           array<value_type> values, array<index_type> col_idxs,
           array<index_type> row_ptrs, std::shared_ptr<strategy_type> strategy);
 
    [[deprecated("please use enum version")]] static std::unique_ptr<const Csr>
    create_const(std::shared_ptr<const Executor> exec, const dim<2>& size,
                 gko::detail::const_array_view<ValueType>&& values,
                 gko::detail::const_array_view<IndexType>&& col_idxs,
                 gko::detail::const_array_view<IndexType>&& row_ptrs,
                 std::shared_ptr<strategy_type> strategy);
 
    GKO_END_DISABLE_DEPRECATION_WARNINGS
 
    std::unique_ptr<Csr<ValueType, IndexType>> create_submatrix(
        const index_set<IndexType>& row_index_set,
        const index_set<IndexType>& column_index_set) const;
 
    std::unique_ptr<Csr<ValueType, IndexType>> create_submatrix(
        const span& row_span, const span& column_span) const;
 
    Csr& operator=(const Csr&);
 
    Csr& operator=(Csr&&);
 
    Csr(const Csr&);
 
    Csr(Csr&&);
 
protected:
    Csr(std::shared_ptr<const Executor> exec, const dim<2>& size = {},
        size_type num_nonzeros = {},
        csr::spmv_strategy strategy = csr::spmv_strategy::automatic);
 
    Csr(std::shared_ptr<const Executor> exec, const dim<2>& size,
        array<value_type> values, array<index_type> col_idxs,
        array<index_type> row_ptrs,
        csr::spmv_strategy strategy = csr::spmv_strategy::automatic);
 
    void apply_impl(const LinOp* b, LinOp* x) const override;
 
    void apply_impl(const LinOp* alpha, const LinOp* b, const LinOp* beta,
                    LinOp* x) const override;
 
    void make_srow();
 
    virtual void scale_impl(const LinOp* alpha);
 
    virtual void inv_scale_impl(const LinOp* alpha);
 
    csr::spmv_strategy get_actual_strategy() const noexcept;
 
private:
    csr::spmv_strategy strategy_;
    array<value_type> values_;
    array<index_type> col_idxs_;
    array<index_type> row_ptrs_;
    array<index_type> srow_;
    index_type max_nnz_per_row_;
 
    void add_scaled_identity_impl(const LinOp* a, const LinOp* b) override;
};
 
 
}  // namespace matrix
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_MATRIX_CSR_HPP_