transformers.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/tmp.hpp"

//The transformers
#include "etl/op/flip_transformers.hpp"
#include "etl/op/rep_transformers.hpp"
#include "etl/op/reduc_transformers.hpp"

namespace etl {

template <typename L, typename R>
struct mm_mul_transformer {
    using left_type  = L;          
    using right_type = R;          
    using value_type = value_t<L>; 

    left_type left;   
    right_type right; 

    static constexpr bool gpu_computable = false;

    mm_mul_transformer(left_type left, right_type right) : left(left), right(right) {
        check_mmul_sizes(left, right);
    }

    value_type operator[](size_t i) const {
        if constexpr (etl::decay_traits<left_type>::storage_order == order::RowMajor) {
            const auto n = etl::dim<1>(right);
            return operator()(i / n, i % n);
        } else {
            const auto m = etl::dim<0>(left);
            return operator()(i % m, i / m);
        }
    }

    value_type read_flat(size_t i) const {
        if constexpr (etl::decay_traits<left_type>::storage_order == order::RowMajor) {
            const auto n = etl::dim<1>(right);
            return operator()(i / n, i % n);
        } else {
            const auto m = etl::dim<0>(left);
            return operator()(i % m, i / m);
        }
    }

    value_type operator()(size_t i, size_t j) const {
        value_type c = 0;

        for (size_t k = 0; k < columns(left); k++) {
            c += left(i, k) * right(k, j);
        }

        return c;
    }

    template <typename E>
    bool alias(const E& rhs) const noexcept {
        return left.alias(rhs) || right.alias(rhs);
    }

    // Internals

    void visit(detail::evaluator_visitor& visitor) const {
        left.visit(visitor);
        right.visit(visitor);
    }

    void ensure_cpu_up_to_date() const {
        // Need to ensure both LHS and RHS
        left.ensure_cpu_up_to_date();
        right.ensure_cpu_up_to_date();
    }

    void ensure_gpu_up_to_date() const {
        // Need to ensure both LHS and RHS
        left.ensure_gpu_up_to_date();
        right.ensure_gpu_up_to_date();
    }

    friend std::ostream& operator<<(std::ostream& os, const mm_mul_transformer& transformer) {
        return os << "mm_mul(" << transformer.left << "," << transformer.right << ")";
    }

private:
    template <typename A, typename B>
    void check_mmul_sizes([[maybe_unused]] const A& a, [[maybe_unused]] const B& b) {
        if constexpr (all_fast<A, B>) {
            static_assert(etl_traits<A>::template dim<1>() == etl_traits<B>::template dim<0>() //interior dimensions
                          ,
                          "Invalid sizes for multiplication");
        } else {
            cpp_assert(dim<1>(a) == dim<0>(b) //interior dimensions
                       ,
                       "Invalid sizes for multiplication");
        }
    }
};

template <etl_1d T>
struct dyn_convmtx_transformer {
    using sub_type   = T;          
    using value_type = value_t<T>; 

    static constexpr bool gpu_computable = false;

    sub_type sub; 
    size_t h;     

    dyn_convmtx_transformer(sub_type expr, size_t h) : sub(expr), h(h) {}

    value_type operator[](size_t i) const {
        if constexpr (decay_traits<sub_type>::storage_order == order::RowMajor) {
            size_t i_i = i / (etl::size(sub) + h - 1);
            size_t i_j = i % (etl::size(sub) + h - 1);
            return operator()(i_i, i_j);
        } else {
            size_t i_i = i % h;
            size_t i_j = i / h;
            return operator()(i_i, i_j);
        }
    }

    value_type read_flat(size_t i) const {
        return operator[](i);
    }

    value_type operator()(size_t i, size_t j) const {
        if (j < i) {
            return value_type(0);
        } else if (j >= etl::size(sub) + i) {
            return value_type(0);
        } else {
            return sub(j - i);
        }
    }

    template <typename E>
    bool alias(const E& rhs) const noexcept {
        return sub.alias(rhs);
    }

    // Internals
    void visit(detail::evaluator_visitor& visitor) const {
        sub.visit(visitor);
    }

