batch_multi_vector.hpp

// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_BASE_BATCH_MULTI_VECTOR_HPP_
#define GKO_PUBLIC_CORE_BASE_BATCH_MULTI_VECTOR_HPP_
 
 
#include <initializer_list>
#include <vector>
 
#include <ginkgo/core/base/array.hpp>
#include <ginkgo/core/base/batch_dim.hpp>
#include <ginkgo/core/base/dim.hpp>
#include <ginkgo/core/base/executor.hpp>
#include <ginkgo/core/base/mtx_io.hpp>
#include <ginkgo/core/base/polymorphic_object.hpp>
#include <ginkgo/core/base/range_accessors.hpp>
#include <ginkgo/core/base/types.hpp>
#include <ginkgo/core/base/utils.hpp>
#include <ginkgo/core/matrix/dense.hpp>
 
 
namespace gko {
namespace batch {
 

template <typename ValueType = default_precision>
class MultiVector
    : public EnablePolymorphicObject<MultiVector<ValueType>>,
      public EnablePolymorphicAssignment<MultiVector<ValueType>>,
#if GINKGO_ENABLE_HALF || GINKGO_ENABLE_BFLOAT16
      public ConvertibleTo<MultiVector<next_precision<ValueType, 2>>>,
#endif
#if GINKGO_ENABLE_HALF && GINKGO_ENABLE_BFLOAT16
      public ConvertibleTo<MultiVector<next_precision<ValueType, 3>>>,
#endif
      public ConvertibleTo<MultiVector<next_precision<ValueType>>> {
    friend class EnablePolymorphicObject<MultiVector>;
    friend class MultiVector<to_complex<ValueType>>;
    friend class MultiVector<previous_precision<ValueType>>;
    GKO_ASSERT_SUPPORTED_VALUE_TYPE;
 
public:
    using EnablePolymorphicAssignment<MultiVector>::convert_to;
    using EnablePolymorphicAssignment<MultiVector>::move_to;
    using ConvertibleTo<MultiVector<next_precision<ValueType>>>::convert_to;
    using ConvertibleTo<MultiVector<next_precision<ValueType>>>::move_to;
 
    using value_type = ValueType;
    using index_type = int32;
    using unbatch_type = gko::matrix::Dense<ValueType>;
    using absolute_type = remove_complex<MultiVector<ValueType>>;
    using complex_type = to_complex<MultiVector<ValueType>>;

    static std::unique_ptr<MultiVector> create_with_config_of(
        ptr_param<const MultiVector> other);
 
    void convert_to(
        MultiVector<next_precision<ValueType>>* result) const override;
 
    void move_to(MultiVector<next_precision<ValueType>>* result) override;
 
#if GINKGO_ENABLE_HALF || GINKGO_ENABLE_BFLOAT16
    friend class MultiVector<previous_precision<ValueType, 2>>;
    using ConvertibleTo<MultiVector<next_precision<ValueType, 2>>>::convert_to;
    using ConvertibleTo<MultiVector<next_precision<ValueType, 2>>>::move_to;
 
    void convert_to(
        MultiVector<next_precision<ValueType, 2>>* result) const override;
 
    void move_to(MultiVector<next_precision<ValueType, 2>>* result) override;
#endif
 
#if GINKGO_ENABLE_HALF && GINKGO_ENABLE_BFLOAT16
    friend class MultiVector<previous_precision<ValueType, 3>>;
    using ConvertibleTo<MultiVector<next_precision<ValueType, 3>>>::convert_to;
    using ConvertibleTo<MultiVector<next_precision<ValueType, 3>>>::move_to;
 
    void convert_to(
        MultiVector<next_precision<ValueType, 3>>* result) const override;
 
    void move_to(MultiVector<next_precision<ValueType, 3>>* 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;

    batch_dim<2> get_size() const { return batch_size_; }

    size_type get_num_batch_items() const
    {
        return batch_size_.get_num_batch_items();
    }

    dim<2> get_common_size() const { return batch_size_.get_common_size(); }

    value_type* get_values() noexcept { return values_.get_data(); }

    const value_type* get_const_values() const noexcept
    {
        return values_.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() + this->get_cumulative_offset(batch_id);
    }

    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() + this->get_cumulative_offset(batch_id);
    }

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

    size_type get_cumulative_offset(size_type batch_id) const
    {
        return batch_id * this->get_common_size()[0] *
               this->get_common_size()[1];
    }

    value_type& at(size_type batch_id, size_type row, size_type col)
    {
        GKO_ASSERT(batch_id < this->get_num_batch_items());
        return values_.get_data()[linearize_index(batch_id, row, col)];
    }

    value_type at(size_type batch_id, size_type row, size_type col) const
    {
        GKO_ASSERT(batch_id < this->get_num_batch_items());
        return values_.get_const_data()[linearize_index(batch_id, row, col)];
    }

    ValueType& at(size_type batch_id, size_type idx) noexcept
    {
        return values_.get_data()[linearize_index(batch_id, idx)];
    }

    ValueType at(size_type batch_id, size_type idx) const noexcept
    {
        return values_.get_const_data()[linearize_index(batch_id, idx)];
    }

    void scale(ptr_param<const MultiVector<ValueType>> alpha);

    void add_scaled(ptr_param<const MultiVector<ValueType>> alpha,
                    ptr_param<const MultiVector<ValueType>> b);

    void compute_dot(ptr_param<const MultiVector<ValueType>> b,
                     ptr_param<MultiVector<ValueType>> result) const;

    void compute_conj_dot(ptr_param<const MultiVector<ValueType>> b,
                          ptr_param<MultiVector<ValueType>> result) const;

    void compute_norm2(
        ptr_param<MultiVector<remove_complex<ValueType>>> result) const;

    static std::unique_ptr<MultiVector> create(
        std::shared_ptr<const Executor> exec,
        const batch_dim<2>& size = batch_dim<2>{});

    static std::unique_ptr<MultiVector> create(
        std::shared_ptr<const Executor> exec, const batch_dim<2>& size,
        array<value_type> values);

    template <typename InputValueType>
    GKO_DEPRECATED(
        "explicitly construct the gko::array argument instead of passing an "
        "initializer list")
    static std::unique_ptr<MultiVector> create(
        std::shared_ptr<const Executor> exec, const batch_dim<2>& size,
        std::initializer_list<InputValueType> values)
    {
        return create(exec, size, array<index_type>{exec, std::move(values)});
    }

    static std::unique_ptr<const MultiVector> create_const(
        std::shared_ptr<const Executor> exec, const batch_dim<2>& sizes,
        gko::detail::const_array_view<ValueType>&& values);

    void fill(ValueType value);
 
private:
    inline size_type compute_num_elems(const batch_dim<2>& size)
    {
        return size.get_num_batch_items() * size.get_common_size()[0] *
               size.get_common_size()[1];
    }
 
protected:
    void set_size(const batch_dim<2>& value) noexcept;
 
    MultiVector(std::shared_ptr<const Executor> exec,
                const batch_dim<2>& size = batch_dim<2>{});
 
    MultiVector(std::shared_ptr<const Executor> exec, const batch_dim<2>& size,
                array<value_type> values);

    std::unique_ptr<MultiVector> create_with_same_config() const;
 
    size_type linearize_index(size_type batch, size_type row,
                              size_type col) const noexcept
    {
        return this->get_cumulative_offset(batch) +
               row * batch_size_.get_common_size()[1] + col;
    }
 
    size_type linearize_index(size_type batch, size_type idx) const noexcept
    {
        return linearize_index(batch, idx / this->get_common_size()[1],
                               idx % this->get_common_size()[1]);
    }
 
private:
    batch_dim<2> batch_size_;
    array<value_type> values_;
};
 
 
}  // namespace batch
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_BASE_BATCH_MULTI_VECTOR_HPP_