vector.hpp

// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_DISTRIBUTED_VECTOR_HPP_
#define GKO_PUBLIC_CORE_DISTRIBUTED_VECTOR_HPP_
 
 
#include <ginkgo/config.hpp>
 
 
#if GINKGO_BUILD_MPI
 
 
#include <ginkgo/core/base/dense_cache.hpp>
#include <ginkgo/core/base/lin_op.hpp>
#include <ginkgo/core/base/mpi.hpp>
#include <ginkgo/core/distributed/base.hpp>
#include <ginkgo/core/matrix/dense.hpp>
 
 
namespace gko {
namespace experimental {
namespace distributed {
namespace detail {
 
 
template <typename ValueType>
class VectorCache;
 
 
}  // namespace detail
 
 
template <typename LocalIndexType, typename GlobalIndexType>
class Partition;
 
 
template <typename ValueType = double>
class Vector
    : public EnableLinOp<Vector<ValueType>>,
      public ConvertibleTo<Vector<next_precision<ValueType>>>,
#if GINKGO_ENABLE_HALF || GINKGO_ENABLE_BFLOAT16
      public ConvertibleTo<Vector<next_precision<ValueType, 2>>>,
#endif
#if GINKGO_ENABLE_HALF && GINKGO_ENABLE_BFLOAT16
      public ConvertibleTo<Vector<next_precision<ValueType, 3>>>,
#endif
      public EnableAbsoluteComputation<remove_complex<Vector<ValueType>>>,
      public DistributedBase {
    friend class EnablePolymorphicObject<Vector, LinOp>;
    friend class Vector<to_complex<ValueType>>;
    friend class Vector<remove_complex<ValueType>>;
    friend class Vector<previous_precision<ValueType>>;
    friend class detail::VectorCache<ValueType>;
    GKO_ASSERT_SUPPORTED_VALUE_TYPE;
 
public:
    using EnableLinOp<Vector>::convert_to;
    using EnableLinOp<Vector>::move_to;
    using ConvertibleTo<Vector<next_precision<ValueType>>>::convert_to;
    using ConvertibleTo<Vector<next_precision<ValueType>>>::move_to;
 
    using value_type = ValueType;
    using absolute_type = remove_complex<Vector>;
    using real_type = absolute_type;
    using complex_type = Vector<to_complex<value_type>>;
    using local_vector_type = gko::matrix::Dense<value_type>;
 
    static std::unique_ptr<Vector> create_with_config_of(
        ptr_param<const Vector> other);
 
 
    static std::unique_ptr<Vector> create_with_type_of(
        ptr_param<const Vector> other, std::shared_ptr<const Executor> exec);
 
    static std::unique_ptr<Vector> create_with_type_of(
        ptr_param<const Vector> other, std::shared_ptr<const Executor> exec,
        const dim<2>& global_size, const dim<2>& local_size, size_type stride);
 
    void read_distributed(const device_matrix_data<ValueType, int64>& data,
                          ptr_param<const Partition<int64, int64>> partition);
 
    void read_distributed(const device_matrix_data<ValueType, int64>& data,
                          ptr_param<const Partition<int32, int64>> partition);
 
    void read_distributed(const device_matrix_data<ValueType, int32>& data,
                          ptr_param<const Partition<int32, int32>> partition);
 
    void read_distributed(const matrix_data<ValueType, int64>& data,
                          ptr_param<const Partition<int64, int64>> partition);
 
    void read_distributed(const matrix_data<ValueType, int64>& data,
                          ptr_param<const Partition<int32, int64>> partition);
 
    void read_distributed(const matrix_data<ValueType, int32>& data,
                          ptr_param<const Partition<int32, int32>> partition);
 
    void convert_to(Vector<next_precision<ValueType>>* result) const override;
 
    void move_to(Vector<next_precision<ValueType>>* result) override;
 
#if GINKGO_ENABLE_HALF || GINKGO_ENABLE_BFLOAT16
    friend class Vector<previous_precision<ValueType, 2>>;
    using ConvertibleTo<Vector<next_precision<ValueType, 2>>>::convert_to;
    using ConvertibleTo<Vector<next_precision<ValueType, 2>>>::move_to;
 
    void convert_to(
        Vector<next_precision<ValueType, 2>>* result) const override;
 
    void move_to(Vector<next_precision<ValueType, 2>>* result) override;
#endif
 
#if GINKGO_ENABLE_HALF && GINKGO_ENABLE_BFLOAT16
    friend class Vector<previous_precision<ValueType, 3>>;
    using ConvertibleTo<Vector<next_precision<ValueType, 3>>>::convert_to;
    using ConvertibleTo<Vector<next_precision<ValueType, 3>>>::move_to;
 
    void convert_to(
        Vector<next_precision<ValueType, 3>>* result) const override;
 
    void move_to(Vector<next_precision<ValueType, 3>>* result) override;
#endif
 
    std::unique_ptr<absolute_type> compute_absolute() const override;
 
    void compute_absolute_inplace() override;
 
    std::unique_ptr<complex_type> make_complex() const;
 
    void make_complex(ptr_param<complex_type> result) const;
 
    std::unique_ptr<real_type> get_real() const;
 
    void get_real(ptr_param<real_type> result) const;
 
    std::unique_ptr<real_type> get_imag() const;
 
    void get_imag(ptr_param<real_type> result) const;
 
    void fill(ValueType value);
 
    void scale(ptr_param<const LinOp> alpha);
 
    void inv_scale(ptr_param<const LinOp> alpha);
 
    void add_scaled(ptr_param<const LinOp> alpha, ptr_param<const LinOp> b);
 
    void sub_scaled(ptr_param<const LinOp> alpha, ptr_param<const LinOp> b);
 
    void compute_dot(ptr_param<const LinOp> b, ptr_param<LinOp> result) const;
 
    void compute_dot(ptr_param<const LinOp> b, ptr_param<LinOp> result,
                     array<char>& tmp) const;
 
    void compute_conj_dot(ptr_param<const LinOp> b,
                          ptr_param<LinOp> result) const;
 
    void compute_conj_dot(ptr_param<const LinOp> b, ptr_param<LinOp> result,
                          array<char>& tmp) const;
 
    void compute_squared_norm2(ptr_param<LinOp> result) const;
 
    void compute_squared_norm2(ptr_param<LinOp> result, array<char>& tmp) const;
 
    void compute_norm2(ptr_param<LinOp> result) const;
 
    void compute_norm2(ptr_param<LinOp> result, array<char>& tmp) const;
 
    void compute_norm1(ptr_param<LinOp> result) const;
 
    void compute_norm1(ptr_param<LinOp> result, array<char>& tmp) const;
 
    void compute_mean(ptr_param<LinOp> result) const;
 
    void compute_mean(ptr_param<LinOp> result, array<char>& tmp) const;
 
    value_type& at_local(size_type row, size_type col) noexcept;
 
    value_type at_local(size_type row, size_type col) const noexcept;
 
    ValueType& at_local(size_type idx) noexcept;
 
    ValueType at_local(size_type idx) const noexcept;
 
    value_type* get_local_values();
 
    const value_type* get_const_local_values() const;
 
    const local_vector_type* get_local_vector() const;
 
    std::unique_ptr<const real_type> create_real_view() const;
 
    std::unique_ptr<real_type> create_real_view();
 
    std::unique_ptr<Vector> create_submatrix(local_span rows,
                                             local_span columns,
                                             dim<2> global_size);
 
    size_type get_stride() const noexcept { return local_.get_stride(); }
 
    static std::unique_ptr<Vector> create(std::shared_ptr<const Executor> exec,
                                          mpi::communicator comm,
                                          dim<2> global_size, dim<2> local_size,
                                          size_type stride);
 
    static std::unique_ptr<Vector> create(std::shared_ptr<const Executor> exec,
                                          mpi::communicator comm,
                                          dim<2> global_size = {},
                                          dim<2> local_size = {});
 
    static std::unique_ptr<Vector> create(
        std::shared_ptr<const Executor> exec, mpi::communicator comm,
        dim<2> global_size, std::unique_ptr<local_vector_type> local_vector);
 
    static std::unique_ptr<Vector> create(
        std::shared_ptr<const Executor> exec, mpi::communicator comm,
        std::unique_ptr<local_vector_type> local_vector);
 
    static std::unique_ptr<const Vector> create_const(
        std::shared_ptr<const Executor> exec, mpi::communicator comm,
        dim<2> global_size,
        std::unique_ptr<const local_vector_type> local_vector);
 
    static std::unique_ptr<const Vector> create_const(
        std::shared_ptr<const Executor> exec, mpi::communicator comm,
        std::unique_ptr<const local_vector_type> local_vector);
 
protected:
    Vector(std::shared_ptr<const Executor> exec, mpi::communicator comm,
           dim<2> global_size, dim<2> local_size, size_type stride);
 
    explicit Vector(std::shared_ptr<const Executor> exec,
                    mpi::communicator comm, dim<2> global_size = {},
                    dim<2> local_size = {});
 
    Vector(std::shared_ptr<const Executor> exec, mpi::communicator comm,
           dim<2> global_size, std::unique_ptr<local_vector_type> local_vector);
 
    Vector(std::shared_ptr<const Executor> exec, mpi::communicator comm,
           std::unique_ptr<local_vector_type> local_vector);
 
    void resize(dim<2> global_size, dim<2> local_size);
 
    template <typename LocalIndexType, typename GlobalIndexType>
    void read_distributed_impl(
        const device_matrix_data<ValueType, GlobalIndexType>& data,
        const Partition<LocalIndexType, GlobalIndexType>* partition);
 
    void apply_impl(const LinOp*, LinOp*) const override;
 
    void apply_impl(const LinOp*, const LinOp*, const LinOp*,
                    LinOp*) const override;
 
    virtual std::unique_ptr<Vector> create_with_same_config() const;
 
    virtual std::unique_ptr<Vector> create_with_type_of_impl(
        std::shared_ptr<const Executor> exec, const dim<2>& global_size,
        const dim<2>& local_size, size_type stride) const;
 
    virtual std::unique_ptr<Vector> create_submatrix_impl(local_span rows,
                                                          local_span columns,
                                                          dim<2> global_size);
 
private:
    local_vector_type local_;
    ::gko::detail::DenseCache<ValueType> host_reduction_buffer_;
    ::gko::detail::DenseCache<remove_complex<ValueType>> host_norm_buffer_;
};
 
 
}  // namespace distributed
}  // namespace experimental
 
 
namespace detail {
 
 
template <typename TargetType>
struct conversion_target_helper;
 
 
template <typename ValueType>
struct conversion_target_helper<experimental::distributed::Vector<ValueType>> {
    using target_type = experimental::distributed::Vector<ValueType>;
    using source_type =
        experimental::distributed::Vector<previous_precision<ValueType>>;
 
