bias_add_4d_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 the implementations
#include "etl/impl/std/bias_add.hpp"
#include "etl/impl/vec/bias_add.hpp"
#include "etl/impl/cudnn/bias_add.hpp"
#include "etl/impl/egblas/bias_add.hpp"

namespace etl {

template <etl_4d A, etl_1d B>
struct bias_add_4d_expr : base_temporary_expr_bin<bias_add_4d_expr<A, B>, A, B> {
    using value_type = value_t<A>;                               
    using this_type  = bias_add_4d_expr<A, B>;                   
    using base_type  = base_temporary_expr_bin<this_type, A, B>; 
    using sub_traits = decay_traits<A>;                          

    static constexpr auto storage_order = sub_traits::storage_order; 

    static constexpr bool gpu_computable = cudnn_enabled && all_floating<A, B> && all_homogeneous<A, B>;

    explicit bias_add_4d_expr(A a, B b) : base_type(a, b) {
        //Nothing else to init
    }

    template <etl_4d C>
    static void check([[maybe_unused]] const A& a, [[maybe_unused]] const B& b, [[maybe_unused]] const C& c) {
        if constexpr (all_fast<A, B, C>) {
            static_assert(etl::dim<1, A>() == etl::dim<0, B>(), "Invalid dimensions for bias_add");

            static_assert(etl::dim<0, A>() == etl::dim<0, C>(), "Invalid dimensions for bias_add");
            static_assert(etl::dim<1, A>() == etl::dim<1, C>(), "Invalid dimensions for bias_add");
            static_assert(etl::dim<2, A>() == etl::dim<2, C>(), "Invalid dimensions for bias_add");
            static_assert(etl::dim<3, A>() == etl::dim<3, C>(), "Invalid dimensions for bias_add");
        } else {
            cpp_assert(etl::dim<1>(a) == etl::dim<0>(b), "Invalid dimensions for bias_add");

            cpp_assert(etl::dim<0>(a) == etl::dim<0>(c), "Invalid dimensions for bias_add");
            cpp_assert(etl::dim<1>(a) == etl::dim<1>(c), "Invalid dimensions for bias_add");
            cpp_assert(etl::dim<2>(a) == etl::dim<2>(c), "Invalid dimensions for bias_add");
            cpp_assert(etl::dim<3>(a) == etl::dim<3>(c), "Invalid dimensions for bias_add");
        }
    }

    // Assignment functions

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

        if (this->is_evaluated()) {
            lhs = this->result();
            return;
        }

        auto& a = this->a();
        auto& b = this->b();

        check(a, b, lhs);

        constexpr_select auto impl = select_impl<L>();

        if
            constexpr_select(impl == bias_add_impl::VEC) {
                inc_counter("impl:vec");
                impl::vec::bias_add_4d(smart_forward(a), smart_forward(b), lhs);
            }
        else if
            constexpr_select(impl == bias_add_impl::STD) {
                inc_counter("impl:std");
                impl::standard::bias_add_4d(smart_forward(a), smart_forward(b), lhs);
            }
        else if
            constexpr_select(impl == bias_add_impl::EGBLAS) {
                inc_counter("impl:egblas");

                decltype(auto) e_x = smart_forward_gpu(a);
                decltype(auto) e_b = smart_forward_gpu(b);
                auto& e_y          = lhs;

                e_x.ensure_gpu_up_to_date();
                e_b.ensure_gpu_up_to_date();
                e_y.ensure_gpu_allocated();

                impl::egblas::bias_add_4d(etl::dim<0>(a), etl::dim<1>(a), etl::dim<2>(a), etl::dim<3>(a), e_x.gpu_memory(), 1, e_b.gpu_memory(), 1,
                                          e_y.gpu_memory(), 1);

                e_y.validate_gpu();
                e_y.invalidate_cpu();
            }
        else if
            constexpr_select(impl == bias_add_impl::CUDNN) {
                inc_counter("impl:cudnn");
                impl::cudnn::bias_add_4d(smart_forward_gpu(a), smart_forward_gpu(b), lhs);
            }
        else {
            cpp_unreachable("Invalid bias_add selection");
        }
    }

    template <etl_4d L>
    void assign_add_to(L&& lhs) const {
        std_add_evaluate(*this, lhs);
    }

    template <etl_4d L>
    void assign_sub_to(L&& lhs) const {
        std_sub_evaluate(*this, lhs);
    }

