jacobi.hpp

// SPDX-FileCopyrightText: 2017 - 2025 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_PRECONDITIONER_JACOBI_HPP_
#define GKO_PUBLIC_CORE_PRECONDITIONER_JACOBI_HPP_
 
 
#include <ginkgo/core/base/array.hpp>
#include <ginkgo/core/base/lin_op.hpp>
#include <ginkgo/core/config/config.hpp>
#include <ginkgo/core/config/registry.hpp>
#include <ginkgo/core/matrix/csr.hpp>
#include <ginkgo/core/matrix/dense.hpp>
#include <ginkgo/core/matrix/diagonal.hpp>
 
 
namespace gko {
namespace preconditioner {
 
 
// TODO: replace this with a custom accessor
template <typename IndexType>
struct block_interleaved_storage_scheme {
    block_interleaved_storage_scheme() = default;
 
    block_interleaved_storage_scheme(IndexType block_offset,
                                     IndexType group_offset, uint32 group_power)
        : block_offset{block_offset},
          group_offset{group_offset},
          group_power{group_power}
    {}

    IndexType block_offset;

    IndexType group_offset;

    uint32 group_power;

    GKO_ATTRIBUTES IndexType get_group_size() const noexcept
    {
        return one<IndexType>() << group_power;
    }

    GKO_ATTRIBUTES size_type
    compute_storage_space(size_type num_blocks) const noexcept
    {
        return (num_blocks + 1 == size_type{0})
                   ? size_type{0}
                   : ceildiv(num_blocks, this->get_group_size()) * group_offset;
    }

    GKO_ATTRIBUTES IndexType get_group_offset(IndexType block_id) const noexcept
    {
        return group_offset * (block_id >> group_power);
    }

    GKO_ATTRIBUTES IndexType get_block_offset(IndexType block_id) const noexcept
    {
        return block_offset * (block_id & (this->get_group_size() - 1));
    }

    GKO_ATTRIBUTES IndexType
    get_global_block_offset(IndexType block_id) const noexcept
    {
        return this->get_group_offset(block_id) +
               this->get_block_offset(block_id);
    }