    void ensure_cpu_up_to_date() const {
        // Need to ensure sub value
        sub.ensure_cpu_up_to_date();
    }

    void ensure_gpu_up_to_date() const {
        // Need to ensure both LHS and RHS
        sub.ensure_gpu_up_to_date();
    }
};

template <etl_2d T>
struct dyn_convmtx2_transformer {
    using sub_type   = T;          
    using value_type = value_t<T>; 

    static constexpr bool gpu_computable = false;

    sub_type sub;          
    const size_t k1;       
    const size_t k2;       
    size_t i1;             
    size_t i2;             
    size_t inner_paddings; 
    size_t inner_padding;  

    dyn_convmtx2_transformer(sub_type sub, size_t k1, size_t k2) : sub(sub), k1(k1), k2(k2) {
        i1 = etl::dim<0>(sub);
        i2 = etl::dim<1>(sub);

        size_t c_height = (i1 + k1 - 1) * (i2 + k2 - 1);
        size_t c_width  = k1 * k2;

        auto max_fill  = c_height - ((i1 + k1 - 1) * ((c_width - 1) / k1) + (c_width - 1) % k1);
        inner_paddings = max_fill - (i1 * i2);
        inner_padding  = inner_paddings / (i2 - 1);
    }

    value_type operator[](size_t i) const {
        size_t i_i = i / (k1 * k2);
        size_t i_j = i % (k1 * k2);
        return (*this)(i_i, i_j);
    }

    value_type read_flat(size_t i) const noexcept {
        size_t i_i = i / (k1 * k2);
        size_t i_j = i % (k1 * k2);
        return (*this)(i_i, i_j);
    }

    value_type operator()(size_t i, size_t j) const {
        auto top_padding = (i1 + k1 - 1) * (j / k1) + j % k1;

        if (i < top_padding || i >= top_padding + (i1 * i2) + inner_paddings) {
            return value_type(0);
        } else {
            auto inner = i - top_padding;
            auto col   = inner % (i1 + inner_padding);
            auto block = inner / (i1 + inner_padding);

            if (col >= i1) {
                return value_type(0);
            } else {
                return sub(col, block);
            }
        }
    }

    template <typename E>
    bool alias(const E& rhs) const noexcept {
        return sub.alias(rhs);
    }

    // Internals

    void visit(detail::evaluator_visitor& visitor) const {
        sub.visit(visitor);
    }

    void ensure_cpu_up_to_date() const {
        // Need to ensure sub value
        sub.ensure_cpu_up_to_date();
    }

