dyn_conv_4d_backward_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.hpp"

namespace etl {

template <etl_4d A, etl_4d B, bool Flipped>
struct dyn_conv_4d_backward_expr : base_temporary_expr_bin<dyn_conv_4d_backward_expr<A, B, Flipped>, A, B> {
    using value_type  = value_t<A>;                               
    using this_type   = dyn_conv_4d_backward_expr<A, B, Flipped>; 
    using base_type   = base_temporary_expr_bin<this_type, A, B>; 
    using left_traits = decay_traits<A>;                          

    static constexpr auto storage_order = left_traits::storage_order; 

    static constexpr bool gpu_computable = cudnn_enabled && impl::cudnn::conv_possible_<A, B>;

    const size_t s1; 
    const size_t s2; 
    const size_t p1; 
    const size_t p2; 

    explicit dyn_conv_4d_backward_expr(A a, B b, size_t s1, size_t s2, size_t p1, size_t p2) : base_type(a, b), s1(s1), s2(s2), p1(p1), p2(p2) {
        //Nothing else to init
    }

    // Assignment functions

    template <etl_4d I, etl_4d K, etl_4d C>
    void check([[maybe_unused]] const I& input, [[maybe_unused]] const K& kernel, [[maybe_unused]] const C& conv) const {
        cpp_assert(etl::dim(conv, 0) == etl::dim(input, 0), "Invalid dimensions for conv4_backward");
        cpp_assert(etl::dim(conv, 1) == etl::dim(kernel, 1), "Invalid dimensions for conv4_backward");
        cpp_assert(etl::dim(input, 1) == etl::dim(kernel, 0), "Invalid dimensions for conv4_backward");

        cpp_assert(etl::dim(conv, 2) == s1 * (etl::dim(input, 2) - 1) + etl::dim(kernel, 2) - 2 * p1, "Invalid dimensions for conv2_backward");
        cpp_assert(etl::dim(conv, 3) == s2 * (etl::dim(input, 3) - 1) + etl::dim(kernel, 3) - 2 * p2, "Invalid dimensions for conv2_backward");
    }

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

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

        check(input, kernel, conv);

        // Need K1 / K2 to compute transposed padding
        const size_t k1 = etl::dim<2>(kernel);
        const size_t k2 = etl::dim<3>(kernel);

        if constexpr (Flipped) {
            // The GPU implementation needs the real forward parameters, not the
            // converted backward parameters
            if constexpr (cudnn_enabled && all_floating<A, B, C>) {
                impl::cudnn::conv4_backward_data_flipped(smart_forward_gpu(input), smart_forward_gpu(kernel), conv, s1, s2, p1, p2);
                return;
            } else {
                // 1. Handle unit strides
                if (s1 == 1 && s2 == 1) {
                    if (p1 == 0 && p2 == 0) {
                        // Unit strides, non-zero padding -> Full convolution
                        detail::conv4_full_flipped_impl::apply(input, kernel, conv);
                    } else {
                        // Unit strides, zero padding -> Valid convolution with the correct padding
                        detail::dyn_conv4_valid_back_flipped_impl::apply(input, kernel, conv, 1, 1, k1 - p1 - 1, k2 - p2 - 1);
                    }
                }
                // 2. Handle non_unit strides
                else {
                    // Fractionally-strided convolution needs inner padding of the input
                    auto strided_input = impl::common::inner_pad(input, s1, s2);

                    if (p1 == 0 && p2 == 0) {
                        // Non-unit strides, non-zero padding -> Fractionally-strided full convolution
                        detail::conv4_full_flipped_impl::apply(strided_input, kernel, conv);
                    } else {
                        // Non-unit strides, zero padding -> Fractionally-strided Valid convolution with the correct padding
                        detail::dyn_conv4_valid_back_flipped_impl::apply(strided_input, kernel, conv, 1, 1, k1 - p1 - 1, k2 - p2 - 1);
                    }
                }
            }
        } else {
            // The GPU implementation needs the real forward parameters, not the
            // converted backward parameters
            if constexpr (cudnn_enabled && all_floating<A, B, C>) {
                impl::cudnn::conv4_backward_data(smart_forward_gpu(input), smart_forward_gpu(kernel), conv, s1, s2, p1, p2);
                return;
            } else {
                // 1. Handle unit strides
                if (s1 == 1 && s2 == 1) {
                    if (p1 == 0 && p2 == 0) {
                        // Unit strides, non-zero padding -> Full convolution
                        detail::conv4_full_impl::apply(input, kernel, conv);
                    } else {
                        // Unit strides, zero padding -> Valid convolution with the correct padding
                        detail::dyn_conv4_valid_back_impl::apply(input, kernel, conv, 1, 1, k1 - p1 - 1, k2 - p2 - 1);
                    }
                }
                // 2. Handle non_unit strides
                else {
                    // Fractionally-strided convolution needs inner padding of the input
                    auto strided_input = impl::common::inner_pad(input, s1, s2);

                    if (p1 == 0 && p2 == 0) {
                        // Non-unit strides, non-zero padding -> Fractionally-strided full convolution
                        detail::conv4_full_impl::apply(strided_input, kernel, conv);
                    } else {
                        // Non-unit strides, zero padding -> Fractionally-strided Valid convolution with the correct padding
                        detail::dyn_conv4_valid_back_impl::apply(strided_input, kernel, conv, 1, 1, k1 - p1 - 1, k2 - p2 - 1);
                    }
                }
            }
        }
    }

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

template <typename A, typename B, bool Flipped>
struct etl_traits<etl::dyn_conv_4d_backward_expr<A, B, Flipped>> {
    using expr_t       = etl::dyn_conv_4d_backward_expr<A, B, Flipped>; 
    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        = 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 = left_traits::storage_order;           

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

