fbcsr.hpp

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

template <typename IndexType>
inline IndexType get_num_blocks(const int block_size, const IndexType size)
{
    GKO_ASSERT_BLOCK_SIZE_CONFORMANT(size, block_size);
    return size / block_size;
}
 
 
}  // namespace detail
 

template <typename ValueType = default_precision, typename IndexType = int32>
class Fbcsr
    : public EnableLinOp<Fbcsr<ValueType, IndexType>>,
      public ConvertibleTo<Fbcsr<next_precision<ValueType>, IndexType>>,
#if GINKGO_ENABLE_HALF || GINKGO_ENABLE_BFLOAT16
      public ConvertibleTo<Fbcsr<next_precision<ValueType, 2>, IndexType>>,
#endif
#if GINKGO_ENABLE_HALF && GINKGO_ENABLE_BFLOAT16
      public ConvertibleTo<Fbcsr<next_precision<ValueType, 3>, IndexType>>,
#endif
      public ConvertibleTo<Dense<ValueType>>,
      public ConvertibleTo<Csr<ValueType, IndexType>>,
      public ConvertibleTo<SparsityCsr<ValueType, IndexType>>,
      public DiagonalExtractable<ValueType>,
      public ReadableFromMatrixData<ValueType, IndexType>,
      public WritableToMatrixData<ValueType, IndexType>,
      public Transposable,
      public EnableAbsoluteComputation<
          remove_complex<Fbcsr<ValueType, IndexType>>> {
    friend class EnablePolymorphicObject<Fbcsr, LinOp>;
    friend class Csr<ValueType, IndexType>;
    friend class Dense<ValueType>;
    friend class SparsityCsr<ValueType, IndexType>;
    friend class FbcsrBuilder<ValueType, IndexType>;
    friend class Fbcsr<to_complex<ValueType>, IndexType>;
    GKO_ASSERT_SUPPORTED_VALUE_AND_INDEX_TYPE;
 
public:
    using value_type = ValueType;
    using index_type = IndexType;
    using transposed_type = Fbcsr<ValueType, IndexType>;
    using mat_data = matrix_data<ValueType, IndexType>;
    using device_mat_data = device_matrix_data<ValueType, IndexType>;
    using absolute_type = remove_complex<Fbcsr>;

    using EnableLinOp<Fbcsr<ValueType, IndexType>>::move_to;

    using EnableLinOp<Fbcsr<ValueType, IndexType>>::convert_to;
 
    using ConvertibleTo<
        Fbcsr<next_precision<ValueType>, IndexType>>::convert_to;
    using ConvertibleTo<Fbcsr<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 ConvertibleTo<SparsityCsr<ValueType, IndexType>>::convert_to;
    using ConvertibleTo<SparsityCsr<ValueType, IndexType>>::move_to;
 
    friend class Fbcsr<previous_precision<ValueType>, IndexType>;
 
    void convert_to(
        Fbcsr<next_precision<ValueType>, IndexType>* result) const override;
 
    void move_to(Fbcsr<next_precision<ValueType>, IndexType>* result) override;
 
#if GINKGO_ENABLE_HALF || GINKGO_ENABLE_BFLOAT16
    friend class Fbcsr<previous_precision<ValueType, 2>, IndexType>;
    using ConvertibleTo<
        Fbcsr<next_precision<ValueType, 2>, IndexType>>::convert_to;
    using ConvertibleTo<
        Fbcsr<next_precision<ValueType, 2>, IndexType>>::move_to;
 
    void convert_to(
        Fbcsr<next_precision<ValueType, 2>, IndexType>* result) const override;
 
    void move_to(
        Fbcsr<next_precision<ValueType, 2>, IndexType>* result) override;
#endif
 
#if GINKGO_ENABLE_HALF && GINKGO_ENABLE_BFLOAT16
    friend class Fbcsr<previous_precision<ValueType, 3>, IndexType>;
    using ConvertibleTo<
        Fbcsr<next_precision<ValueType, 3>, IndexType>>::convert_to;
    using ConvertibleTo<
        Fbcsr<next_precision<ValueType, 3>, IndexType>>::move_to;
 
    void convert_to(
        Fbcsr<next_precision<ValueType, 3>, IndexType>* result) const override;
 
    void move_to(
        Fbcsr<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>* result) const override;
 
    void move_to(Csr<ValueType, IndexType>* result) override;

    void convert_to(SparsityCsr<ValueType, IndexType>* result) const override;
 
    void move_to(SparsityCsr<ValueType, IndexType>* result) 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<LinOp> transpose() const override;
 
    std::unique_ptr<LinOp> conj_transpose() 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;

    void sort_by_column_index();

    bool is_sorted_by_column_index() 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_row_ptrs() noexcept { return row_ptrs_.get_data(); }

    const index_type* get_const_row_ptrs() const noexcept
    {
        return row_ptrs_.get_const_data();
    }

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

    size_type get_num_stored_blocks() const noexcept
    {
        return col_idxs_.get_size();
    }

    int get_block_size() const noexcept { return bs_; }

    index_type get_num_block_rows() const noexcept
    {
        return this->get_size()[0] / bs_;
    }

    index_type get_num_block_cols() const noexcept
    {
        return this->get_size()[1] / bs_;
    }

    static std::unique_ptr<Fbcsr> create(std::shared_ptr<const Executor> exec,
                                         int block_size = 1);

    static std::unique_ptr<Fbcsr> create(std::shared_ptr<const Executor> exec,
                                         const dim<2>& size,
                                         size_type num_nonzeros,
                                         int block_size);

    static std::unique_ptr<Fbcsr> create(std::shared_ptr<const Executor> exec,
                                         const dim<2>& size, int block_size,
                                         array<value_type> values,
                                         array<index_type> col_idxs,
                                         array<index_type> row_ptrs);

    template <typename InputValueType, typename InputColumnIndexType,
              typename InputRowPtrType>
    GKO_DEPRECATED(
        "explicitly construct the gko::array argument instead of passing "
        "initializer lists")
    static std::unique_ptr<Fbcsr> create(
        std::shared_ptr<const Executor> exec, const dim<2>& size,
        int block_size, std::initializer_list<InputValueType> values,
        std::initializer_list<InputColumnIndexType> col_idxs,
        std::initializer_list<InputRowPtrType> row_ptrs)
    {
        return create(exec, size, block_size,
                      array<value_type>{exec, std::move(values)},
                      array<index_type>{exec, std::move(col_idxs)},
                      array<index_type>{exec, std::move(row_ptrs)});
    }

    static std::unique_ptr<const Fbcsr> create_const(
        std::shared_ptr<const Executor> exec, const dim<2>& size, int blocksize,
        gko::detail::const_array_view<ValueType>&& values,
        gko::detail::const_array_view<IndexType>&& col_idxs,
        gko::detail::const_array_view<IndexType>&& row_ptrs);

    Fbcsr& operator=(const Fbcsr&);

    Fbcsr& operator=(Fbcsr&&);

    Fbcsr(const Fbcsr&);

    Fbcsr(Fbcsr&&);
 
protected:
    Fbcsr(std::shared_ptr<const Executor> exec, int block_size = 1);
 
    Fbcsr(std::shared_ptr<const Executor> exec, const dim<2>& size,
          size_type num_nonzeros, int block_size);
 
    Fbcsr(std::shared_ptr<const Executor> exec, const dim<2>& size,
          int block_size, array<value_type> values, array<index_type> col_idxs,
          array<index_type> row_ptrs);
 
    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:
    int bs_;                      
    array<value_type> values_;    
    array<index_type> col_idxs_;  
    array<index_type> row_ptrs_;  
};
 
 
}  // namespace matrix
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_MATRIX_FBCSR_HPP_