    void ensure_gpu_up_to_date() const {
        // Need to ensure both LHS and RHS
        sub.ensure_gpu_up_to_date();
    }
};

//TODO Make a temporary_expr out of this

template <typename A, typename M>
void convmtx2_direct_t(M& m, A&& sub, size_t k1, size_t k2) {
    const size_t i1 = etl::dim<0>(sub);
    const size_t i2 = etl::dim<1>(sub);

    const size_t c_height = (i1 + k1 - 1) * (i2 + k2 - 1);
    const size_t c_width  = k1 * k2;

    const auto max_fill       = c_height - ((i1 + k1 - 1) * ((c_width - 1) / k1) + (c_width - 1) % k1);
    const auto inner_paddings = max_fill - (i1 * i2);
    const auto inner_padding  = inner_paddings / (i2 - 1);

    auto* __restrict mm = m.memory_start();
    auto* __restrict ss = sub.memory_start();

    std::fill(mm, mm + etl::size(m), 0.0);

    for (size_t j = 0; j < c_width; ++j) {
        size_t big_i = (i1 + k1 - 1) * (j / k1) + j % k1;

        for (size_t ii = 0; ii < etl::dim<1>(sub); ++ii) {
            for (size_t jj = 0; jj < etl::dim<0>(sub); ++jj) {
                mm[j * c_width + big_i] = ss[jj * i2 + ii];
                ++big_i;
            }
            big_i += inner_padding;
        }
    }
}

//TODO Adapt this to an expression

template <typename A, typename M>
void im2col_direct(M& m, A&& sub, size_t k1, size_t k2) {
    if constexpr (all_dma<A, M>) {
        // This is a direct memory version
        // On gcc and clang, this is significantly faster
        const size_t i1 = etl::dim<0>(sub);
        const size_t i2 = etl::dim<1>(sub);

        const auto m_width = (i1 - k1 + 1) * (i2 - k2 + 1);

        const auto mm = m.memory_start();
        const auto ss = sub.memory_start();

        for (size_t b = 0; b < m_width; ++b) {
            auto s_i = b % (i1 - k1 + 1);
            auto s_j = b / (i1 - k1 + 1);

            for (size_t b_i = 0; b_i < k1; ++b_i) {
                for (size_t b_j = 0; b_j < k2; ++b_j) {
                    mm[(b_j * k1 + b_i) * m_width + b] = ss[(s_i + b_i) * i2 + s_j + b_j];
                }
            }
        }
    } else {
        const size_t i1 = etl::dim<0>(sub);
        const size_t i2 = etl::dim<1>(sub);

        const size_t m_width = (i1 - k1 + 1) * (i2 - k2 + 1);

        for (size_t b = 0; b < m_width; ++b) {
            auto s_i = b % (i1 - k1 + 1);
            auto s_j = b / (i1 - k1 + 1);

            for (size_t b_i = 0; b_i < k1; ++b_i) {
                for (size_t b_j = 0; b_j < k2; ++b_j) {
                    m(b_j * k1 + b_i, b) = sub(s_i + b_i, s_j + b_j);
                }
            }
        }
    }
}

//im2col version without any need for transpose

template <etl_dma A, etl_dma M>
void im2col_direct_tr(M& m, A&& sub, size_t k1, size_t k2) {
    const size_t i1 = etl::dim<0>(sub);
    const size_t i2 = etl::dim<1>(sub);

    const auto height = i1 - k1 + 1;
    const auto width  = i2 - k2 + 1;

    const auto mm = m.memory_start();
    const auto ss = sub.memory_start();

    for (size_t c = 0; c < k1 * k2; ++c) {
        const size_t w_source = c % k2;
        const size_t h_source = (c / k2) % k1;
        const size_t c_source = c / (k1 * k2);

        for (size_t h = 0; h < height; ++h) {
            const size_t block_source = (c_source * i1 + h + h_source) * i2 + w_source;
            const size_t block_target = (c * height + h) * width;

            direct_copy_n(ss + block_source, mm + block_target, width);
        }
    }
}

template <etl_dma A, etl_dma M>
void im2col_direct_tr_multi(M& m, A&& sub, size_t k1, size_t k2) {
    const auto N  = etl::dim<0>(sub);
    const auto i1 = etl::dim<1>(sub);
    const auto i2 = etl::dim<2>(sub);

    const auto height = i1 - k1 + 1;
    const auto width  = i2 - k2 + 1;

    const auto mm = m.memory_start();
    const auto ss = sub.memory_start();

    for (size_t w = 0; w < k1 * k2; ++w) {
        const auto w_source = w % k2;
        const auto h_source = (w / k2) % k1;
        const auto c_source = w / (k1 * k2);

        for (size_t i = 0; i < N; ++i) {
            for (size_t h = 0; h < height; ++h) {
                const auto block_source = ((c_source * i1 + h + h_source) * i2 + w_source) + (i) * (i1 * i2);
                const auto block_target = (w * N + i) * (height * width) + h * width;

                etl::direct_copy_n(ss + block_source, mm + block_target, width);
            }
        }
    }
}

template <typename LE, typename RE>
struct etl_traits<mm_mul_transformer<LE, RE>> {
    using expr_t       = etl::mm_mul_transformer<LE, RE>; 
    using left_expr_t  = std::decay_t<LE>;                
    using right_expr_t = std::decay_t<RE>;                
    using value_type   = value_t<left_expr_t>;            
    using l_traits     = etl_traits<left_expr_t>;         
    using r_traits     = etl_traits<right_expr_t>;        

