hybrid.hpp

// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_MATRIX_HYBRID_HPP_
#define GKO_PUBLIC_CORE_MATRIX_HYBRID_HPP_
 
 
#include <algorithm>
 
#include <ginkgo/core/base/array.hpp>
#include <ginkgo/core/base/lin_op.hpp>
#include <ginkgo/core/matrix/coo.hpp>
#include <ginkgo/core/matrix/csr.hpp>
#include <ginkgo/core/matrix/ell.hpp>
 
 
namespace gko {
namespace matrix {
 
 
template <typename ValueType>
class Dense;
 
template <typename ValueType, typename IndexType>
class Csr;
 
 
template <typename ValueType = default_precision, typename IndexType = int32>
class Hybrid
    : public EnableLinOp<Hybrid<ValueType, IndexType>>,
      public ConvertibleTo<Hybrid<next_precision<ValueType>, IndexType>>,
#if GINKGO_ENABLE_HALF || GINKGO_ENABLE_BFLOAT16
      public ConvertibleTo<Hybrid<next_precision<ValueType, 2>, IndexType>>,
#endif
#if GINKGO_ENABLE_HALF && GINKGO_ENABLE_BFLOAT16
      public ConvertibleTo<Hybrid<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<Hybrid<ValueType, IndexType>>> {
    friend class EnablePolymorphicObject<Hybrid, LinOp>;
    friend class Dense<ValueType>;
    friend class Csr<ValueType, IndexType>;
    friend class Hybrid<to_complex<ValueType>, IndexType>;
 
 
public:
    using EnableLinOp<Hybrid>::convert_to;
    using EnableLinOp<Hybrid>::move_to;
    using ConvertibleTo<
        Hybrid<next_precision<ValueType>, IndexType>>::convert_to;
    using ConvertibleTo<Hybrid<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 coo_type = Coo<ValueType, IndexType>;
    using ell_type = Ell<ValueType, IndexType>;
    using absolute_type = remove_complex<Hybrid>;
 
 
    class strategy_type {
    public:
        strategy_type()
            : ell_num_stored_elements_per_row_(zero<size_type>()),
              coo_nnz_(zero<size_type>())
        {}
 
        void compute_hybrid_config(const array<size_type>& row_nnz,
                                   size_type* ell_num_stored_elements_per_row,
                                   size_type* coo_nnz)
        {
            array<size_type> ref_row_nnz(row_nnz.get_executor()->get_master(),
                                         row_nnz.get_size());
            ref_row_nnz = row_nnz;
            ell_num_stored_elements_per_row_ =
                this->compute_ell_num_stored_elements_per_row(&ref_row_nnz);
            coo_nnz_ = this->compute_coo_nnz(ref_row_nnz);
            *ell_num_stored_elements_per_row = ell_num_stored_elements_per_row_;
            *coo_nnz = coo_nnz_;
        }
 
        size_type get_ell_num_stored_elements_per_row() const noexcept
        {
            return ell_num_stored_elements_per_row_;
        }
 
        size_type get_coo_nnz() const noexcept { return coo_nnz_; }
 
        virtual size_type compute_ell_num_stored_elements_per_row(
            array<size_type>* row_nnz) const = 0;
 
    protected:
        size_type compute_coo_nnz(const array<size_type>& row_nnz) const
        {
            size_type coo_nnz = 0;
            auto row_nnz_val = row_nnz.get_const_data();
            for (size_type i = 0; i < row_nnz.get_size(); i++) {
                if (row_nnz_val[i] > ell_num_stored_elements_per_row_) {
                    coo_nnz +=
                        row_nnz_val[i] - ell_num_stored_elements_per_row_;
                }
            }
            return coo_nnz;
        }
 
    private:
        size_type ell_num_stored_elements_per_row_;
        size_type coo_nnz_;
    };
 
    class column_limit : public strategy_type {
    public:
        explicit column_limit(size_type num_column = 0)
            : num_columns_(num_column)
        {}
 
        size_type compute_ell_num_stored_elements_per_row(
            array<size_type>* row_nnz) const override
        {
            return num_columns_;
        }
 
        auto get_num_columns() const { return num_columns_; }
 
    private:
        size_type num_columns_;
    };
 
    class imbalance_limit : public strategy_type {
    public:
        explicit imbalance_limit(double percent = 0.8) : percent_(percent)
        {
            percent_ = std::min(percent_, 1.0);
            percent_ = std::max(percent_, 0.0);
        }
 
        size_type compute_ell_num_stored_elements_per_row(
            array<size_type>* row_nnz) const override
        {
            auto row_nnz_val = row_nnz->get_data();
            auto num_rows = row_nnz->get_size();
            if (num_rows == 0) {
                return 0;
            }
            std::sort(row_nnz_val, row_nnz_val + num_rows);
            if (percent_ < 1) {
                auto percent_pos = static_cast<size_type>(num_rows * percent_);
                return row_nnz_val[percent_pos];
            } else {
                return row_nnz_val[num_rows - 1];
            }
        }
 
        auto get_percentage() const { return percent_; }
 
    private:
        double percent_;
    };
 
    class imbalance_bounded_limit : public strategy_type {
    public:
        imbalance_bounded_limit(double percent = 0.8, double ratio = 0.0001)
            : strategy_(imbalance_limit(percent)), ratio_(ratio)
        {}
 
        size_type compute_ell_num_stored_elements_per_row(
            array<size_type>* row_nnz) const override
        {
            auto num_rows = row_nnz->get_size();
            auto ell_cols =
                strategy_.compute_ell_num_stored_elements_per_row(row_nnz);
            return std::min(ell_cols,
                            static_cast<size_type>(num_rows * ratio_));
        }
 
        auto get_percentage() const { return strategy_.get_percentage(); }
 
        auto get_ratio() const { return ratio_; }
 
    private:
        imbalance_limit strategy_;
        double ratio_;
    };
 
 
    class minimal_storage_limit : public strategy_type {
    public:
        minimal_storage_limit()
            : strategy_(
                  imbalance_limit(static_cast<double>(sizeof(IndexType)) /
                                  (sizeof(ValueType) + 2 * sizeof(IndexType))))
        {}
 
        size_type compute_ell_num_stored_elements_per_row(
            array<size_type>* row_nnz) const override
        {
            return strategy_.compute_ell_num_stored_elements_per_row(row_nnz);
        }
 
        auto get_percentage() const { return strategy_.get_percentage(); }
 
    private:
        imbalance_limit strategy_;
    };
 
 
    class automatic : public strategy_type {
    public:
        automatic() : strategy_(imbalance_bounded_limit(1.0 / 3.0, 0.001)) {}
 
        size_type compute_ell_num_stored_elements_per_row(
            array<size_type>* row_nnz) const override
        {
            return strategy_.compute_ell_num_stored_elements_per_row(row_nnz);
        }
 
    private:
        imbalance_bounded_limit strategy_;
    };
 
    friend class Hybrid<previous_precision<ValueType>, IndexType>;
 
    void convert_to(
        Hybrid<next_precision<ValueType>, IndexType>* result) const override;
 
    void move_to(Hybrid<next_precision<ValueType>, IndexType>* result) override;
 
#if GINKGO_ENABLE_HALF || GINKGO_ENABLE_BFLOAT16
    friend class Hybrid<previous_precision<ValueType, 2>, IndexType>;
    using ConvertibleTo<
        Hybrid<next_precision<ValueType, 2>, IndexType>>::convert_to;
    using ConvertibleTo<
        Hybrid<next_precision<ValueType, 2>, IndexType>>::move_to;
 
    void convert_to(
        Hybrid<next_precision<ValueType, 2>, IndexType>* result) const override;
 
    void move_to(
        Hybrid<next_precision<ValueType, 2>, IndexType>* result) override;
#endif
 
#if GINKGO_ENABLE_HALF && GINKGO_ENABLE_BFLOAT16
    friend class Hybrid<previous_precision<ValueType, 3>, IndexType>;
    using ConvertibleTo<
        Hybrid<next_precision<ValueType, 3>, IndexType>>::convert_to;
    using ConvertibleTo<
        Hybrid<next_precision<ValueType, 3>, IndexType>>::move_to;
 
    void convert_to(
        Hybrid<next_precision<ValueType, 3>, IndexType>* result) const override;
 
    void move_to(
        Hybrid<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_ell_values() noexcept { return ell_->get_values(); }
 
    const value_type* get_const_ell_values() const noexcept
    {
        return ell_->get_const_values();
    }
 
    index_type* get_ell_col_idxs() noexcept { return ell_->get_col_idxs(); }
 
    const index_type* get_const_ell_col_idxs() const noexcept
    {
        return ell_->get_const_col_idxs();
    }
 
    size_type get_ell_num_stored_elements_per_row() const noexcept
    {
        return ell_->get_num_stored_elements_per_row();
    }
 
    size_type get_ell_stride() const noexcept { return ell_->get_stride(); }
 
    size_type get_ell_num_stored_elements() const noexcept
    {
        return ell_->get_num_stored_elements();
    }
 
    value_type& ell_val_at(size_type row, size_type idx) noexcept
    {
        return ell_->val_at(row, idx);
    }
 
    value_type ell_val_at(size_type row, size_type idx) const noexcept
    {
        return ell_->val_at(row, idx);
    }
 
    index_type& ell_col_at(size_type row, size_type idx) noexcept
    {
        return ell_->col_at(row, idx);
    }
 
    index_type ell_col_at(size_type row, size_type idx) const noexcept
    {
        return ell_->col_at(row, idx);
    }
 
    const ell_type* get_ell() const noexcept { return ell_.get(); }
 
    value_type* get_coo_values() noexcept { return coo_->get_values(); }
 
    const value_type* get_const_coo_values() const noexcept
    {
        return coo_->get_const_values();
    }
 
    index_type* get_coo_col_idxs() noexcept { return coo_->get_col_idxs(); }
 
    const index_type* get_const_coo_col_idxs() const noexcept
    {
        return coo_->get_const_col_idxs();
    }
 
    index_type* get_coo_row_idxs() noexcept { return coo_->get_row_idxs(); }
 
    const index_type* get_const_coo_row_idxs() const noexcept
    {
        return coo_->get_const_row_idxs();
    }
 
    size_type get_coo_num_stored_elements() const noexcept
    {
        return coo_->get_num_stored_elements();
    }
 
    const coo_type* get_coo() const noexcept { return coo_.get(); }
 
    size_type get_num_stored_elements() const noexcept
    {
        return coo_->get_num_stored_elements() +
               ell_->get_num_stored_elements();
    }
 
    std::shared_ptr<strategy_type> get_strategy() const noexcept
    {
        return strategy_;
    }
 
    template <typename HybType>
    std::shared_ptr<typename HybType::strategy_type> get_strategy() const;
 
    static std::unique_ptr<Hybrid> create(
        std::shared_ptr<const Executor> exec,
        std::shared_ptr<strategy_type> strategy =
            std::make_shared<automatic>());
 
    static std::unique_ptr<Hybrid> create(
        std::shared_ptr<const Executor> exec, const dim<2>& size,
        std::shared_ptr<strategy_type> strategy =
            std::make_shared<automatic>());
 
    static std::unique_ptr<Hybrid> create(
        std::shared_ptr<const Executor> exec, const dim<2>& size,
        size_type num_stored_elements_per_row,
        std::shared_ptr<strategy_type> strategy =
            std::make_shared<automatic>());
 
    static std::unique_ptr<Hybrid> create(
        std::shared_ptr<const Executor> exec, const dim<2>& size,
        size_type num_stored_elements_per_row, size_type stride,
        std::shared_ptr<strategy_type> strategy);
 
    static std::unique_ptr<Hybrid> create(
        std::shared_ptr<const Executor> exec, const dim<2>& size,
        size_type num_stored_elements_per_row, size_type stride,
        size_type num_nonzeros = {},
        std::shared_ptr<strategy_type> strategy =
            std::make_shared<automatic>());
 
    Hybrid& operator=(const Hybrid&);
 
    Hybrid& operator=(Hybrid&&);
 
    Hybrid(const Hybrid&);
 
    Hybrid(Hybrid&&);
 
protected:
    Hybrid(std::shared_ptr<const Executor> exec, const dim<2>& size = {},
           size_type num_stored_elements_per_row = 0, size_type stride = 0,
           size_type num_nonzeros = 0,
           std::shared_ptr<strategy_type> strategy =
               std::make_shared<automatic>());
 
    void resize(dim<2> new_size, size_type ell_row_nnz, size_type coo_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;
 
private:
    std::unique_ptr<ell_type> ell_;
    std::unique_ptr<coo_type> coo_;
    std::shared_ptr<strategy_type> strategy_;
};
 
 
template <typename ValueType, typename IndexType>
template <typename HybType>
std::shared_ptr<typename HybType::strategy_type>
Hybrid<ValueType, IndexType>::get_strategy() const
{
    static_assert(
        std::is_same<HybType, Hybrid<typename HybType::value_type,
                                     typename HybType::index_type>>::value,
        "The given `HybType` type must be of type `matrix::Hybrid`!");
 
    std::shared_ptr<typename HybType::strategy_type> strategy;
    if (std::dynamic_pointer_cast<automatic>(strategy_)) {
        strategy = std::make_shared<typename HybType::automatic>();
    } else if (auto temp = std::dynamic_pointer_cast<minimal_storage_limit>(
                   strategy_)) {
        // minimal_storage_limit is related to ValueType and IndexType size.
        if (sizeof(value_type) == sizeof(typename HybType::value_type) &&
            sizeof(index_type) == sizeof(typename HybType::index_type)) {
            strategy =
                std::make_shared<typename HybType::minimal_storage_limit>();
        } else {
            strategy = std::make_shared<typename HybType::imbalance_limit>(
                temp->get_percentage());
        }
    } else if (auto temp = std::dynamic_pointer_cast<imbalance_bounded_limit>(
                   strategy_)) {
        strategy = std::make_shared<typename HybType::imbalance_bounded_limit>(
            temp->get_percentage(), temp->get_ratio());
    } else if (auto temp =
                   std::dynamic_pointer_cast<imbalance_limit>(strategy_)) {
        strategy = std::make_shared<typename HybType::imbalance_limit>(
            temp->get_percentage());
    } else if (auto temp = std::dynamic_pointer_cast<column_limit>(strategy_)) {
        strategy = std::make_shared<typename HybType::column_limit>(
            temp->get_num_columns());
    } else {
        GKO_NOT_SUPPORTED(strategy_);
    }
    return strategy;
}
 
 
}  // namespace matrix
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_MATRIX_HYBRID_HPP_