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, size_t C1, size_t C2, size_t S1, size_t S2, size_t P1, size_t P2, bool Max>
struct pool_upsample_2d_expr : base_temporary_expr_tern<pool_upsample_2d_expr<A, B, C, C1, C2, S1, S2, P1, P2, Max>, A, B, C> {
    using value_type = value_t<A>;                                                  
    using sub_traits = etl::decay_traits<A>;                                        
    using this_type  = pool_upsample_2d_expr<A, B, C, C1, C2, S1, S2, P1, P2, 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>;

    friend struct etl_traits<pool_upsample_2d_expr>;

    pool_upsample_2d_expr(A a, B b, C c) : base_type(a, b, c) {
        //Nothing else to init
    }

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

        if constexpr (all_fast<A, B, C, R>) {
            static_assert(etl::decay_traits<R>::size() == etl::decay_traits<A>::size(), "max_pool_upsample_2d:A and R must have the same size");
            static_assert(etl::decay_traits<B>::size() == etl::decay_traits<C>::size(), "max_pool_upsample_2d:B and C must have the same size");

            static_assert(etl::decay_traits<A>::template dim<D - 2>() == S1 * (etl::decay_traits<B>::template dim<D - 2>() - 1) + C1 - 2 * P1,
                          "Invalid pooling dimensions for max_pool_upsample_2d");
            static_assert(etl::decay_traits<A>::template dim<D - 1>() == S2 * (etl::decay_traits<B>::template dim<D - 1>() - 1) + C2 - 2 * P2,
                          "Invalid pooling dimensions for max_pool_upsample_2d");
        } else {
            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) + C1 - 2 * P1, "Invalid pooling dimensions for max_pool_upsample_2d");
            cpp_assert(etl::dim<D - 1>(a) == S2 * (etl::dim<D - 1>(b) - 1) + C2 - 2 * P2, "Invalid pooling dimensions for max_pool_upsample_2d");
        }
    }

    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<C1, C2, S1, S2, P1, P2>(smart_forward(a), smart_forward(b), smart_forward(c), result);
                }
            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<C1, C2, S1, S2, P1, P2>(smart_forward(a), smart_forward(b), smart_forward(c), result);
                }
            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 pool_upsample_2d_expr& expr) {
        return os << "max_pool_upsample2(" << expr._a << ", " << expr._b << ", " << expr._c << ")";
    }
};

template <typename A, typename B, typename C, size_t C1, size_t C2, size_t S1, size_t S2, size_t P1, size_t P2, bool Max>
struct etl_traits<etl::pool_upsample_2d_expr<A, B, C, C1, C2, S1, S2, P1, P2, Max>> {
    using expr_t     = etl::pool_upsample_2d_expr<A, B, C, C1, C2, S1, S2, P1, P2, 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        = 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 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 decay_traits<A>::template dim<DD>();
    }

    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 size() {
        return decay_traits<A>::size();
    }

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

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

template <size_t C1, size_t C2, size_t S1 = C1, size_t S2 = C2, size_t P1 = 0, size_t P2 = 0, etl_expr A, etl_expr B, etl_expr C>
pool_upsample_2d_expr<detail::build_type<A>, detail::build_type<B>, detail::build_type<C>, C1, C2, S1, S2, P1, P2, true> max_pool_upsample_2d(A&& input,
                                                                                                                                              B&& output,
                                                                                                                                              C&& errors) {
    return {input, output, errors};
}

template <size_t C1, size_t C2, size_t S1 = C1, size_t S2 = C2, size_t P1 = 0, size_t P2 = 0, etl_expr A, etl_expr B, etl_expr C>
pool_upsample_2d_expr<detail::build_type<A>, detail::build_type<B>, detail::build_type<C>, C1, C2, S1, S2, P1, P2, false> avg_pool_upsample_2d(A&& input,
                                                                                                                                               B&& output,
                                                                                                                                               C&& errors) {
    return {input, output, errors};
}

} //end of namespace etl