batch_jacobi.hpp

// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_PRECONDITIONER_BATCH_JACOBI_HPP_
#define GKO_PUBLIC_CORE_PRECONDITIONER_BATCH_JACOBI_HPP_
 
 
#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/exception_helpers.hpp>
#include <ginkgo/core/base/lin_op.hpp>
#include <ginkgo/core/base/types.hpp>
#include <ginkgo/core/matrix/batch_csr.hpp>
#include <ginkgo/core/matrix/batch_dense.hpp>
#include <ginkgo/core/matrix/csr.hpp>
 
 
namespace gko {
namespace batch {
namespace preconditioner {
 

template <typename ValueType = default_precision, typename IndexType = int32>
class Jacobi final : public EnableBatchLinOp<Jacobi<ValueType, IndexType>> {
    friend class EnableBatchLinOp<Jacobi>;
    friend class EnablePolymorphicObject<Jacobi, BatchLinOp>;
    GKO_ASSERT_SUPPORTED_VALUE_AND_INDEX_TYPE;
 
public:
    using EnableBatchLinOp<Jacobi>::convert_to;
    using EnableBatchLinOp<Jacobi>::move_to;
    using value_type = ValueType;
    using index_type = IndexType;
    using matrix_type = batch::matrix::Csr<ValueType, IndexType>;

    const index_type* get_const_block_pointers() const noexcept
    {
        return block_pointers_.get_const_data();
    }

    const index_type* get_const_map_block_to_row() const noexcept
    {
        return map_block_to_row_.get_const_data();
    }

    const index_type* get_const_blocks_cumulative_offsets() const noexcept
    {
        return blocks_cumulative_offsets_.get_const_data();
    }

    uint32 get_max_block_size() const noexcept
    {
        return parameters_.max_block_size;
    }

    size_type get_num_blocks() const noexcept { return num_blocks_; }

    const value_type* get_const_blocks() const noexcept
    {
        return blocks_.get_const_data();
    }

    size_type get_num_stored_elements() const noexcept
    {
        return blocks_.get_size();
    }
 
    GKO_CREATE_FACTORY_PARAMETERS(parameters, Factory)
    {
        uint32 GKO_FACTORY_PARAMETER_SCALAR(max_block_size, 8u);

        gko::array<index_type> GKO_FACTORY_PARAMETER_VECTOR(block_pointers,
                                                            nullptr);
    };
    GKO_ENABLE_BATCH_LIN_OP_FACTORY(Jacobi, parameters, Factory);
    GKO_ENABLE_BUILD_METHOD(Factory);
 
private:
    explicit Jacobi(std::shared_ptr<const Executor> exec);
 
    explicit Jacobi(const Factory* factory,
                    std::shared_ptr<const BatchLinOp> system_matrix);
 
    void generate_precond(const BatchLinOp* const system_matrix);
 
    size_type compute_storage_space(const size_type num_batch) const noexcept;
 
    void detect_blocks(
        const gko::matrix::Csr<ValueType, IndexType>* system_matrix);
 
    array<index_type> block_pointers_;
    size_type num_blocks_;
    array<value_type> blocks_;
    array<index_type> map_block_to_row_;
    array<index_type> blocks_cumulative_offsets_;
};
 
 
}  // namespace preconditioner
}  // namespace batch
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_PRECONDITIONER_BATCH_JACOBI_HPP_