    template <etl_4d L>
    void assign_mul_to(L&& lhs) const {
        std_mul_evaluate(*this, lhs);
    }

    template <etl_4d L>
    void assign_div_to(L&& lhs) const {
        std_div_evaluate(*this, lhs);
    }

    template <etl_4d L>
    void assign_mod_to(L&& lhs) const {
        std_mod_evaluate(*this, lhs);
    }

    friend std::ostream& operator<<(std::ostream& os, const bias_add_4d_expr& expr) {
        return os << "bias_add(" << expr._a << "," << expr._b << ")";
    }

private:
    template <typename C>
    static constexpr etl::bias_add_impl select_default_impl(bool no_gpu) {
        constexpr bool homo           = all_homogeneous<A, B, C>;
        constexpr bool vec_possible   = vec_enabled && vectorize_impl && all_vectorizable<vector_mode, A, B, C> && homo;
        constexpr bool cudnn_possible = cudnn_enabled && all_floating<A, B, C> && homo;

        if (homo && is_single_precision<A> && impl::egblas::has_sbias_add_4d) {
            return etl::bias_add_impl::EGBLAS;
        }

        if (homo && is_double_precision<A> && impl::egblas::has_dbias_add_4d) {
            return etl::bias_add_impl::EGBLAS;
        }

        if (cudnn_possible && !no_gpu) {
            return etl::bias_add_impl::CUDNN;
        }

        if (vec_possible) {
            return etl::bias_add_impl::VEC;
        }

        return etl::bias_add_impl::STD;
    }

#ifdef ETL_MANUAL_SELECT

    template <typename C>
    static etl::bias_add_impl select_impl() {
        auto def = select_default_impl<C>(local_context().cpu);

        if (local_context().bias_add_selector.forced) {
            auto forced = local_context().bias_add_selector.impl;

            switch (forced) {
                // EGBLAS cannot always be used
                case bias_add_impl::EGBLAS:
                    if (!all_homogeneous<
                            A, B,
                            C> || !((is_single_precision<A> && impl::egblas::has_sbias_add_4d) || (is_double_precision<A> && impl::egblas::has_sbias_add_4d))
                        || local_context().cpu) {
                        std::cerr << "Forced selection to EGBLAS bias_add implementation, but not possible for this expression" << std::endl;
                        return def;
                    }

                    return forced;

                //CUDNN cannot always be used
                case bias_add_impl::CUDNN:
                    if (!cudnn_enabled || !all_floating<A, B, C> || !all_homogeneous<A, B, C> || local_context().cpu) {
                        std::cerr << "Forced selection to cUDNN bias_add implementation, but not possible for this expression" << std::endl;
                        return def;
                    }

                    return forced;

                //VEC cannot always be used
                case bias_add_impl::VEC:
                    if (!vec_enabled || !vectorize_impl || !all_vectorizable<vector_mode, A, B, C> || !all_homogeneous<A, B, C>) {
                        std::cerr << "Forced selection to VEC bias_add implementation, but not possible for this expression" << std::endl;
                        return def;
                    }

                    return forced;

                //In other cases, simply use the forced impl
                default:
                    return forced;
            }
        }

        return def;
    }

#else

    template <typename C>
    static constexpr etl::bias_add_impl select_impl() {
        return select_default_impl<C>(false);
    }

#endif
};

template <typename A, typename B>
struct etl_traits<etl::bias_add_4d_expr<A, B>> {
    using expr_t     = etl::bias_add_4d_expr<A, B>; 
    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        = all_fast<A, B>;                       
    static constexpr bool is_linear      = true;                                 
    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 bool gpu_computable = is_gpu_t<value_type> && cuda_enabled; 
    static constexpr order storage_order = sub_traits::storage_order;            

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

    template <size_t DD>
    static constexpr size_t dim() {
        return sub_traits::template dim<DD>();
    }

    static size_t dim(const expr_t& e, size_t d) {
        return sub_traits::dim(e._a, d);
    }

    static size_t size(const expr_t& e) {
        return sub_traits::size(e._a);
    }

    static constexpr size_t size() {
        return sub_traits::size();
    }

    static constexpr size_t dimensions() {
        return 4;
    }

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

template <etl_4d E, etl_1d B>
bias_add_4d_expr<detail::build_type<E>, detail::build_type<B>> bias_add_4d(const E& x, const B& biases) {
    return bias_add_4d_expr<detail::build_type<E>, detail::build_type<B>>{x, biases};
}

} //end of namespace etl