jacobi.hpp

/*******************************<GINKGO LICENSE>******************************
Copyright (c) 2017-2019, the Ginkgo authors
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions
are met:

1. Redistributions of source code must retain the above copyright
notice, this list of conditions and the following disclaimer.

2. Redistributions in binary form must reproduce the above copyright
notice, this list of conditions and the following disclaimer in the
documentation and/or other materials provided with the distribution.

3. Neither the name of the copyright holder nor the names of its
contributors may be used to endorse or promote products derived from
this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED
TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A
PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
******************************<GINKGO LICENSE>*******************************/

#ifndef GKO_CORE_PRECONDITIONER_JACOBI_HPP_
#define GKO_CORE_PRECONDITIONER_JACOBI_HPP_


#include <ginkgo/core/base/array.hpp>
#include <ginkgo/core/base/lin_op.hpp>
#include <ginkgo/core/matrix/dense.hpp>


namespace gko {
namespace preconditioner {


// TODO: replace this with a custom accessor
template <typename IndexType>
struct block_interleaved_storage_scheme {
    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> {
    friend class EnableLinOp<Jacobi>;
    friend class EnablePolymorphicObject<Jacobi, LinOp>;

public:
    using EnableLinOp<Jacobi>::convert_to;
    using EnableLinOp<Jacobi>::move_to;
    using value_type = ValueType;
    using index_type = IndexType;
    using mat_data = matrix_data<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_num_elems();
    }

    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;

    GKO_CREATE_FACTORY_PARAMETERS(parameters, Factory)
    {
        uint32 GKO_FACTORY_PARAMETER(max_block_size, 32u);

        gko::Array<index_type> GKO_FACTORY_PARAMETER(block_pointers, nullptr);

    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_num_elems() > 0},
                  block_wise{block_wise_opt}
            {}

            storage_optimization_type(
                Array<precision_reduction> &&block_wise_opt)
                : is_block_wise{block_wise_opt.get_num_elems() > 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(
            storage_optimization, precision_reduction(0, 0));

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

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(),
                              transpose(system_matrix->get_size())),
          parameters_{factory->get_parameters()},
          storage_scheme_{compute_storage_scheme(parameters_.max_block_size)},
          num_blocks_{parameters_.block_pointers.get_num_elems() - 1},
          blocks_(factory->get_executor(),
                  storage_scheme_.compute_storage_space(
                      parameters_.block_pointers.get_num_elems() - 1)),
          conditioning_(factory->get_executor())
    {
        if (parameters_.max_block_size >= 32 ||
            parameters_.max_block_size < 1) {
            GKO_NOT_SUPPORTED(this);
        }
        parameters_.block_pointers.set_executor(this->get_executor());
        parameters_.storage_optimization.block_wise.set_executor(
            this->get_executor());
        this->generate(lend(system_matrix));
    }

    static constexpr size_type max_block_stride_ = 32;

    static block_interleaved_storage_scheme<index_type> compute_storage_scheme(
        uint32 max_block_size) noexcept
    {
        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);

    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_CORE_PRECONDITIONER_JACOBI_HPP_