batch_identity.hpp

// SPDX-FileCopyrightText: 2017 - 2024 The Ginkgo authors
//
// SPDX-License-Identifier: BSD-3-Clause
 
#ifndef GKO_PUBLIC_CORE_MATRIX_BATCH_IDENTITY_HPP_
#define GKO_PUBLIC_CORE_MATRIX_BATCH_IDENTITY_HPP_
 
 
#include <ginkgo/core/base/batch_lin_op.hpp>
#include <ginkgo/core/base/batch_multi_vector.hpp>
#include <ginkgo/core/base/executor.hpp>
#include <ginkgo/core/base/types.hpp>
#include <ginkgo/core/base/utils.hpp>
#include <ginkgo/core/matrix/identity.hpp>
 
 
namespace gko {
namespace batch {
namespace matrix {
 
 
template <typename ValueType = default_precision>
class Identity final : public EnableBatchLinOp<Identity<ValueType>> {
    friend class EnablePolymorphicObject<Identity, BatchLinOp>;
 
public:
    using value_type = ValueType;
    using index_type = int32;
    using unbatch_type = gko::matrix::Identity<ValueType>;
    using absolute_type = remove_complex<Identity>;
    using complex_type = to_complex<Identity>;
 
    Identity* apply(ptr_param<const MultiVector<value_type>> b,
                    ptr_param<MultiVector<value_type>> x);
 
    Identity* apply(ptr_param<const MultiVector<value_type>> alpha,
                    ptr_param<const MultiVector<value_type>> b,
                    ptr_param<const MultiVector<value_type>> beta,
                    ptr_param<MultiVector<value_type>> x);
 
    const Identity* apply(ptr_param<const MultiVector<value_type>> b,
                          ptr_param<MultiVector<value_type>> x) const;
 
    const Identity* apply(ptr_param<const MultiVector<value_type>> alpha,
                          ptr_param<const MultiVector<value_type>> b,
                          ptr_param<const MultiVector<value_type>> beta,
                          ptr_param<MultiVector<value_type>> x) const;
    static std::unique_ptr<Identity> create(
        std::shared_ptr<const Executor> exec,
        const batch_dim<2>& size = batch_dim<2>{});
 
private:
    Identity(std::shared_ptr<const Executor> exec,
             const batch_dim<2>& size = batch_dim<2>{});
 
    void apply_impl(const MultiVector<value_type>* b,
                    MultiVector<value_type>* x) const;
 
    void apply_impl(const MultiVector<value_type>* alpha,
                    const MultiVector<value_type>* b,
                    const MultiVector<value_type>* beta,
                    MultiVector<value_type>* x) const;
};
 
 
}  // namespace matrix
}  // namespace batch
}  // namespace gko
 
 
#endif  // GKO_PUBLIC_CORE_MATRIX_BATCH_IDENTITY_HPP_