    GKO_ATTRIBUTES IndexType get_stride() const noexcept
    {
        return block_offset << group_power;
    }
};
 

template <typename ValueType = default_precision, typename IndexType = int32>
class Jacobi : public EnableLinOp<Jacobi<ValueType, IndexType>>,
               public ConvertibleTo<matrix::Dense<ValueType>>,
               public WritableToMatrixData<ValueType, IndexType>,
               public Transposable {
    friend class EnableLinOp<Jacobi>;
    friend class EnablePolymorphicObject<Jacobi, LinOp>;
    GKO_ASSERT_SUPPORTED_VALUE_AND_INDEX_TYPE;
 
public:
    using EnableLinOp<Jacobi>::convert_to;
    using EnableLinOp<Jacobi>::move_to;
    using ConvertibleTo<matrix::Dense<ValueType>>::convert_to;
    using ConvertibleTo<matrix::Dense<ValueType>>::move_to;
    using value_type = ValueType;
    using index_type = IndexType;
    using mat_data = matrix_data<ValueType, IndexType>;
    using transposed_type = Jacobi<ValueType, IndexType>;

    size_type get_num_blocks() const noexcept { return num_blocks_; }

    const block_interleaved_storage_scheme<index_type>& get_storage_scheme()
        const noexcept
    {
        return storage_scheme_;
    }

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

    const remove_complex<value_type>* get_conditioning() const noexcept
    {
        return conditioning_.get_const_data();
    }

    size_type get_num_stored_elements() const noexcept
    {
        return blocks_.get_size();
    }
 
    void convert_to(matrix::Dense<value_type>* result) const override;
 
    void move_to(matrix::Dense<value_type>* result) override;
 
    void write(mat_data& data) const override;
 
    std::unique_ptr<LinOp> transpose() const override;
 
    std::unique_ptr<LinOp> conj_transpose() const override;

    Jacobi& operator=(const Jacobi& other);

    Jacobi& operator=(Jacobi&& other);

    Jacobi(const Jacobi& other);

    Jacobi(Jacobi&& other);
 
    GKO_CREATE_FACTORY_PARAMETERS(parameters, Factory)
    {
        uint32 GKO_FACTORY_PARAMETER_SCALAR(max_block_size, 32u);

        uint32 GKO_FACTORY_PARAMETER_SCALAR(max_block_stride, 0u);

        bool GKO_FACTORY_PARAMETER_SCALAR(skip_sorting, false);

        gko::array<index_type> GKO_FACTORY_PARAMETER_VECTOR(block_pointers,
                                                            nullptr);

        bool GKO_FACTORY_PARAMETER_SCALAR(aggregate_l1, false);
 
    private:
        // See documentation of storage_optimization parameter for details about
        // this class
        struct storage_optimization_type {
            storage_optimization_type(precision_reduction p)
                : is_block_wise{false}, of_all_blocks{p}
            {}
 
            storage_optimization_type(
                const array<precision_reduction>& block_wise_opt)
                : is_block_wise{block_wise_opt.get_size() > 0},
                  block_wise{block_wise_opt}
            {}
 
            storage_optimization_type(
                array<precision_reduction>&& block_wise_opt)
                : is_block_wise{block_wise_opt.get_size() > 0},
                  block_wise{std::move(block_wise_opt)}
            {}
 
            operator precision_reduction() { return of_all_blocks; }
 
            bool is_block_wise;
            precision_reduction of_all_blocks;
            gko::array<precision_reduction> block_wise;
        };
 
    public:
        storage_optimization_type GKO_FACTORY_PARAMETER_VECTOR(
            storage_optimization, precision_reduction(0, 0));

        remove_complex<value_type> GKO_FACTORY_PARAMETER_SCALAR(
            accuracy, static_cast<remove_complex<value_type>>(1e-1));
    };
    GKO_ENABLE_LIN_OP_FACTORY(Jacobi, parameters, Factory);
    GKO_ENABLE_BUILD_METHOD(Factory);

    static parameters_type parse(
        const config::pnode& config, const config::registry& context,
        const config::type_descriptor& td_for_child =
            config::make_type_descriptor<ValueType, IndexType>());
 
protected:
    explicit Jacobi(std::shared_ptr<const Executor> exec)
        : EnableLinOp<Jacobi>(exec),
          num_blocks_{},
          blocks_(exec),
          conditioning_(exec)
    {
        parameters_.block_pointers.set_executor(exec);
        parameters_.storage_optimization.block_wise.set_executor(exec);
    }

    explicit Jacobi(const Factory* factory,
                    std::shared_ptr<const LinOp> system_matrix)
        : EnableLinOp<Jacobi>(factory->get_executor(),
                              gko::transpose(system_matrix->get_size())),
          parameters_{factory->get_parameters()},
          storage_scheme_{this->compute_storage_scheme(
              parameters_.max_block_size, parameters_.max_block_stride)},
          num_blocks_{parameters_.block_pointers.get_size() - 1},
          blocks_(factory->get_executor(),
                  storage_scheme_.compute_storage_space(
                      parameters_.block_pointers.get_size() - 1)),
          conditioning_(factory->get_executor())
    {
        parameters_.block_pointers.set_executor(this->get_executor());
        parameters_.storage_optimization.block_wise.set_executor(
            this->get_executor());
        this->generate(system_matrix.get(), parameters_.skip_sorting);
    }

    block_interleaved_storage_scheme<index_type> compute_storage_scheme(
        uint32 max_block_size, uint32 param_max_block_stride)
    {
        uint32 default_block_stride = 32;
        // If the executor is hip, the warp size is 32 or 64
        if (auto hip_exec = std::dynamic_pointer_cast<const gko::HipExecutor>(
                this->get_executor())) {
            default_block_stride = hip_exec->get_warp_size();
        }
        uint32 max_block_stride = default_block_stride;
        if (param_max_block_stride != 0) {
            // if parameter max_block_stride is not zero, set max_block_stride =
            // param_max_block_stride
            max_block_stride = param_max_block_stride;
            if (this->get_executor() != this->get_executor()->get_master() &&
                max_block_stride != default_block_stride) {
                // only support the default value on the gpu device
                GKO_NOT_SUPPORTED(this);
            }
        }
        if (parameters_.max_block_size > max_block_stride ||
            parameters_.max_block_size < 1) {
            GKO_NOT_SUPPORTED(this);
        }
        const auto group_size = static_cast<uint32>(
            max_block_stride / get_superior_power(uint32{2}, max_block_size));
        const auto block_offset = max_block_size;
        const auto block_stride = group_size * block_offset;
        const auto group_offset = max_block_size * block_stride;
        return {static_cast<index_type>(block_offset),
                static_cast<index_type>(group_offset),
                get_significant_bit(group_size)};
    }

    void generate(const LinOp* system_matrix, bool skip_sorting);

    void detect_blocks(const matrix::Csr<ValueType, IndexType>* system_matrix);
 
    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:
    block_interleaved_storage_scheme<index_type> storage_scheme_{};
    size_type num_blocks_;
    array<value_type> blocks_;
    array<remove_complex<value_type>> conditioning_;
};
 
 
}  // namespace preconditioner
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_PRECONDITIONER_JACOBI_HPP_