bias_batch_mean_2d_expr.hpp

//=======================================================================
// Copyright (c) 2014-2023 Baptiste Wicht
// Distributed under the terms of the MIT License.
// (See accompanying file LICENSE or copy at
//  http://opensource.org/licenses/MIT)
//=======================================================================

#pragma once

#include "etl/expr/base_temporary_expr.hpp"

#include "etl/impl/cudnn/bias_batch_mean.hpp"
#include "etl/impl/egblas/bias_batch_sum.hpp"

namespace etl {

template <etl_2d A, bool Mean>
struct bias_batch_mean_2d_expr : base_temporary_expr_un<bias_batch_mean_2d_expr<A, Mean>, A> {
    using value_type = value_t<A>;                           
    using this_type  = bias_batch_mean_2d_expr<A, Mean>;     
    using base_type  = base_temporary_expr_un<this_type, A>; 
    using sub_traits = decay_traits<A>;                      

    static constexpr auto storage_order = sub_traits::storage_order; 

    static constexpr bool gpu_computable = (!Mean && cudnn_enabled && is_floating<A>)
            || (impl::egblas::has_sbias_batch_sum && all_row_major<A> && all_single_precision<A>)
            || (impl::egblas::has_dbias_batch_sum && all_row_major<A> && all_double_precision<A>);

    explicit bias_batch_mean_2d_expr(A a) : base_type(a) {
        //Nothing else to init
    }

    template <etl_1d C>
    static void check([[maybe_unused]] const A& a, [[maybe_unused]] const C& c) {
        if constexpr (all_fast<A, C>) {
            static_assert(etl::dim<1, A>() == etl::dim<0, C>(), "Invalid dimensions for bias_batch_mean_2d");
        } else {
            cpp_assert(etl::dim<1>(a) == etl::dim<0>(c), "Invalid dimensions for bias_batch_mean_2d");
        }
    }

    // Assignment functions

    template <etl_expr L>
    void assign_to(L&& lhs) const {
        inc_counter("temp:assign");

        auto& a = this->a();

        using T = value_t<A>;

        check(a, lhs);

        if constexpr (!Mean && impl::egblas::has_sbias_batch_sum && all_row_major<A> && all_floating<A, L>) {
            const auto N = etl::dim<0>(a);
            const auto K = etl::dim<1>(a);

            decltype(auto) t1 = smart_forward_gpu(a);
            t1.ensure_gpu_up_to_date();

            lhs.ensure_gpu_allocated();

            impl::egblas::bias_batch_sum(N, K, t1.gpu_memory(), 1, lhs.gpu_memory(), 1);

            lhs.validate_gpu();
            lhs.invalidate_cpu();
        } else if constexpr (Mean && impl::egblas::has_sbias_batch_mean && all_row_major<A> && all_floating<A, L>) {
            const auto N = etl::dim<0>(a);
            const auto K = etl::dim<1>(a);

            decltype(auto) t1 = smart_forward_gpu(a);
            t1.ensure_gpu_up_to_date();

            lhs.ensure_gpu_allocated();

            impl::egblas::bias_batch_mean(N, K, t1.gpu_memory(), 1, lhs.gpu_memory(), 1);

            lhs.validate_gpu();
            lhs.invalidate_cpu();
        } else if constexpr (!Mean && cudnn_enabled && all_floating<A, L>) {
            impl::cudnn::bias_batch_mean_2d(smart_forward_gpu(a), lhs);
        } else {
            const auto N = etl::dim<0>(a);
            const auto K = etl::dim<1>(a);

            standard_evaluator::pre_assign_rhs(a);

            auto batch_fun_k = [&](const size_t first, const size_t last) {
                for (size_t k = first; k < last; ++k) {
                    T mean(0);

                    for (size_t b = 0; b < N; ++b) {
                        mean += a(b, k);
                    }

                    if constexpr (Mean) {
                        lhs(k) = mean / static_cast<T>(N);
                    } else {
                        lhs(k) = mean;
                    }
                }
            };

            engine_dispatch_1d_serial(batch_fun_k, 0, K, 4UL);
        }
    }

