batch_ell.hpp

// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_MATRIX_BATCH_ELL_HPP_
#define GKO_PUBLIC_CORE_MATRIX_BATCH_ELL_HPP_
 
 
#include <initializer_list>
#include <vector>
 
#include <ginkgo/core/base/array.hpp>
#include <ginkgo/core/base/batch_lin_op.hpp>
#include <ginkgo/core/base/batch_multi_vector.hpp>
#include <ginkgo/core/base/executor.hpp>
#include <ginkgo/core/base/mtx_io.hpp>
#include <ginkgo/core/base/range_accessors.hpp>
#include <ginkgo/core/base/types.hpp>
#include <ginkgo/core/base/utils.hpp>
#include <ginkgo/core/base/utils_helper.hpp>
#include <ginkgo/core/matrix/ell.hpp>
 
 
namespace gko {
namespace batch {
namespace matrix {
 

template <typename ValueType = default_precision, typename IndexType = int32>
class Ell final
    : public EnableBatchLinOp<Ell<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<Ell<next_precision<ValueType>, IndexType>> {
    friend class EnablePolymorphicObject<Ell, BatchLinOp>;
    friend class Ell<to_complex<ValueType>, IndexType>;
    friend class Ell<previous_precision<ValueType>, IndexType>;
    GKO_ASSERT_SUPPORTED_VALUE_TYPE;
    static_assert(std::is_same<IndexType, int32>::value,
                  "IndexType must be a 32 bit integer");
 
public:
    using EnableBatchLinOp<Ell>::convert_to;
    using EnableBatchLinOp<Ell>::move_to;
 
    using value_type = ValueType;
    using index_type = IndexType;
    using unbatch_type = gko::matrix::Ell<value_type, index_type>;
    using absolute_type = remove_complex<Ell>;
    using complex_type = to_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

    std::unique_ptr<unbatch_type> create_view_for_item(size_type item_id);

    std::unique_ptr<const unbatch_type> create_const_view_for_item(
        size_type item_id) const;

    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();
    }

    index_type get_num_stored_elements_per_row() const noexcept
    {
        return num_elems_per_row_;
    }

    size_type get_num_stored_elements() const noexcept
    {
        return values_.get_size();
    }

    size_type get_num_elements_per_item() const noexcept
    {
        return this->get_num_stored_elements() / this->get_num_batch_items();
    }

    index_type* get_col_idxs_for_item(size_type batch_id) noexcept
    {
        GKO_ASSERT(batch_id < this->get_num_batch_items());
        return col_idxs_.get_data();
    }

    const index_type* get_const_col_idxs_for_item(
        size_type batch_id) const noexcept
    {
        GKO_ASSERT(batch_id < this->get_num_batch_items());
        return col_idxs_.get_const_data();
    }

    value_type* get_values_for_item(size_type batch_id) noexcept
    {
        GKO_ASSERT(batch_id < this->get_num_batch_items());
        return values_.get_data() +
               batch_id * this->get_num_elements_per_item();
    }

    const value_type* get_const_values_for_item(
        size_type batch_id) const noexcept
    {
        GKO_ASSERT(batch_id < this->get_num_batch_items());
        return values_.get_const_data() +
               batch_id * this->get_num_elements_per_item();
    }