batch_lin_op.hpp

// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_BASE_BATCH_LIN_OP_HPP_
#define GKO_PUBLIC_CORE_BASE_BATCH_LIN_OP_HPP_
 
 
#include <memory>
#include <type_traits>
#include <utility>
 
#include <ginkgo/core/base/abstract_factory.hpp>
#include <ginkgo/core/base/batch_multi_vector.hpp>
#include <ginkgo/core/base/dim.hpp>
#include <ginkgo/core/base/exception_helpers.hpp>
#include <ginkgo/core/base/math.hpp>
#include <ginkgo/core/base/matrix_assembly_data.hpp>
#include <ginkgo/core/base/matrix_data.hpp>
#include <ginkgo/core/base/polymorphic_object.hpp>
#include <ginkgo/core/base/types.hpp>
#include <ginkgo/core/base/utils.hpp>
#include <ginkgo/core/log/logger.hpp>
 
 
namespace gko {
namespace batch {
 
 
class BatchLinOp : public EnableAbstractPolymorphicObject<BatchLinOp> {
public:
    size_type get_num_batch_items() const noexcept
    {
        return get_size().get_num_batch_items();
    }
 
    dim<2> get_common_size() const { return get_size().get_common_size(); }
 
    const batch_dim<2>& get_size() const noexcept { return size_; }
 
    template <typename ValueType>
    void validate_application_parameters(const MultiVector<ValueType>* b,
                                         MultiVector<ValueType>* x) const
    {
        GKO_ASSERT_EQ(b->get_num_batch_items(), this->get_num_batch_items());
        GKO_ASSERT_EQ(this->get_num_batch_items(), x->get_num_batch_items());
 
        GKO_ASSERT_CONFORMANT(this->get_common_size(), b->get_common_size());
        GKO_ASSERT_EQUAL_ROWS(this->get_common_size(), x->get_common_size());
        GKO_ASSERT_EQUAL_COLS(b->get_common_size(), x->get_common_size());
    }
 
    template <typename ValueType>
    void validate_application_parameters(const MultiVector<ValueType>* alpha,
                                         const MultiVector<ValueType>* b,
                                         const MultiVector<ValueType>* beta,
                                         MultiVector<ValueType>* x) const
    {
        GKO_ASSERT_EQ(b->get_num_batch_items(), this->get_num_batch_items());
        GKO_ASSERT_EQ(this->get_num_batch_items(), x->get_num_batch_items());
 
        GKO_ASSERT_CONFORMANT(this->get_common_size(), b->get_common_size());
        GKO_ASSERT_EQUAL_ROWS(this->get_common_size(), x->get_common_size());
        GKO_ASSERT_EQUAL_COLS(b->get_common_size(), x->get_common_size());
        GKO_ASSERT_EQUAL_DIMENSIONS(alpha->get_common_size(),
                                    gko::dim<2>(1, 1));
        GKO_ASSERT_EQUAL_DIMENSIONS(beta->get_common_size(), gko::dim<2>(1, 1));
    }
 
protected:
    void set_size(const batch_dim<2>& size) { size_ = size; }
 
    explicit BatchLinOp(std::shared_ptr<const Executor> exec,
                        const batch_dim<2>& batch_size)
        : EnableAbstractPolymorphicObject<BatchLinOp>(exec), size_{batch_size}
    {}
 
    explicit BatchLinOp(std::shared_ptr<const Executor> exec,
                        const size_type num_batch_items = 0,
                        const dim<2>& common_size = dim<2>{})
        : BatchLinOp{std::move(exec),
                     num_batch_items > 0
                         ? batch_dim<2>(num_batch_items, common_size)
                         : batch_dim<2>{}}
    {}
 
private:
    batch_dim<2> size_{};
};
 
 
class BatchLinOpFactory
    : public AbstractFactory<BatchLinOp, std::shared_ptr<const BatchLinOp>> {
public:
    using AbstractFactory<BatchLinOp,
                          std::shared_ptr<const BatchLinOp>>::AbstractFactory;
 
    std::unique_ptr<BatchLinOp> generate(
        std::shared_ptr<const BatchLinOp> input) const
    {
        this->template log<
            gko::log::Logger::batch_linop_factory_generate_started>(
            this, input.get());
        const auto exec = this->get_executor();
        std::unique_ptr<BatchLinOp> generated;
        if (input->get_executor() == exec) {
            generated = this->AbstractFactory::generate(input);
        } else {
            generated =
                this->AbstractFactory::generate(gko::clone(exec, input));
        }
        this->template log<
            gko::log::Logger::batch_linop_factory_generate_completed>(
            this, input.get(), generated.get());
        return generated;
    }
};
 
 
template <typename ConcreteBatchLinOp, typename PolymorphicBase = BatchLinOp>
class EnableBatchLinOp
    : public EnablePolymorphicObject<ConcreteBatchLinOp, PolymorphicBase>,
      public EnablePolymorphicAssignment<ConcreteBatchLinOp> {
public:
    using EnablePolymorphicObject<ConcreteBatchLinOp,
                                  PolymorphicBase>::EnablePolymorphicObject;
};
 
 
template <typename ConcreteFactory, typename ConcreteBatchLinOp,
          typename ParametersType, typename PolymorphicBase = BatchLinOpFactory>
using EnableDefaultBatchLinOpFactory =
    EnableDefaultFactory<ConcreteFactory, ConcreteBatchLinOp, ParametersType,
                         PolymorphicBase>;
 
 
#define GKO_ENABLE_BATCH_LIN_OP_FACTORY(_batch_lin_op, _parameters_name,     \
                                        _factory_name)                       \
public:                                                                      \
    const _parameters_name##_type& get_##_parameters_name() const            \
    {                                                                        \
        return _parameters_name##_;                                          \
    }                                                                        \
                                                                             \
    class _factory_name                                                      \
        : public ::gko::batch::EnableDefaultBatchLinOpFactory<               \
              _factory_name, _batch_lin_op, _parameters_name##_type> {       \
        friend class ::gko::EnablePolymorphicObject<                         \
            _factory_name, ::gko::batch::BatchLinOpFactory>;                 \
        friend class ::gko::enable_parameters_type<_parameters_name##_type,  \
                                                   _factory_name>;           \
        explicit _factory_name(std::shared_ptr<const ::gko::Executor> exec)  \
            : ::gko::batch::EnableDefaultBatchLinOpFactory<                  \
                  _factory_name, _batch_lin_op, _parameters_name##_type>(    \
                  std::move(exec))                                           \
        {}                                                                   \
        explicit _factory_name(std::shared_ptr<const ::gko::Executor> exec,  \
                               const _parameters_name##_type& parameters)    \
            : ::gko::batch::EnableDefaultBatchLinOpFactory<                  \
                  _factory_name, _batch_lin_op, _parameters_name##_type>(    \
                  std::move(exec), parameters)                               \
        {}                                                                   \
    };                                                                       \
    friend ::gko::batch::EnableDefaultBatchLinOpFactory<                     \
        _factory_name, _batch_lin_op, _parameters_name##_type>;              \
                                                                             \
                                                                             \
private:                                                                     \
    _parameters_name##_type _parameters_name##_;                             \
                                                                             \
public:                                                                      \
    static_assert(true,                                                      \
                  "This assert is used to counter the false positive extra " \
                  "semi-colon warnings")
 
 
}  // namespace batch
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_BASE_BATCH_LIN_OP_HPP_