    template <etl_expr L>
    void assign_add_to(L&& lhs) const {
        if constexpr (all_floating<A, L> && ((!Mean && cudnn_enabled) || (all_row_major<A, L> && impl::egblas::has_sbias_batch_sum))) {
            std_add_evaluate(*this, lhs);
        } else {
            auto& a = this->a();

            standard_evaluator::pre_assign_rhs(a);

            const auto N = etl::size(a) / etl::size(lhs);
            const auto K = etl::size(lhs);

            using T = value_t<A>;

            check(a, lhs);

            auto batch_fun_k = [&](const size_t first, const size_t last) {
                for (size_t k = first; k < last; ++k) {
                    T mean(0);

                    for (size_t b = 0; b < N; ++b) {
                        mean += a(b, k);
                    }

                    if constexpr (Mean) {
                        lhs(k) += mean / static_cast<T>(N);
                    } else {
                        lhs(k) += mean;
                    }
                }
            };

            engine_dispatch_1d_serial(batch_fun_k, 0, K, 4UL);
        }
    }

    template <etl_expr L>
    void assign_sub_to(L&& lhs) const {
        if constexpr (all_floating<A, L> && ((!Mean && cudnn_enabled) || (all_row_major<A, L> && impl::egblas::has_sbias_batch_sum))) {
            std_sub_evaluate(*this, lhs);
        } else {
            auto& a = this->a();

            standard_evaluator::pre_assign_rhs(a);

            const auto N = etl::size(a) / etl::size(lhs);
            const auto K = etl::size(lhs);

            using T = value_t<A>;

            check(a, lhs);

            auto batch_fun_k = [&](const size_t first, const size_t last) {
                for (size_t k = first; k < last; ++k) {
                    T mean(0);

                    for (size_t b = 0; b < N; ++b) {
                        mean += a(b, k);
                    }

                    if constexpr (Mean) {
                        lhs(k) -= mean / static_cast<T>(N);
                    } else {
                        lhs(k) -= mean;
                    }
                }
            };

            engine_dispatch_1d_serial(batch_fun_k, 0, K, 4UL);
        }
    }

    template <etl_expr L>
    void assign_mul_to(L&& lhs) const {
        if constexpr (all_floating<A, L> && ((!Mean && cudnn_enabled) || (all_row_major<A, L> && impl::egblas::has_sbias_batch_sum))) {
            std_mul_evaluate(*this, lhs);
        } else {
            auto& a = this->a();

            standard_evaluator::pre_assign_rhs(a);

            const auto N = etl::size(a) / etl::size(lhs);
            const auto K = etl::size(lhs);

            using T = value_t<A>;

            check(a, lhs);

            auto batch_fun_k = [&](const size_t first, const size_t last) {
                for (size_t k = first; k < last; ++k) {
                    T mean(0);

                    for (size_t b = 0; b < N; ++b) {
                        mean += a(b, k);
                    }

                    if constexpr (Mean) {
                        lhs(k) *= mean / static_cast<T>(N);
                    } else {
                        lhs(k) *= mean;
                    }
                }
            };

            engine_dispatch_1d_serial(batch_fun_k, 0, K, 4UL);
        }
    }

    template <etl_expr L>
    void assign_div_to(L&& lhs) const {
        if constexpr (all_floating<A, L> && ((!Mean && cudnn_enabled) || (all_row_major<A, L> && impl::egblas::has_sbias_batch_sum))) {
            std_div_evaluate(*this, lhs);
        } else {
            auto& a = this->a();

            standard_evaluator::pre_assign_rhs(a);

            const auto N = etl::size(a) / etl::size(lhs);
            const auto K = etl::size(lhs);

            using T = value_t<A>;

            check(a, lhs);

            auto batch_fun_k = [&](const size_t first, const size_t last) {
                for (size_t k = first; k < last; ++k) {
                    T mean(0);

                    for (size_t b = 0; b < N; ++b) {
                        mean += a(b, k);
                    }

                    if constexpr (Mean) {
                        lhs(k) /= mean / static_cast<T>(N);
                    } else {
                        lhs(k) /= mean;
                    }
                }
            };

            engine_dispatch_1d_serial(batch_fun_k, 0, K, 4UL);
        }
    }