    static constexpr bool is_etl         = true;                                             
    static constexpr bool is_transformer = true;                                             
    static constexpr bool is_view        = false;                                            
    static constexpr bool is_magic_view  = false;                                            
    static constexpr bool is_fast        = l_traits::is_fast && r_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      = false;                                            
    static constexpr bool is_generator   = false;                                            
    static constexpr bool is_padded      = false;                                            
    static constexpr bool is_aligned     = false;                                            
    static constexpr bool is_temporary   = l_traits::is_temporary || r_traits::is_temporary; 
    static constexpr bool gpu_computable = false;                                            
    static constexpr order storage_order = l_traits::is_generator ? r_traits::storage_order : l_traits::storage_order; 

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

    static size_t size(const expr_t& v) {
        return dim(v, 0) * dim(v, 1);
    }

    static size_t dim(const expr_t& v, size_t d) {
        if (d == 0) {
            return etl::dim(v.left, 0);
        } else {
            cpp_assert(d == 1, "Only 2D mmul are supported");

            return etl::dim(v.right, 1);
        }
    }

    static constexpr size_t size() {
        return etl_traits<left_expr_t>::template dim<0>() * etl_traits<right_expr_t>::template dim<1>();
    }

    template <size_t D>
    static constexpr size_t dim() requires(D < 2) {
        return D == 0 ? etl_traits<left_expr_t>::template dim<0>() : etl_traits<right_expr_t>::template dim<1>();
    }

    static constexpr size_t dimensions() {
        return 2;
    }

    static constexpr int complexity() noexcept {
        return 5;
    }
};

template <typename E>
struct etl_traits<dyn_convmtx_transformer<E>> {
    using expr_t     = etl::dyn_convmtx_transformer<E>; 
    using sub_expr_t = std::decay_t<E>;                 
    using value_type = value_t<sub_expr_t>;             

    static constexpr bool is_etl         = true;                                  
    static constexpr bool is_transformer = true;                                  
    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      = false;                                 
    static constexpr bool is_generator   = false;                                 
    static constexpr bool is_padded      = false;                                 
    static constexpr bool is_aligned     = false;                                 
    static constexpr bool is_temporary   = etl_traits<sub_expr_t>::is_temporary;  
    static constexpr bool gpu_computable = false;                                 
    static constexpr order storage_order = etl_traits<sub_expr_t>::storage_order; 

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

    static size_t size(const expr_t& v) {
        return v.h * (etl::size(v.sub) + v.h - 1);
    }

    static size_t dim(const expr_t& v, size_t d) {
        if (d == 0) {
            return v.h;
        } else {
            return etl::size(v.sub) + v.h - 1;
        }
    }

    static constexpr size_t dimensions() {
        return 2;
    }

    static constexpr int complexity() noexcept {
        return 5;
    }
};

template <typename E>
struct etl_traits<dyn_convmtx2_transformer<E>> {
    using expr_t     = etl::dyn_convmtx2_transformer<E>; 
    using sub_expr_t = std::decay_t<E>;                  
    using value_type = value_t<sub_expr_t>;              

    static constexpr bool is_etl         = true;                                  
    static constexpr bool is_transformer = true;                                  
    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      = false;                                 
    static constexpr bool is_generator   = false;                                 
    static constexpr bool is_padded      = false;                                 
    static constexpr bool is_aligned     = false;                                 
    static constexpr bool is_temporary   = etl_traits<sub_expr_t>::is_temporary;  
    static constexpr bool gpu_computable = false;                                 
    static constexpr order storage_order = etl_traits<sub_expr_t>::storage_order; 

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

    static size_t size(const expr_t& v) {
        auto c_height = (etl::dim<0>(v.sub) + v.k1 - 1) * (etl::dim<1>(v.sub) + v.k2 - 1);
        auto c_width  = v.k1 * v.k2;
        return c_height * c_width;
    }

    static size_t dim(const expr_t& v, size_t d) {
        if (d == 0) {
            return (etl::dim<0>(v.sub) + v.k1 - 1) * (etl::dim<1>(v.sub) + v.k2 - 1);
        } else {
            return v.k1 * v.k2;
        }
    }

    static constexpr size_t dimensions() {
        return 2;
    }

    static constexpr int complexity() noexcept {
        return 5;
    }
};

} //end of namespace etl