    static std::unique_ptr<target_type> create_empty(const source_type* source)
    {
        return target_type::create(source->get_executor(),
                                   source->get_communicator());
    }
 
    // Allow to create_empty of the same type
    // For distributed case, next<next<V>> will be V in the candidate list.
    // TODO: decide to whether to add this or add condition to the list
    static std::unique_ptr<target_type> create_empty(const target_type* source)
    {
        return target_type::create(source->get_executor(),
                                   source->get_communicator());
    }
 
#if GINKGO_ENABLE_HALF || GINKGO_ENABLE_BFLOAT16
    using snd_source_type =
        experimental::distributed::Vector<previous_precision<ValueType, 2>>;
 
    static std::unique_ptr<target_type> create_empty(
        const snd_source_type* source)
    {
        return target_type::create(source->get_executor(),
                                   source->get_communicator());
    }
#endif
#if GINKGO_ENABLE_HALF && GINKGO_ENABLE_BFLOAT16
    using trd_source_type =
        experimental::distributed::Vector<previous_precision<ValueType, 3>>;
 
    static std::unique_ptr<target_type> create_empty(
        const trd_source_type* source)
    {
        return target_type::create(source->get_executor(),
                                   source->get_communicator());
    }
#endif
};
 
 
}  // namespace detail
}  // namespace gko
 
 
#endif  // GINKGO_BUILD_MPI
 
 
#endif  // GKO_PUBLIC_CORE_DISTRIBUTED_VECTOR_HPP_