    template <etl_expr L>
    void assign_mod_to(L&& lhs) const {
        if constexpr (all_floating<A, L> && ((!Mean && cudnn_enabled) || (all_row_major<A, L> && impl::egblas::has_sbias_batch_sum))) {
            std_mod_evaluate(*this, lhs);
        } else {
            auto& a = this->a();

            standard_evaluator::pre_assign_rhs(a);

            const auto N = etl::size(a) / etl::size(lhs);
            const auto K = etl::size(lhs);

            using T = value_t<A>;

            check(a, lhs);

            auto batch_fun_k = [&](const size_t first, const size_t last) {
                for (size_t k = first; k < last; ++k) {
                    T mean(0);

                    for (size_t b = 0; b < N; ++b) {
                        mean += a(b, k);
                    }

                    if constexpr (Mean) {
                        lhs(k) %= mean / static_cast<T>(N);
                    } else {
                        lhs(k) %= mean;
                    }
                }
            };

            engine_dispatch_1d_serial(batch_fun_k, 0, K, 4UL);
        }
    }

    friend std::ostream& operator<<(std::ostream& os, const bias_batch_mean_2d_expr& expr) {
        if (Mean) {
            return os << "bias_batch_mean_2d(" << expr._a << ")";
        } else {
            return os << "bias_batch_sum_2d(" << expr._a << ")";
        }
    }
};

template <typename A, bool Mean>
struct etl_traits<etl::bias_batch_mean_2d_expr<A, Mean>> {
    using expr_t     = etl::bias_batch_mean_2d_expr<A, Mean>; 
    using sub_expr_t = std::decay_t<A>;                       
    using sub_traits = etl_traits<sub_expr_t>;                
    using value_type = value_t<A>;                            

    static constexpr bool is_etl         = true;                                 
    static constexpr bool is_transformer = false;                                
    static constexpr bool is_view        = false;                                
    static constexpr bool is_magic_view  = false;                                
    static constexpr bool is_fast        = sub_traits::is_fast;                  
    static constexpr bool is_linear      = false;                                
    static constexpr bool is_thread_safe = true;                                 
    static constexpr bool is_value       = false;                                
    static constexpr bool is_direct      = true;                                 
    static constexpr bool is_generator   = false;                                
    static constexpr bool is_padded      = false;                                
    static constexpr bool is_aligned     = true;                                 
    static constexpr bool is_temporary   = true;                                 
    static constexpr order storage_order = sub_traits::storage_order;            
    static constexpr bool gpu_computable = is_gpu_t<value_type> && cuda_enabled; 

    template <vector_mode_t V>
    static constexpr bool vectorizable = true;

    template <size_t DD>
    static constexpr size_t dim() requires(DD == 0) {
        return decay_traits<A>::template dim<1>();
    }

    static size_t dim(const expr_t& e, [[maybe_unused]] size_t d) {
        cpp_assert(d == 0, "Invalid dimensions access");

        return etl::dim<1>(e._a);
    }

    static size_t size(const expr_t& e) {
        return etl::dim<1>(e._a);
    }

    static constexpr size_t size() {
        return decay_traits<A>::template dim<1>();
    }

    static constexpr size_t dimensions() {
        return 1;
    }

    static constexpr int complexity() noexcept {
        return -1;
    }
};

template <etl_2d E>
bias_batch_mean_2d_expr<detail::build_type<E>, true> bias_batch_mean_2d(const E& value) {
    return bias_batch_mean_2d_expr<detail::build_type<E>, true>{value};
}

template <etl_2d E>
bias_batch_mean_2d_expr<detail::build_type<E>, false> bias_batch_sum_2d(const E& value) {
    return bias_batch_mean_2d_expr<detail::build_type<E>, false>{value};
}

} //end of namespace etl