dyn_pool_upsample_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"

//Get the implementations
#include "etl/impl/std/max_pooling_upsample.hpp"
#include "etl/impl/std/avg_pooling_upsample.hpp"
#include "etl/impl/cudnn/pooling_upsample.hpp"

namespace etl {

template <etl_expr A, same_dimensions<A> B, same_dimensions<A> C, bool Max>
struct dyn_pool_upsample_2d_expr : base_temporary_expr_tern<dyn_pool_upsample_2d_expr<A, B, C, Max>, A, B, C> {
    using value_type = value_t<A>;                                   
    using sub_traits = etl::decay_traits<A>;                         
    using this_type  = dyn_pool_upsample_2d_expr<A, B, C, Max>;      
    using base_type  = base_temporary_expr_tern<this_type, A, B, C>; 

    static constexpr auto storage_order = sub_traits::storage_order; 

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

private:
    const size_t c1; 
    const size_t c2; 

    const size_t s1; 
    const size_t s2; 

    const size_t p1; 
    const size_t p2; 

    friend struct etl_traits<dyn_pool_upsample_2d_expr>;

public:
    dyn_pool_upsample_2d_expr(A a, B b, C c, size_t c1, size_t c2, size_t s1, size_t s2, size_t p1, size_t p2) :
            base_type(a, b, c), c1(c1), c2(c2), s1(s1), s2(s2), p1(p1), p2(p2) {
        //Nothing else to init
    }

    template <same_dimensions<A> R>
    void check([[maybe_unused]] const A& a, [[maybe_unused]] const B& b, [[maybe_unused]] const C& c, [[maybe_unused]] const R& result) const {
        static constexpr size_t D = etl::decay_traits<A>::dimensions();

        cpp_assert(etl::size(result) == etl::size(a), "max_pool_upsample_2d:A and R must have the same size");
        cpp_assert(etl::size(b) == etl::size(c), "max_pool_upsample_2d:B and C must have the same size");

        cpp_assert(etl::dim<D - 2>(a) == s1 * (etl::dim<D - 2>(b) - 1) - 2 * p1 + c1, "Invalid pooling dimensions for max_pool_upsample_2d");
        cpp_assert(etl::dim<D - 1>(a) == s2 * (etl::dim<D - 1>(b) - 1) - 2 * p2 + c2, "Invalid pooling dimensions for max_pool_upsample_2d");
    }

    // Assignment functions

    template <typename R>
    static constexpr etl::pool_impl select_default_impl(bool no_gpu) {
        if (cudnn_enabled && all_floating<A, B, C, R> && !no_gpu) {
            return etl::pool_impl::CUDNN;
        }

        return etl::pool_impl::STD;
    }

#ifdef ETL_MANUAL_SELECT

    template <typename R>
    static etl::pool_impl select_impl() {
        if (local_context().pool_selector.forced) {
            auto forced = local_context().pool_selector.impl;

            switch (forced) {
                // CUDNN cannot always be used
                case pool_impl::CUDNN:
                    if (!cudnn_enabled || !all_floating<A, B, C, R> || local_context().cpu) {                                            //COVERAGE_EXCLUDE_LINE
                        std::cerr << "Forced selection to CUDNN pool implementation, but not possible for this expression" << std::endl; //COVERAGE_EXCLUDE_LINE
                        return select_default_impl<R>(local_context().cpu);                                                              //COVERAGE_EXCLUDE_LINE
                    }                                                                                                                    //COVERAGE_EXCLUDE_LINE

                    return forced;

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

        return select_default_impl<R>(local_context().cpu);
    }

#else

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

#endif

    // Assignment functions

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

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

        check(a, b, c, result);

        constexpr_select auto impl = select_impl<R>();

        if constexpr (Max) {
            if
                constexpr_select(impl == pool_impl::STD) {
                    inc_counter("impl:std");
                    impl::standard::max_pool_upsample_2d::apply(smart_forward(a), smart_forward(b), smart_forward(c), result, c1, c2, s1, s2, p1, p2);
                }
            else if
                constexpr_select(impl == pool_impl::CUDNN) {
                    inc_counter("impl:cudnn");
                    impl::cudnn::max_pool_upsample_2d::apply(smart_forward_gpu(a), smart_forward_gpu(b), smart_forward_gpu(c), result, c1, c2, s1, s2, p1, p2);
                }
            else {
                cpp_unreachable("Invalid pool implementation");
            }
        } else {
            if
                constexpr_select(impl == pool_impl::STD) {
                    inc_counter("impl:std");
                    impl::standard::avg_pool_upsample_2d::apply(smart_forward(a), smart_forward(b), smart_forward(c), result, c1, c2, s1, s2, p1, p2);
                }
            else if
                constexpr_select(impl == pool_impl::CUDNN) {
                    inc_counter("impl:cudnn");
                    impl::cudnn::avg_pool_upsample_2d::apply(smart_forward_gpu(a), smart_forward_gpu(b), smart_forward_gpu(c), result, c1, c2, s1, s2, p1, p2);
                }
            else {
                cpp_unreachable("Invalid pool implementation");
            }
        }
    }

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

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

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

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

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

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

template <typename A, typename B, typename C, bool Max>
struct etl_traits<etl::dyn_pool_upsample_2d_expr<A, B, C, Max>> {
    using expr_t     = etl::dyn_pool_upsample_2d_expr<A, B, C, Max>; 
    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        = false;                                
    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 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;

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

    static size_t size(const expr_t& e) {
        return etl::size(e.a());
    }

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

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

template <typename A, typename B, typename C>
dyn_pool_upsample_2d_expr<detail::build_type<A>, detail::build_type<B>, detail::build_type<C>, true> max_pool_upsample_2d(
    A&& input, B&& output, C&& errors, size_t c1, size_t c2) {
    return {input, output, errors, c1, c2, c1, c2, 0, 0};
}

template <typename A, typename B, typename C>
dyn_pool_upsample_2d_expr<detail::build_type<A>, detail::build_type<B>, detail::build_type<C>, true> max_pool_upsample_2d(
    A&& input, B&& output, C&& errors, size_t c1, size_t c2, size_t s1, size_t s2, size_t p1 = 0, size_t p2 = 0) {
    return {input, output, errors, c1, c2, s1, s2, p1, p2};
}

template <typename A, typename B, typename C>
dyn_pool_upsample_2d_expr<detail::build_type<A>, detail::build_type<B>, detail::build_type<C>, false> avg_pool_upsample_2d(
    A&& input, B&& output, C&& errors, size_t c1, size_t c2) {
    return {input, output, errors, c1, c2, c1, c2, 0, 0};
}

template <typename A, typename B, typename C>
dyn_pool_upsample_2d_expr<detail::build_type<A>, detail::build_type<B>, detail::build_type<C>, false> avg_pool_upsample_2d(
    A&& input, B&& output, C&& errors, size_t c1, size_t c2, size_t s1, size_t s2, size_t p1 = 0, size_t p2 = 0) {
    return {input, output, errors, c1, c2, s1, s2, p1, p2};
}

} //end of namespace etl