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

namespace etl {

template <etl_expr A, size_t C1, size_t C2, size_t C3, size_t S1, size_t S2, size_t S3, size_t P1, size_t P2, size_t P3, typename Impl>
struct pool_3d_expr : base_temporary_expr_un<pool_3d_expr<A, C1, C2, C3, S1, S2, S3, P1, P2, P3, Impl>, A> {
    using value_type = value_t<A>;                                                
    using this_type  = pool_3d_expr<A, C1, C2, C3, S1, S2, S3, P1, P2, P3, Impl>; 
    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 = Impl::template gpu_computable<A>;

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

    // Assignment functions

    template <same_dimensions<A> C>
    void assign_to(C&& c) const {
        inc_counter("temp:assign");

        auto& a = this->a();

        Impl::template apply<C1, C2, C3, S1, S2, S3, P1, P2, P3>(a, c);
    }

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

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

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

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

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

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

template <etl_expr A, size_t C1, size_t C2, size_t C3, size_t S1, size_t S2, size_t S3, size_t P1, size_t P2, size_t P3, typename Impl>
struct etl_traits<etl::pool_3d_expr<A, C1, C2, C3, S1, S2, S3, P1, P2, P3, Impl>> {
    using expr_t     = etl::pool_3d_expr<A, C1, C2, C3, S1, S2, S3, P1, P2, P3, Impl>; 
    using sub_expr_t = std::decay_t<A>;                                                
    using sub_traits = etl_traits<sub_expr_t>;                                         
    using value_type = value_t<A>;                                                     

    static constexpr size_t D = sub_traits::dimensions(); 

    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 DD == D - 3 ? (decay_traits<A>::template dim<DD>() - C1 + 2 * P1) / S1 + 1
                           : DD == D - 2 ? (decay_traits<A>::template dim<DD>() - C2 + 2 * P2) / S2 + 1
                                         : DD == D - 1 ? (decay_traits<A>::template dim<DD>() - C3 + 2 * P3) / S3 + 1 : decay_traits<A>::template dim<DD>();
    }

    static size_t dim(const expr_t& e, size_t d) {
        if (d == D - 3) {
            return (etl::dim(e._a, d) - C1 + 2 * P1) / S1 + 1;
        } else if (d == D - 2) {
            return (etl::dim(e._a, d) - C2 + 2 * P2) / S2 + 1;
        } else if (d == D - 1) {
            return (etl::dim(e._a, d) - C3 + 2 * P3) / S3 + 1;
        } else {
            return etl::dim(e._a, d);
        }
    }

    static size_t size(const expr_t& e) {
        size_t acc = 1;
        for (size_t i = 0; i < D; ++i) {
            acc *= dim(e, i);
        }
        return acc;
    }

    template <size_t... I>
    static constexpr size_t size_mul(const std::index_sequence<I...>& /*seq*/) {
        return (dim<I>() * ...);
    }

    static constexpr size_t size() {
        return size_mul(std::make_index_sequence<D>());
    }

    static constexpr size_t dimensions() {
        return D;
    }

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

template <size_t C1, size_t C2, size_t C3, size_t S1 = C1, size_t S2 = C2, size_t S3 = C3, size_t P1 = 0, size_t P2 = 0, size_t P3 = 0, etl_expr E>
pool_3d_expr<detail::build_type<E>, C1, C2, C3, S1, S2, S3, P1, P2, P3, impl::max_pool_3d> max_pool_3d(E&& value) {
    return pool_3d_expr<detail::build_type<E>, C1, C2, C3, S1, S2, S3, P1, P2, P3, impl::max_pool_3d>{value};
}

template <size_t C1, size_t C2, size_t C3, size_t S1 = C1, size_t S2 = C2, size_t S3 = C3, size_t P1 = 0, size_t P2 = 0, size_t P3 = 0, etl_expr E>
pool_3d_expr<detail::build_type<E>, C1, C2, C3, S1, S2, S3, P1, P2, P3, impl::avg_pool_3d> avg_pool_3d(E&& value) {
    return pool_3d_expr<detail::build_type<E>, C1, C2, C3, S1, S2, S3, P1, P2, P3, impl::avg_pool_3d>{value};
}

} //end of namespace etl