conv_1d_full_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_select.hpp"
#include "etl/impl/conv.hpp"

namespace etl {

template <typename A, typename B>
struct conv_1d_full_expr : base_temporary_expr_bin<conv_1d_full_expr<A, B>, A, B> {
    using value_type  = value_t<A>;                               
    using this_type   = conv_1d_full_expr<A, B>;                  
    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 = false;

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

    // Assignment functions

    template <etl_1d I, etl_1d K, etl_1d 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<0, C>() == etl::dim<0, I>() + etl::dim<0, K>() - 1, "Invalid dimensions for conv1_full");
            static_assert(etl::dim<0, I>() >= etl::dim<0, K>(), "Invalid dimensions for conv1_full");
        } else {
            cpp_assert(etl::dim(conv, 0) == etl::dim(input, 0) + etl::dim(kernel, 0) - 1, "Invalid dimensions for conv1_full");
            cpp_assert(etl::dim(input, 0) >= etl::dim(kernel, 0), "Invalid dimensions for conv1_full");
        }
    }

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

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

        check(input_raw, kernel_raw, conv);

        // Execute the correct implementation

        constexpr_select const auto impl = detail::select_conv1_impl_new<conv_type::FULL, A, B, C>();

        if constexpr_select (impl == etl::conv_impl::VEC) {
            inc_counter("inc:vec");

            if constexpr (parallel_support) {
                bool parallel_dispatch = detail::select_parallel(input_raw, kernel_raw, conv);

                if (parallel_dispatch) {
                    decltype(auto) input  = smart_forward(input_raw);
                    decltype(auto) kernel = smart_forward(kernel_raw);

                    engine_dispatch_1d(
                            [&](size_t first, size_t last) { impl::vec::conv1_full(input, kernel, conv, first, last); }, 0, etl::size(conv), parallel_dispatch);
                } else {
                    impl::vec::conv1_full(smart_forward(input_raw), smart_forward(kernel_raw), conv, 0, etl::size(conv));
                }
            } else {
                impl::standard::conv1_full(smart_forward(input_raw), smart_forward(kernel_raw), conv, 0, etl::size(conv));
            }
        } else if constexpr_select (impl == etl::conv_impl::STD) {
            inc_counter("inc:std");

            if constexpr (parallel_support) {
                bool parallel_dispatch = detail::select_parallel(input_raw, kernel_raw, conv);

                if (parallel_dispatch) {
                    decltype(auto) input  = smart_forward(input_raw);
                    decltype(auto) kernel = smart_forward(kernel_raw);

                    engine_dispatch_1d([&](size_t first, size_t last) { impl::standard::conv1_full(input, kernel, conv, first, last); },
                                       0,
                                       etl::size(conv),
                                       parallel_dispatch);
                } else {
                    impl::standard::conv1_full(smart_forward(input_raw), smart_forward(kernel_raw), conv, 0, etl::size(conv));
                }
            } else {
                impl::standard::conv1_full(smart_forward(input_raw), smart_forward(kernel_raw), conv, 0, etl::size(conv));
            }
        } else if constexpr_select (impl == etl::conv_impl::FFT_STD) {
            inc_counter("inc:fft_std");
            impl::standard::conv1_full_fft(smart_forward(input_raw), smart_forward(kernel_raw), conv);
        } else if constexpr_select (impl == etl::conv_impl::FFT_MKL) {
            inc_counter("inc:fft_mkl");
            impl::blas::conv1_full(smart_forward(input_raw), smart_forward(kernel_raw), conv);
        } else if constexpr_select (impl == etl::conv_impl::FFT_CUFFT) {
            inc_counter("inc:fft_cufft");
            impl::cufft::conv1_full(smart_forward_gpu(input_raw), smart_forward_gpu(kernel_raw), conv);
        } else if constexpr_select (impl == etl::conv_impl::EGBLAS) {
            if constexpr (all_homogeneous<A, B, C>) {
                decltype(auto) input  = smart_forward_gpu(input_raw);
                decltype(auto) kernel = smart_forward_gpu(kernel_raw);

                input.ensure_gpu_up_to_date();
                kernel.ensure_gpu_up_to_date();

                conv.ensure_gpu_allocated();

                impl::egblas::conv1_full(
                        etl::size(input_raw), etl::size(kernel_raw), value_type(1), input.gpu_memory(), 1, kernel.gpu_memory(), 1, conv.gpu_memory(), 1);

                conv.validate_gpu();
                conv.invalidate_cpu();
            } else {
                cpp_unreachable("Invalid conv implementation selection");
            }
        } else {
            cpp_unreachable("Invalid conv implementation selection");
        }
    }

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

template <typename A, typename B>
struct etl_traits<etl::conv_1d_full_expr<A, B>> {
    using expr_t       = etl::conv_1d_full_expr<A, B>; 
    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;           

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

    template <size_t DD>
    static constexpr size_t dim() {
        return etl::dim<0, A>() + etl::dim<0, B>() - 1;
    }

    static size_t dim(const expr_t& e, [[maybe_unused]] size_t d) {
        return etl::dim(e._a, 0) + etl::dim(e._b, 0) - 1;
    }

    static size_t size(const expr_t& e) {
        return etl::dim(e._a, 0) + etl::dim(e._b, 0) - 1;
    }

    static constexpr size_t size() {
        return etl::dim<0, A>() + etl::dim<0, B>() - 1;
    }

    static constexpr size_t dimensions() {
        return 1;
    }

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

template <etl_expr A, etl_expr B>
conv_1d_full_expr<detail::build_type<A>, detail::build_type<B>> conv_1d_full(A&& a, B&& b) {
    return conv_1d_full_expr<detail::build_type<A>, detail::build_type<B>>{a, b};
}

template <etl_expr A, etl_expr B, etl_expr C>
auto conv_1d_full(A&& a, B&& b, C&& c) {
    c = conv_1d_full(a, b);
    return c;
}

} //end of namespace etl