    static size_t dim(const expr_t& e, size_t d) {
        if (d == 0) {
            return etl::dim(e._a, 0);
        } else if (d == 1) {
            return etl::dim(e._b, 1);
        } else if (d == 2) {
            return e.s1 * (etl::dim(e._a, 2) - 1) + etl::dim(e._b, 2) - 2 * e.p1;
        } else {
            return e.s2 * (etl::dim(e._a, 3) - 1) + etl::dim(e._b, 3) - 2 * e.p2;
        }
    }

    static size_t size(const expr_t& e) {
        return etl::dim(e._a, 0) * etl::dim(e._b, 1) * (e.s1 * (etl::dim(e._a, 2) - 1) + etl::dim(e._b, 2) - 2 * e.p1)
               * (e.s2 * (etl::dim(e._a, 3) - 1) + etl::dim(e._b, 3) - 2 * e.p2);
    }

    static constexpr size_t dimensions() {
        return 4;
    }

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

template <etl_expr A, etl_expr B>
dyn_conv_4d_backward_expr<detail::build_type<A>, detail::build_type<B>, false> conv_4d_backward(A&& a, B&& b, size_t s1, size_t s2, size_t p1, size_t p2) {
    return dyn_conv_4d_backward_expr<detail::build_type<A>, detail::build_type<B>, false>{a, b, s1, s2, p1, p2};
}

template <etl_expr A, etl_expr B, etl_expr C>
auto conv_4d_backward(A&& a, B&& b, C&& c, size_t s1, size_t s2, size_t p1, size_t p2) {
    c = conv_4d_backward(a, b, s1, s2, p1, p2);

    return c;
}

template <etl_expr A, etl_expr B>
dyn_conv_4d_backward_expr<detail::build_type<A>, detail::build_type<B>, true> conv_4d_backward_flipped(
    A&& a, B&& b, size_t s1, size_t s2, size_t p1, size_t p2) {
    return dyn_conv_4d_backward_expr<detail::build_type<A>, detail::build_type<B>, true>{a, b, s1, s2, p1, p2};
}

template <etl_expr A, etl_expr B, etl_expr C>
auto conv_4d_backward_flipped(A&& a, B&& b, C&& c, size_t s1, size_t s2, size_t p1, size_t p2) {
    c = conv_4d_backward_flipped(a, b, s1, s2, p1, p2);

    return c;
}

} //end of namespace etl