conv_2d_valid_deep_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/conv_deep.hpp"

namespace etl {

template <etl_expr A, etl_expr B, size_t S1, size_t S2, size_t P1, size_t P2, bool Flipped>
struct conv_2d_valid_deep_expr : base_temporary_expr_bin<conv_2d_valid_deep_expr<A, B, S1, S2, P1, P2, Flipped>, A, B> {
    using value_type  = value_t<A>;                                             
    using this_type   = conv_2d_valid_deep_expr<A, B, S1, S2, P1, P2, Flipped>; 
    using base_type   = base_temporary_expr_bin<this_type, A, B>;               
    using left_traits = decay_traits<A>;                                        

    static constexpr size_t D           = left_traits::dimensions();  
    static constexpr auto storage_order = left_traits::storage_order; 

    static constexpr bool gpu_computable = false;

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

    // Assignment functions

    template <same_dimensions<A> I, same_dimensions<A> K, same_dimensions<A> C>
    static void check([[maybe_unused]] const I& input, [[maybe_unused]] const K& kernel, [[maybe_unused]] const C& conv) {
        if constexpr (all_fast<A, B, C>) {
            static_assert(etl::dim<D - 2, C>() == (etl::dim<D - 2, I>() + etl::dim<D - 2, K>() + 2 * P1) / S1 - 1, "Invalid dimensions for conv2_valid_deep");
            static_assert(etl::dim<D - 1, C>() == (etl::dim<D - 1, I>() + etl::dim<D - 1, K>() + 2 * P2) / S2 - 1, "Invalid dimensions for conv2_valid_deep");
        } else {
            cpp_assert(etl::dim(conv, D - 2) == (etl::dim(input, D - 2) + etl::dim(kernel, D - 2) + 2 * P1) / S1 - 1,
                       "Invalid dimensions for conv2_valid_deep");
            cpp_assert(etl::dim(conv, D - 1) == (etl::dim(input, D - 1) + etl::dim(kernel, D - 1) + 2 * P2) / S2 - 1,
                       "Invalid dimensions for conv2_valid_deep");
        }
    }

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

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

        check(a, b, c);

        if constexpr (Flipped) {
            detail::conv2_valid_flipped_deep_impl<S1, S2, P1, P2>::apply(smart_forward(a), smart_forward(b), c);
        } else {
            detail::conv2_valid_deep_impl<S1, S2, P1, P2>::apply(smart_forward(a), smart_forward(b), c);
        }
    }

    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 conv_2d_valid_deep_expr& expr) {
        return os << "conv2_valid_deep(" << expr._a << ", " << expr._b << ")";
    }
};

template <typename A, typename B, size_t S1, size_t S2, size_t P1, size_t P2, bool Flipped>
struct etl_traits<etl::conv_2d_valid_deep_expr<A, B, S1, S2, P1, P2, Flipped>> {
    using expr_t       = etl::conv_2d_valid_deep_expr<A, B, S1, S2, P1, P2, Flipped>; 
    using this_type    = etl_traits<expr_t>;                                          
    using left_expr_t  = std::decay_t<A>;                                             
    using right_expr_t = std::decay_t<B>;                                             
    using left_traits  = etl_traits<left_expr_t>;                                     
    using right_traits = etl_traits<right_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      = 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 = left_traits::storage_order;           

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

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

    template <size_t DD>
    static constexpr size_t dim() {
        return DD == D - 2 ? (etl::dim<DD, A>() + etl::dim<DD, B>() + 2 * P1) / S1 - 1
                           : DD == D - 1 ? (etl::dim<DD, A>() + etl::dim<DD, B>() + 2 * P2) / S2 - 1 : etl::dim<DD, A>();
    }

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

    static size_t size(const expr_t& e) {
        size_t s = 1;
        for (size_t d = 0; d < D; ++d) {
            s *= this_type::dim(e, d);
        }
        return s;
    }

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

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

    static constexpr size_t dimensions() {
        return D;
    }

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

template <size_t S1 = 1, size_t S2 = 1, size_t P1 = 0, size_t P2 = 0, etl_expr A, etl_expr B>
conv_2d_valid_deep_expr<detail::build_type<A>, detail::build_type<B>, S1, S2, P1, P2, false> conv_2d_valid_deep(A&& a, B&& b) {
    return conv_2d_valid_deep_expr<detail::build_type<A>, detail::build_type<B>, S1, S2, P1, P2, false>{a, b};
}

template <size_t S1 = 1, size_t S2 = 1, size_t P1 = 0, size_t P2 = 0, etl_expr A, etl_expr B, etl_expr C>
auto conv_2d_valid_deep(A&& a, B&& b, C&& c) {
    c = conv_2d_valid_deep<S1, S2, P1, P2>(a, b);

    return c;
}

template <size_t S1 = 1, size_t S2 = 1, size_t P1 = 0, size_t P2 = 0, etl_expr A, etl_expr B>
conv_2d_valid_deep_expr<detail::build_type<A>, detail::build_type<B>, S1, S2, P1, P2, true> conv_2d_valid_deep_flipped(A&& a, B&& b) {
    return conv_2d_valid_deep_expr<detail::build_type<A>, detail::build_type<B>, S1, S2, P1, P2, true>{a, b};
}

template <size_t S1 = 1, size_t S2 = 1, size_t P1 = 0, size_t P2 = 0, etl_expr A, etl_expr B, etl_expr C>
auto conv_2d_valid_deep_flipped(A&& a, B&& b, C&& c) {
    c = conv_2d_valid_deep_flipped<S1, S2, P1, P2>(a, b);

    return c;
}

} //end of namespace etl