ell.hpp

// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_MATRIX_ELL_HPP_
#define GKO_PUBLIC_CORE_MATRIX_ELL_HPP_
 
 
#include <ginkgo/core/base/array.hpp>
#include <ginkgo/core/base/lin_op.hpp>
 
 
namespace gko {
namespace matrix {
 
 
template <typename ValueType>
class Dense;
 
template <typename ValueType, typename IndexType>
class Coo;
 
template <typename ValueType, typename IndexType>
class Csr;
 
template <typename ValueType, typename IndexType>
class Hybrid;
 
 
template <typename ValueType = default_precision, typename IndexType = int32>
class Ell : public EnableLinOp<Ell<ValueType, IndexType>>,
            public ConvertibleTo<Ell<next_precision<ValueType>, IndexType>>,
#if GINKGO_ENABLE_HALF || GINKGO_ENABLE_BFLOAT16
            public ConvertibleTo<Ell<next_precision<ValueType, 2>, IndexType>>,
#endif
#if GINKGO_ENABLE_HALF && GINKGO_ENABLE_BFLOAT16
            public ConvertibleTo<Ell<next_precision<ValueType, 3>, IndexType>>,
#endif
            public ConvertibleTo<Dense<ValueType>>,
            public ConvertibleTo<Csr<ValueType, IndexType>>,
            public DiagonalExtractable<ValueType>,
            public ReadableFromMatrixData<ValueType, IndexType>,
            public WritableToMatrixData<ValueType, IndexType>,
            public EnableAbsoluteComputation<
                remove_complex<Ell<ValueType, IndexType>>> {
    friend class EnablePolymorphicObject<Ell, LinOp>;
    friend class Dense<ValueType>;
    friend class Coo<ValueType, IndexType>;
    friend class Csr<ValueType, IndexType>;
    friend class Ell<to_complex<ValueType>, IndexType>;
    friend class Ell<previous_precision<ValueType>, IndexType>;
    friend class Hybrid<ValueType, IndexType>;
 
public:
    using EnableLinOp<Ell>::convert_to;
    using EnableLinOp<Ell>::move_to;
    using ConvertibleTo<Ell<next_precision<ValueType>, IndexType>>::convert_to;
    using ConvertibleTo<Ell<next_precision<ValueType>, IndexType>>::move_to;
    using ConvertibleTo<Dense<ValueType>>::convert_to;
    using ConvertibleTo<Dense<ValueType>>::move_to;
    using ConvertibleTo<Csr<ValueType, IndexType>>::convert_to;
    using ConvertibleTo<Csr<ValueType, IndexType>>::move_to;
    using ReadableFromMatrixData<ValueType, IndexType>::read;
 
    using value_type = ValueType;
    using index_type = IndexType;
    using mat_data = matrix_data<ValueType, IndexType>;
    using device_mat_data = device_matrix_data<ValueType, IndexType>;
    using absolute_type = remove_complex<Ell>;
 
    void convert_to(
        Ell<next_precision<ValueType>, IndexType>* result) const override;
 
    void move_to(Ell<next_precision<ValueType>, IndexType>* result) override;
 
#if GINKGO_ENABLE_HALF || GINKGO_ENABLE_BFLOAT16
    friend class Ell<previous_precision<ValueType, 2>, IndexType>;
    using ConvertibleTo<
        Ell<next_precision<ValueType, 2>, IndexType>>::convert_to;
    using ConvertibleTo<Ell<next_precision<ValueType, 2>, IndexType>>::move_to;
 
    void convert_to(
        Ell<next_precision<ValueType, 2>, IndexType>* result) const override;
 
    void move_to(Ell<next_precision<ValueType, 2>, IndexType>* result) override;
#endif
 
#if GINKGO_ENABLE_HALF && GINKGO_ENABLE_BFLOAT16
    friend class Ell<previous_precision<ValueType, 3>, IndexType>;
    using ConvertibleTo<
        Ell<next_precision<ValueType, 3>, IndexType>>::convert_to;
    using ConvertibleTo<Ell<next_precision<ValueType, 3>, IndexType>>::move_to;
 
    void convert_to(
        Ell<next_precision<ValueType, 3>, IndexType>* result) const override;
 
    void move_to(Ell<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(Csr<ValueType, IndexType>* other) const override;
 
    void move_to(Csr<ValueType, IndexType>* other) 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;
 
    std::unique_ptr<Diagonal<ValueType>> extract_diagonal() const override;
 
    std::unique_ptr<absolute_type> compute_absolute() const override;
 
    void compute_absolute_inplace() override;
 
    value_type* get_values() noexcept { return values_.get_data(); }
 
    const value_type* get_const_values() const noexcept
    {
        return values_.get_const_data();
    }
 
    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();
    }
 
    size_type get_num_stored_elements_per_row() const noexcept
    {
        return num_stored_elements_per_row_;
    }
 
    size_type get_stride() const noexcept { return stride_; }
 
    size_type get_num_stored_elements() const noexcept
    {
        return values_.get_size();
    }
 
    value_type& val_at(size_type row, size_type idx) noexcept
    {
        return values_.get_data()[this->linearize_index(row, idx)];
    }
 
    value_type val_at(size_type row, size_type idx) const noexcept
    {
        return values_.get_const_data()[this->linearize_index(row, idx)];
    }
 
    index_type& col_at(size_type row, size_type idx) noexcept
    {
        return this->get_col_idxs()[this->linearize_index(row, idx)];
    }
 
    index_type col_at(size_type row, size_type idx) const noexcept
    {
        return this->get_const_col_idxs()[this->linearize_index(row, idx)];
    }
 
    static std::unique_ptr<Ell> create(
        std::shared_ptr<const Executor> exec, const dim<2>& size = {},
        size_type num_stored_elements_per_row = 0, size_type stride = 0);
 
    static std::unique_ptr<Ell> create(std::shared_ptr<const Executor> exec,
                                       const dim<2>& size,
                                       array<value_type> values,
                                       array<index_type> col_idxs,
                                       size_type num_stored_elements_per_row,
                                       size_type stride);
 
    template <typename InputValueType, typename InputColumnIndexType>
    GKO_DEPRECATED(
        "explicitly construct the gko::array argument instead of passing "
        "initializer lists")
    static std::unique_ptr<Ell> create(
        std::shared_ptr<const Executor> exec, const dim<2>& size,
        std::initializer_list<InputValueType> values,
        std::initializer_list<InputColumnIndexType> col_idxs,
        size_type num_stored_elements_per_row, size_type stride)
    {
        return create(exec, size, array<value_type>{exec, std::move(values)},
                      array<index_type>{exec, std::move(col_idxs)},
                      num_stored_elements_per_row, stride);
    }
 
    static std::unique_ptr<const Ell> 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,
        size_type num_stored_elements_per_row, size_type stride);
 
    Ell& operator=(const Ell&);
 
    Ell& operator=(Ell&&);
 
    Ell(const Ell&);
 
    Ell(Ell&&);
 
protected:
    Ell(std::shared_ptr<const Executor> exec, const dim<2>& size = {},
        size_type num_stored_elements_per_row = 0, size_type stride = 0);
 
    Ell(std::shared_ptr<const Executor> exec, const dim<2>& size,
        array<value_type> values, array<index_type> col_idxs,
        size_type num_stored_elements_per_row, size_type stride);
 
    void resize(dim<2> new_size, size_type max_row_nnz);
 
    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;
 
    size_type linearize_index(size_type row, size_type col) const noexcept
    {
        return row + stride_ * col;
    }
 
private:
    size_type num_stored_elements_per_row_;
    size_type stride_;
    array<value_type> values_;
    array<index_type> col_idxs_;
};
 
 
}  // namespace matrix
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_MATRIX_ELL_HPP_