diagonal.hpp

// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_MATRIX_DIAGONAL_HPP_
#define GKO_PUBLIC_CORE_MATRIX_DIAGONAL_HPP_
 
 
#include <ginkgo/core/base/array.hpp>
#include <ginkgo/core/base/lin_op.hpp>
 
 
namespace gko {
namespace matrix {
 
 
template <typename ValueType, typename IndexType>
class Csr;
 
template <typename ValueType>
class Dense;
 

template <typename ValueType = default_precision>
class Diagonal
    : public EnableLinOp<Diagonal<ValueType>>,
      public ConvertibleTo<Csr<ValueType, int32>>,
      public ConvertibleTo<Csr<ValueType, int64>>,
      public ConvertibleTo<Diagonal<next_precision<ValueType>>>,
#if GINKGO_ENABLE_HALF || GINKGO_ENABLE_BFLOAT16
      public ConvertibleTo<Diagonal<next_precision<ValueType, 2>>>,
#endif
#if GINKGO_ENABLE_HALF && GINKGO_ENABLE_BFLOAT16
      public ConvertibleTo<Diagonal<next_precision<ValueType, 3>>>,
#endif
      public Transposable,
      public WritableToMatrixData<ValueType, int32>,
      public WritableToMatrixData<ValueType, int64>,
      public ReadableFromMatrixData<ValueType, int32>,
      public ReadableFromMatrixData<ValueType, int64>,
      public EnableAbsoluteComputation<remove_complex<Diagonal<ValueType>>> {
    friend class EnablePolymorphicObject<Diagonal, LinOp>;
    friend class Csr<ValueType, int32>;
    friend class Csr<ValueType, int64>;
    friend class Diagonal<to_complex<ValueType>>;
    GKO_ASSERT_SUPPORTED_VALUE_TYPE;
 
public:
    using EnableLinOp<Diagonal>::convert_to;
    using EnableLinOp<Diagonal>::move_to;
    using ConvertibleTo<Csr<ValueType, int32>>::convert_to;
    using ConvertibleTo<Csr<ValueType, int32>>::move_to;
    using ConvertibleTo<Csr<ValueType, int64>>::convert_to;
    using ConvertibleTo<Csr<ValueType, int64>>::move_to;
    using ConvertibleTo<Diagonal<next_precision<ValueType>>>::convert_to;
    using ConvertibleTo<Diagonal<next_precision<ValueType>>>::move_to;
 
    using value_type = ValueType;
    using index_type = int64;
    using mat_data = matrix_data<ValueType, int64>;
    using mat_data32 = matrix_data<ValueType, int32>;
    using device_mat_data = device_matrix_data<ValueType, int64>;
    using device_mat_data32 = device_matrix_data<ValueType, int32>;
    using absolute_type = remove_complex<Diagonal>;
 
    friend class Diagonal<previous_precision<ValueType>>;
 
    std::unique_ptr<LinOp> transpose() const override;
 
    std::unique_ptr<LinOp> conj_transpose() const override;
 
    void convert_to(Diagonal<next_precision<ValueType>>* result) const override;
 
    void move_to(Diagonal<next_precision<ValueType>>* result) override;
 
#if GINKGO_ENABLE_HALF || GINKGO_ENABLE_BFLOAT16
    friend class Diagonal<previous_precision<ValueType, 2>>;
    using ConvertibleTo<Diagonal<next_precision<ValueType, 2>>>::convert_to;
    using ConvertibleTo<Diagonal<next_precision<ValueType, 2>>>::move_to;
 
    void convert_to(
        Diagonal<next_precision<ValueType, 2>>* result) const override;
 
    void move_to(Diagonal<next_precision<ValueType, 2>>* result) override;
#endif
 
#if GINKGO_ENABLE_HALF && GINKGO_ENABLE_BFLOAT16
    friend class Diagonal<previous_precision<ValueType, 3>>;
    using ConvertibleTo<Diagonal<next_precision<ValueType, 3>>>::convert_to;
    using ConvertibleTo<Diagonal<next_precision<ValueType, 3>>>::move_to;
 
    void convert_to(
        Diagonal<next_precision<ValueType, 3>>* result) const override;
 
    void move_to(Diagonal<next_precision<ValueType, 3>>* result) override;
#endif
 
    void convert_to(Csr<ValueType, int32>* result) const override;
 
    void move_to(Csr<ValueType, int32>* result) override;
 
    void convert_to(Csr<ValueType, int64>* result) const override;
 
    void move_to(Csr<ValueType, int64>* result) 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();
    }

    void rapply(ptr_param<const LinOp> b, ptr_param<LinOp> x) const
    {
        GKO_ASSERT_REVERSE_CONFORMANT(this, b);
        GKO_ASSERT_EQUAL_ROWS(b, x);
        GKO_ASSERT_EQUAL_COLS(this, x);
 
        this->rapply_impl(b.get(), x.get());
    }

    void inverse_apply(ptr_param<const LinOp> b, ptr_param<LinOp> x) const
    {
        GKO_ASSERT_CONFORMANT(this, b);
        GKO_ASSERT_EQUAL_ROWS(b, x);
        GKO_ASSERT_EQUAL_ROWS(this, x);
 
        this->inverse_apply_impl(b.get(), x.get());
    }
 
    void read(const mat_data& data) override;
 
    void read(const mat_data32& data) override;
 
    void read(const device_mat_data& data) override;
 
    void read(const device_mat_data32& data) override;
 
    void read(device_mat_data&& data) override;
 
    void read(device_mat_data32&& data) override;
 
    void write(mat_data& data) const override;
 
    void write(mat_data32& data) const override;

    static std::unique_ptr<Diagonal> create(
        std::shared_ptr<const Executor> exec, size_type size = 0);

    static std::unique_ptr<Diagonal> create(
        std::shared_ptr<const Executor> exec, const size_type 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<Diagonal> create(
        std::shared_ptr<const Executor> exec, const size_type size,
        std::initializer_list<InputValueType> values)
    {
        return create(exec, size, array<value_type>{exec, std::move(values)});
    }

    static std::unique_ptr<const Diagonal> create_const(
        std::shared_ptr<const Executor> exec, size_type size,
        gko::detail::const_array_view<ValueType>&& values);
 
protected:
    Diagonal(std::shared_ptr<const Executor> exec, size_type size = 0);
 
    Diagonal(std::shared_ptr<const Executor> exec, const size_type size,
             array<value_type> values);
 
    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;
 
    void rapply_impl(const LinOp* b, LinOp* x) const;
 
    void inverse_apply_impl(const LinOp* b, LinOp* x) const;
 
private:
    array<value_type> values_;
};
 
 
}  // namespace matrix
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_MATRIX_DIAGONAL_HPP_